#!/usr/bin/env python3
"""
OpenManus Integration Fixer

This script fixes common issues with OpenManus installations and ensures 
proper integration with Casibase and other Atlas components.
"""

import os
import sys
import json
import shutil
import subprocess
import importlib
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('openmanus_fix.log')
    ]
)

logger = logging.getLogger(__name__)

class OpenManusFixer:
    """Utility class to fix OpenManus integration issues."""
    
    def __init__(self, openmanus_path: Optional[str] = None):
        """Initialize the fixer with the path to OpenManus."""
        if openmanus_path:
            self.openmanus_path = Path(openmanus_path)
        else:
            # Try to find OpenManus in the current directory or parent
            current_dir = Path.cwd()
            if (current_dir / "app").exists() and (current_dir / "requirements.txt").exists():
                self.openmanus_path = current_dir
            else:
                self.openmanus_path = Path(os.path.expanduser("~/OpenManus"))
        
        logger.info(f"OpenManus path: {self.openmanus_path}")
        
        if not self.openmanus_path.exists():
            raise FileNotFoundError(f"OpenManus directory not found at {self.openmanus_path}")
    
    def fix_requirements(self) -> None:
        """Fix the requirements.txt file to ensure all dependencies are properly listed."""
        req_file = self.openmanus_path / "requirements.txt"
        
        if not req_file.exists():
            logger.error(f"Requirements file not found at {req_file}")
            return
        
        logger.info("Fixing requirements.txt...")
        
        # Read existing requirements
        with open(req_file, "r") as f:
            requirements = f.readlines()
        
        # Clean requirements by removing comments and invalid entries
        cleaned_requirements = []
        for req in requirements:
            req = req.strip()
            
            # Skip empty lines and comments
            if not req or req.startswith('#'):
                cleaned_requirements.append(req)
                continue
            
            # Fix any malformed requirements
            if ' #' in req:
                req = req.split(' #')[0].strip()
            
            cleaned_requirements.append(req)
        
        # Add missing dependencies
        required_deps = {
            "flask": "flask~=2.3.3",
            "flask-cors": "flask-cors~=4.0.0",
            "requests": "requests~=2.31.0",
            "pydantic": "pydantic~=2.10.6",
            "openai": "openai~=1.66.3",
            "fastapi": "fastapi~=0.115.11",
            "uvicorn": "uvicorn~=0.34.0",
            "httpx": "httpx>=0.27.0",
            "tiktoken": "tiktoken~=0.9.0"
        }
        
        # Check for each required dependency
        for dep_name, dep_req in required_deps.items():
            found = False
            for i, req in enumerate(cleaned_requirements):
                if req.startswith(f"{dep_name}"):
                    found = True
                    # Update the dependency if needed
                    cleaned_requirements[i] = dep_req
                    break
            
            if not found:
                cleaned_requirements.append(dep_req)
        
        # Write back the fixed requirements
        with open(req_file, "w") as f:
            f.write("\n".join(cleaned_requirements) + "\n")
        
        logger.info("Requirements file fixed")
    
    def install_dependencies(self) -> None:
        """Install or upgrade required dependencies."""
        logger.info("Installing/upgrading dependencies...")
        
        try:
            subprocess.run([
                sys.executable, "-m", "pip", "install", "-r", 
                str(self.openmanus_path / "requirements.txt"),
                "--upgrade"
            ], check=True)
            logger.info("Dependencies installed successfully")
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to install dependencies: {e}")
    
    def ensure_casibase_connector(self) -> None:
        """Ensure that the Casibase connector exists and is properly configured."""
        connector_path = self.openmanus_path / "app" / "casibase_connector.py"
        
        if not connector_path.exists():
            logger.info("Creating Casibase connector...")
            
            # Check if the directory exists
            connector_path.parent.mkdir(exist_ok=True)
            
            # Create the connector file
            self.create_casibase_connector(connector_path)
        else:
            logger.info("Updating existing Casibase connector...")
            self.create_casibase_connector(connector_path)
        
        logger.info("Casibase connector updated")
    
    def create_casibase_connector(self, path: Path) -> None:
        """Create or update the Casibase connector file."""
        connector_code = """
\"\"\"
Casibase Connector for OpenManus
This module provides integration with Casibase for document storage, retrieval, and querying.
\"\"\"

import os
import json
import requests
from typing import Dict, List, Optional, Union, Any
import logging
from urllib.parse import urljoin

# Import from current package if possible, otherwise create placeholders
try:
    from .exceptions import CasibaseConnectionError
    from .config import get_config
    from .logger import get_logger
    
    logger = get_logger(__name__)
except ImportError:
    # Create minimal implementations if the imports fail
    class CasibaseConnectionError(Exception):
        \"\"\"Exception raised for Casibase connection errors.\"\"\"
        pass
    
    def get_config():
        \"\"\"Get configuration values.\"\"\"
        return {}
    
    # Set up basic logger
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

class CasibaseConnector:
    \"\"\"
    Connector for integrating with Casibase services.
    \"\"\"
    
    def __init__(
        self,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        store_name: Optional[str] = None,
        timeout: int = 30
    ):
        \"\"\"
        Initialize the Casibase connector.
        
        Args:
            base_url: Base URL for the Casibase API
            api_key: API key for authentication
            store_name: Default store name to use
            timeout: Request timeout in seconds
        \"\"\"
        config = get_config()
        
        self.base_url = base_url or config.get("CASIBASE_URL", "http://localhost:7001")
        self.base_url = self.base_url.rstrip('/')  # Remove trailing slash if present
        
        self.api_key = api_key or config.get("CASIBASE_API_KEY", "")
        self.store_name = store_name or config.get("CASIBASE_STORE", "default")
        self.timeout = timeout
        
        self.headers = {
            "Content-Type": "application/json",
            "Accept": "application/json"
        }
        
        if self.api_key:
            self.headers["Authorization"] = f"Bearer {self.api_key}"
        
        logger.info(f"Initialized Casibase connector with base_url: {self.base_url}")
    
    def _make_request(
        self,
        method: str,
        endpoint: str,
        data: Optional[Dict[str, Any]] = None,
        files: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        \"\"\"
        Make a request to the Casibase API.
        
        Args:
            method: HTTP method (GET, POST, etc.)
            endpoint: API endpoint
            data: Request data
            files: Files to upload
            params: Query parameters
            
        Returns:
            Response data as dictionary
            
        Raises:
            CasibaseConnectionError: If the request fails
        \"\"\"
        url = urljoin(self.base_url, endpoint)
        
        # If sending files, don't include Content-Type header
        request_headers = self.headers.copy()
        if files:
            request_headers.pop("Content-Type", None)
        
        try:
            if method.upper() == "GET":
                response = requests.get(
                    url,
                    params=params,
                    headers=request_headers,
                    timeout=self.timeout
                )
            elif method.upper() == "POST":
                if files:
                    # Multipart form data for file uploads
                    response = requests.post(
                        url,
                        data=data,  # Form data
                        files=files,
                        headers=request_headers,
                        timeout=self.timeout
                    )
                else:
                    # JSON data
                    response = requests.post(
                        url,
                        json=data,
                        params=params,
                        headers=request_headers,
                        timeout=self.timeout
                    )
            else:
                raise CasibaseConnectionError(f"Unsupported HTTP method: {method}")
            
            # Check for errors
            if response.status_code >= 400:
                error_msg = f"Casibase API error: {response.status_code} - {response.text}"
                logger.error(error_msg)
                raise CasibaseConnectionError(error_msg)
            
            # Return JSON response
            if response.content:
                return response.json()
            else:
                return {"status": "success"}
                
        except requests.RequestException as e:
            error_msg = f"Casibase API request failed: {str(e)}"
            logger.error(error_msg)
            raise CasibaseConnectionError(error_msg) from e
    
    def query(
        self,
        query_text: str,
        store_name: Optional[str] = None,
        model_name: Optional[str] = None,
        top_k: int = 5,
        system_prompt: Optional[str] = None
    ) -> Dict[str, Any]:
        \"\"\"
        Query Casibase with the given text.
        
        Args:
            query_text: Query text
            store_name: Store name to query (defaults to instance store_name)
            model_name: Model name to use for the query
            top_k: Number of documents to retrieve
            system_prompt: Optional system prompt to include
            
        Returns:
            Query response
        \"\"\"
        store = store_name or self.store_name
        
        payload = {
            "question": query_text,
            "store": store,
            "topK": top_k
        }
        
        if model_name:
            payload["modelName"] = model_name
            
        if system_prompt:
            payload["systemPrompt"] = system_prompt
        
        logger.debug(f"Querying Casibase: {payload}")
        return self._make_request("POST", "/api/query", data=payload)
    
    def upload_file(
        self,
        file_path: str,
        store_name: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        \"\"\"
        Upload a file to Casibase.
        
        Args:
            file_path: Path to the file
            store_name: Store name to upload to (defaults to instance store_name)
            metadata: Additional metadata to associate with the file
            
        Returns:
            Upload response
        \"\"\"
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        store = store_name or self.store_name
        metadata = metadata or {}
        
        with open(file_path, "rb") as f:
            files = {"file": (os.path.basename(file_path), f)}
            data = {"store": store}
            
            if metadata:
                data["metadata"] = json.dumps(metadata)
            
            logger.debug(f"Uploading file to Casibase: {file_path} to store {store}")
            return self._make_request("POST", "/api/upload-file", data=data, files=files)
    
    def get_stores(self) -> List[Dict[str, Any]]:
        \"\"\"
        Get available Casibase stores.
        
        Returns:
            List of stores
        \"\"\"
        response = self._make_request("GET", "/api/get-stores")
        return response.get("data", [])
    
    def search_similar(
        self,
        query_text: str,
        store_name: Optional[str] = None,
        limit: int = 5
    ) -> List[Dict[str, Any]]:
        \"\"\"
        Search for similar documents in Casibase.
        
        Args:
            query_text: Text to search for
            store_name: Store to search in (defaults to instance store_name)
            limit: Maximum number of results
            
        Returns:
            List of similar documents
        \"\"\"
        store = store_name or self.store_name
        
        payload = {
            "query": query_text,
            "store": store,
            "limit": limit
        }
        
        logger.debug(f"Searching similar documents in Casibase: {payload}")
        response = self._make_request("POST", "/api/search", data=payload)
        return response.get("data", [])
    
    def get_health(self) -> Dict[str, Any]:
        \"\"\"
        Check if the Casibase service is healthy.
        
        Returns:
            Health status
        \"\"\"
        try:
            return self._make_request("GET", "/api/health-check")
        except CasibaseConnectionError:
            return {"status": "error", "message": "Casibase service is unreachable"}
"""
        
        with open(path, "w") as f:
            f.write(connector_code.lstrip())
    
    def update_api_server(self) -> None:
        """Update the API server to include Casibase integration."""
        api_server_path = self.openmanus_path / "app" / "api_server.py"
        
        if not api_server_path.exists():
            logger.error(f"API server file not found at {api_server_path}")
            return
        
        logger.info("Checking API server for Casibase integration...")
        
        with open(api_server_path, "r") as f:
            content = f.read()
        
        # Check if Casibase is already integrated
        if "CasibaseConnector" in content and "casibase_connector" in content:
            logger.info("Casibase already integrated in API server")
            return
        
        # Add imports
        if "import argparse" in content:
            content = content.replace(
                "import argparse", 
                "import argparse\nfrom .casibase_connector import CasibaseConnector"
            )
        
        # Add endpoint for Casibase bridge
        if "def create_app():" in content:
            casibase_routes = """
    # Casibase integration routes
    @app.route('/api/casibase/health', methods=['GET'])
    def casibase_health():
        try:
            casibase = CasibaseConnector()
            health_status = casibase.get_health()
            return jsonify(health_status)
        except Exception as e:
            return jsonify({"status": "error", "message": str(e)}), 500
    
    @app.route('/api/casibase/query', methods=['POST'])
    def casibase_query():
        try:
            data = request.json
            query_text = data.get('query_text')
            store_name = data.get('store_name')
            model_name = data.get('model_name')
            
            if not query_text:
                return jsonify({"status": "error", "message": "query_text is required"}), 400
            
            casibase = CasibaseConnector()
            result = casibase.query(
                query_text=query_text,
                store_name=store_name,
                model_name=model_name
            )
            
            return jsonify(result)
        except Exception as e:
            return jsonify({"status": "error", "message": str(e)}), 500
    
    @app.route('/api/casibase/search', methods=['POST'])
    def casibase_search():
        try:
            data = request.json
            query_text = data.get('query_text')
            store_name = data.get('store_name')
            limit = data.get('limit', 5)
            
            if not query_text:
                return jsonify({"status": "error", "message": "query_text is required"}), 400
            
            casibase = CasibaseConnector()
            result = casibase.search_similar(
                query_text=query_text,
                store_name=store_name,
                limit=limit
            )
            
            return jsonify({"status": "success", "data": result})
        except Exception as e:
            return jsonify({"status": "error", "message": str(e)}), 500
"""
            
            # Insert the Casibase routes before the return statement
            if "return app" in content:
                content = content.replace(
                    "return app", 
                    f"{casibase_routes}\n    return app"
                )
        
        # Add imports for Flask if needed
        if "from flask import " in content:
            if "jsonify" not in content:
                content = content.replace(
                    "from flask import ", 
                    "from flask import jsonify, "
                )
            if "request" not in content:
                content = content.replace(
                    "from flask import ", 
                    "from flask import request, "
                )
        
        # Write back the updated content
        with open(api_server_path, "w") as f:
            f.write(content)
        
        logger.info("API server updated with Casibase integration")
    
    def run_fixes(self) -> None:
        """Run all fixes."""
        logger.info("Starting OpenManus fixes...")
        
        self.fix_requirements()
        self.install_dependencies()
        self.ensure_casibase_connector()
        self.update_api_server()
        
        logger.info("OpenManus fixes completed successfully")

def main():
    """Main entry point."""
    try:
        fixer = OpenManusFixer()
        fixer.run_fixes()
        return 0
    except Exception as e:
        logger.error(f"Error: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())