import os
import sys
import json
import importlib.util
from pathlib import Path

# Define paths
DRIVE_ROOT = "/content/drive/MyDrive/AI_Tools_Platform"
TOOLS_DIR = os.path.join(DRIVE_ROOT, "tools")
INSTALLED_TOOLS_JSON = os.path.join(DRIVE_ROOT, "installed_tools.json")

# Import the Tool class from main_ui
from main_ui import Tool, MainUI

def setup_environment():
    print("Setting up environment...")
    
    # Create tools directory in Google Drive if it doesn't exist
    os.makedirs(TOOLS_DIR, exist_ok=True)
    print(f"Tools directory in Drive: {TOOLS_DIR}")
    
    # Create installed_tools.json in Google Drive if it doesn't exist
    if not os.path.exists(INSTALLED_TOOLS_JSON):
        with open(INSTALLED_TOOLS_JSON, 'w') as f:
            json.dump({"tools": []}, f, indent=4)
        print(f"Created initial installed_tools.json in Drive")
    else:
        print(f"Using existing installed_tools.json from Drive")
    
    # Verify the content of installed_tools.json
    try:
        with open(INSTALLED_TOOLS_JSON, 'r') as f:
            tools_data = json.load(f)
            if not isinstance(tools_data, dict) or "tools" not in tools_data:
                print("Warning: installed_tools.json has incorrect format. Fixing...")
                with open(INSTALLED_TOOLS_JSON, 'w') as f:
                    json.dump({"tools": []}, f, indent=4)
    except (json.JSONDecodeError, UnicodeDecodeError):
        print("Error: installed_tools.json is not valid JSON. Fixing...")
        with open(INSTALLED_TOOLS_JSON, 'w') as f:
            json.dump({"tools": []}, f, indent=4)
    
    # Add tools directory to Python path
    if TOOLS_DIR not in sys.path:
        sys.path.append(TOOLS_DIR)
    if os.getcwd() not in sys.path:
        sys.path.append(os.getcwd())
    
    print("\nEnvironment set up successfully!")

# Override MainUI methods to use direct file paths
class ColabMainUI(MainUI):
    def __init__(self):
        # Set up paths before calling parent's __init__
        self.tools = {}  # Initialize tools dictionary
        self.tools_dir = TOOLS_DIR
        self.installed_tools_json = INSTALLED_TOOLS_JSON
        self.marketplace_url = "https://marketplace.vidplus.app/api/marketplace"
        self.version = "0.1"
        
        # Create tools directory if it doesn't exist
        os.makedirs('tools', exist_ok=True)
        
        # Load tools manually instead of calling parent's __init__
        self.load_installed_tools()
    
    def load_installed_tools(self):
        """Load installed tools directly from Google Drive"""
        print(f"Loading tools from {self.installed_tools_json}")
        try:
            if os.path.exists(self.installed_tools_json):
                with open(self.installed_tools_json, 'r') as f:
                    data = json.load(f)
                    
                for tool_data in data.get('tools', []):
                    # If ui_module is specified, dynamically import it
                    if tool_data.get('ui_module'):
                        try:
                            # Modify the module path to use the direct path
                            module_name = tool_data['ui_module']
                            if module_name.startswith('tools.'):
                                # Extract the relative path and use the full path
                                rel_path = module_name.replace('tools.', '')
                                module_path = os.path.join(TOOLS_DIR, rel_path.replace('.', '/') + '.py')
                                print(f"Loading module from: {module_path}")
                                
                                # Use importlib to import the module
                                spec = importlib.util.spec_from_file_location(module_name, module_path)
                                if spec:
                                    module = importlib.util.module_from_spec(spec)
                                    spec.loader.exec_module(module)
                                    ui_class = getattr(module, tool_data['ui_class'])
                                    
                                    self.tools[tool_data['id']] = Tool(
                                        name=tool_data['name'],
                                        description=tool_data['description'],
                                        ui_class=ui_class,
                                        icon=tool_data['icon'],
                                        version=tool_data['version'],
                                        ui_module=tool_data['ui_module']
                                    )
                                else:
                                    print(f"Could not find module: {module_path}")
                            else:
                                # Standard import for non-tools modules
                                module = importlib.import_module(module_name)
                                ui_class = getattr(module, tool_data['ui_class'])
                                
                                self.tools[tool_data['id']] = Tool(
                                    name=tool_data['name'],
                                    description=tool_data['description'],
                                    ui_class=ui_class,
                                    icon=tool_data['icon'],
                                    version=tool_data['version'],
                                    ui_module=tool_data['ui_module']
                                )
                        except (ImportError, AttributeError) as e:
                            print(f"Error loading tool {tool_data['name']}: {e}")
                            import traceback
                            traceback.print_exc()
        except Exception as e:
            print(f"Error loading installed tools: {e}")
            import traceback
            traceback.print_exc()
    
    def save_installed_tools(self):
        """Save installed tools directly to Google Drive"""
        data = {
            "tools": [
                {
                    "id": tool_id,
                    "name": tool.name,
                    "description": tool.description,
                    "icon": tool.icon,
                    "version": tool.version,
                    "ui_class": tool.ui_class.__name__,
                    "ui_module": tool.ui_module
                }
                for tool_id, tool in self.tools.items()
            ]
        }
        
        with open(self.installed_tools_json, 'w') as f:
            json.dump(data, f, indent=4)
        print(f"Saved tools configuration to {self.installed_tools_json}")

if __name__ == "__main__":
    setup_environment()
    
    # Import and run the modified UI
    print("\nStarting AI Tools Platform...")
    
    ui = ColabMainUI()
    demo = ui.create_ui()
    demo.queue()
    demo.title = "AI Tools Platform"
    demo.launch(debug=True, share=True)  # share=True to get a public URL