"""
GAIA Agent Deployment Script

This script prepares and deploys the GAIA agent to Hugging Face Space.
It handles packaging, file preparation, and deployment of the latest version.
"""

import os
import sys
import shutil
import subprocess
import argparse
import logging
from datetime import datetime

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("deployment.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("gaia_deploy")

def setup_environment():
    """
    Set up the deployment environment and verify dependencies.
    
    Returns:
        bool: True if setup succeeds, False otherwise
    """
    logger.info("Setting up deployment environment")
    
    # Check for required environmental variables
    required_env_vars = ["HF_TOKEN"]
    missing_vars = [var for var in required_env_vars if not os.environ.get(var)]
    
    if missing_vars:
        logger.error(f"Missing required environment variables: {', '.join(missing_vars)}")
        logger.error("Please set the HF_TOKEN environment variable with your Hugging Face API token")
        return False
    
    # Check for required dependencies
    try:
        # Check for git
        subprocess.run(['git', '--version'], check=True, capture_output=True)
        
        # Check for huggingface_hub
        import huggingface_hub
        logger.info(f"Using huggingface_hub version: {huggingface_hub.__version__}")
        
    except (subprocess.SubprocessError, ImportError) as e:
        logger.error(f"Missing required dependency: {str(e)}")
        logger.error("Please install required dependencies: pip install huggingface_hub")
        return False
    
    logger.info("Environment setup successful")
    return True

def run_tests():
    """
    Run the test suite to ensure everything is working correctly.
    
    Returns:
        bool: True if all tests pass, False otherwise
    """
    logger.info("Running GAIA agent tests")
    
    try:
        # Change to the tests directory
        os.chdir("src/gaia/agent/tests")
        
        # Run the tests
        result = subprocess.run(
            [sys.executable, "run_tests.py"],
            check=True,
            capture_output=True,
            text=True
        )
        
        logger.info("Test output:")
        logger.info(result.stdout)
        
        # Go back to the original directory
        os.chdir("../../../..")
        
        logger.info("All tests passed successfully")
        return True
        
    except subprocess.CalledProcessError as e:
        logger.error("Tests failed")
        logger.error(e.stdout)
        logger.error(e.stderr)
        
        # Go back to the original directory
        os.chdir("../../../..")
        
        return False

def prepare_deployment_files():
    """
    Prepare files for deployment by ensuring all necessary files are present.
    
    Returns:
        bool: True if preparation succeeds, False otherwise
    """
    logger.info("Preparing deployment files")
    
    # Check for required files
    required_files = [
        "app.py",
        "requirements.txt",
        "README.md",
        "src/gaia/agent/agent_integrated.py",
        "src/gaia/agent/answer_formatter.py",
        "src/gaia/agent/multimodal_processor.py"
    ]
    
    missing_files = [file for file in required_files if not os.path.exists(file)]
    
    if missing_files:
        logger.error(f"Missing required files: {', '.join(missing_files)}")
        return False
    
    # Update requirements.txt if needed
    update_requirements()
    
    # Create version info file
    create_version_info()
    
    logger.info("Deployment files prepared successfully")
    return True

def update_requirements():
    """Update requirements.txt to ensure all dependencies are listed."""
    logger.info("Updating requirements.txt")
    
    # Read existing requirements
    with open("requirements.txt", "r") as f:
        requirements = f.read().splitlines()
    
    # Add any missing requirements
    required_packages = [
        "langchain",
        "langchain_core",
        "langgraph",
        "openai",
        "pillow",
        "supabase",
        "duckduckgo-search",
        "youtube-transcript-api",
        "pydantic",
        "gradio"
    ]
    
    # Check if each package is in the requirements
    new_requirements = requirements.copy()
    for package in required_packages:
        if not any(req.startswith(f"{package}==") or req == package for req in requirements):
            new_requirements.append(package)
            logger.info(f"Added {package} to requirements.txt")
    
    # Write updated requirements
    if new_requirements != requirements:
        with open("requirements.txt", "w") as f:
            f.write("\n".join(new_requirements))
        logger.info("Updated requirements.txt")

def create_version_info():
    """Create a version info file with deployment timestamp."""
    logger.info("Creating version info file")
    
    version_info = {
        "version": "1.0.0",
        "deployed_at": datetime.now().isoformat(),
        "components": [
            "integrated_agent",
            "answer_formatter",
            "multimodal_processor",
            "component_analyzers"
        ]
    }
    
    # Write version info to a file
    version_file = "src/gaia/version.py"
    with open(version_file, "w") as f:
        f.write(f"VERSION = \"{version_info['version']}\"\n")
        f.write(f"DEPLOYED_AT = \"{version_info['deployed_at']}\"\n")
        f.write(f"COMPONENTS = {version_info['components']}\n")
    
    logger.info(f"Created version info at {version_file}")

def deploy_to_huggingface():
    """
    Deploy the agent to Hugging Face Space.
    
    Returns:
        bool: True if deployment succeeds, False otherwise
    """
    logger.info("Deploying to Hugging Face Space")
    
    try:
        from huggingface_hub import HfApi
        
        api = HfApi()
        
        # Get Hugging Face token from environment
        token = os.environ.get("HF_TOKEN")
        
        # Set repository ID (username/space-name)
        # TODO: Replace with your actual repository ID
        repo_id = "your-username/gaia-agent"
        
        # Files to upload
        files_to_upload = [
            "app.py",
            "requirements.txt",
            "README.md",
            ".gitignore",
            "src/"
        ]
        
        # Upload each file/directory
        for file_path in files_to_upload:
            logger.info(f"Uploading {file_path}")
            
            if os.path.isdir(file_path):
                # Upload directory recursively
                api.upload_folder(
                    folder_path=file_path,
                    repo_id=repo_id,
                    repo_type="space",
                    token=token
                )
            else:
                # Upload single file
                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=file_path,
                    repo_id=repo_id,
                    repo_type="space",
                    token=token
                )
        
        logger.info("Deployment to Hugging Face successful")
        return True
        
    except Exception as e:
        logger.error(f"Deployment failed: {str(e)}")
        logger.error(traceback.format_exc())
        return False

def main():
    """Main deployment function."""
    parser = argparse.ArgumentParser(description="Deploy GAIA agent to Hugging Face Space")
    parser.add_argument("--skip-tests", action="store_true", help="Skip running tests")
    parser.add_argument("--prepare-only", action="store_true", help="Only prepare files without deploying")
    args = parser.parse_args()
    
    logger.info("Starting GAIA agent deployment process")
    
    # Set up environment
    if not setup_environment():
        logger.error("Environment setup failed. Aborting deployment.")
        return 1
    
    # Run tests
    if not args.skip_tests:
        if not run_tests():
            logger.error("Tests failed. Aborting deployment.")
            return 1
    else:
        logger.info("Skipping tests as requested")
    
    # Prepare deployment files
    if not prepare_deployment_files():
        logger.error("Failed to prepare deployment files. Aborting deployment.")
        return 1
    
    # Deploy to Hugging Face
    if not args.prepare_only:
        if not deploy_to_huggingface():
            logger.error("Deployment failed.")
            return 1
    else:
        logger.info("Files prepared successfully. Skipping deployment as requested.")
    
    logger.info("Deployment process completed successfully")
    return 0

if __name__ == "__main__":
    # Add traceback import for exception handling
    import traceback
    
    try:
        sys.exit(main())
    except Exception as e:
        logger.error(f"Unhandled exception: {str(e)}")
        logger.error(traceback.format_exc())
        sys.exit(1)