#!/usr/bin/env python3
"""
Minimal version of the AI Knowledge Distillation Platform
For testing and debugging purposes
"""

import os
import sys
import logging
from datetime import datetime
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from fastapi import FastAPI, Request, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# Setup basic logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="AI Knowledge Distillation Platform",
    description="Minimal version for testing",
    version="2.0.0-minimal"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Create directories
for directory in ["static", "templates", "cache", "database", "logs"]:
    Path(directory).mkdir(exist_ok=True)

# Mount static files and templates
try:
    app.mount("/static", StaticFiles(directory="static"), name="static")
    templates = Jinja2Templates(directory="templates")
except Exception as e:
    logger.warning(f"Could not mount static files: {e}")
    templates = None

# Initialize components with error handling
memory_manager = None
token_manager = None
medical_dataset_manager = None

try:
    from src.core.memory_manager import AdvancedMemoryManager
    memory_manager = AdvancedMemoryManager(max_memory_gb=14.0)
    logger.info("✅ Memory manager initialized")
except Exception as e:
    logger.warning(f"⚠️ Could not initialize memory manager: {e}")

try:
    from src.core.token_manager import TokenManager
    token_manager = TokenManager()
    logger.info("✅ Token manager initialized")
except Exception as e:
    logger.warning(f"⚠️ Could not initialize token manager: {e}")

try:
    from src.medical.medical_datasets import MedicalDatasetManager
    if memory_manager:
        medical_dataset_manager = MedicalDatasetManager(memory_manager)
        logger.info("✅ Medical dataset manager initialized")
except Exception as e:
    logger.warning(f"⚠️ Could not initialize medical dataset manager: {e}")

@app.get("/", response_class=HTMLResponse)
async def read_root():
    """Serve the main web interface"""
    if templates:
        try:
            return templates.TemplateResponse("index.html", {"request": {}})
        except Exception as e:
            logger.error(f"Template error: {e}")
            return HTMLResponse("<h1>AI Knowledge Distillation Platform</h1><p>Minimal version running</p>")
    else:
        return HTMLResponse("<h1>AI Knowledge Distillation Platform</h1><p>Minimal version running</p>")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    try:
        status = {
            "status": "healthy",
            "version": "2.0.0-minimal",
            "timestamp": datetime.now().isoformat(),
            "components": {
                "memory_manager": memory_manager is not None,
                "token_manager": token_manager is not None,
                "medical_datasets": medical_dataset_manager is not None,
                "templates": templates is not None
            }
        }
        
        if memory_manager:
            try:
                memory_info = memory_manager.get_memory_info()
                status["memory"] = {
                    "usage_percent": memory_info.get("process_memory_percent", 0),
                    "available_gb": memory_info.get("system_memory_available_gb", 0)
                }
            except Exception as e:
                status["memory"] = {"error": str(e)}
        
        return status
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        return {
            "status": "unhealthy",
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }

@app.get("/tokens")
async def token_management_page(request: Request):
    """Token management page"""
    if templates:
        try:
            return templates.TemplateResponse("token-management.html", {"request": request})
        except Exception as e:
            logger.error(f"Template error: {e}")
            return HTMLResponse("<h1>Token Management</h1><p>Template not available</p>")
    else:
        return HTMLResponse("<h1>Token Management</h1><p>Templates not available</p>")

@app.get("/medical-datasets")
async def medical_datasets_page(request: Request):
    """Medical datasets page"""
    if templates:
        try:
            return templates.TemplateResponse("medical-datasets.html", {"request": request})
        except Exception as e:
            logger.error(f"Template error: {e}")
            return HTMLResponse("<h1>Medical Datasets</h1><p>Template not available</p>")
    else:
        return HTMLResponse("<h1>Medical Datasets</h1><p>Templates not available</p>")

@app.get("/api/tokens")
async def list_tokens():
    """List all saved tokens"""
    if token_manager:
        try:
            tokens = token_manager.list_tokens()
            return {"tokens": tokens}
        except Exception as e:
            logger.error(f"Error listing tokens: {e}")
            raise HTTPException(status_code=500, detail=str(e))
    else:
        return {"tokens": [], "error": "Token manager not available"}

@app.get("/api/medical-datasets")
async def list_medical_datasets():
    """List supported medical datasets"""
    if medical_dataset_manager:
        try:
            datasets = medical_dataset_manager.list_supported_datasets()
            return {"datasets": datasets}
        except Exception as e:
            logger.error(f"Error listing medical datasets: {e}")
            raise HTTPException(status_code=500, detail=str(e))
    else:
        return {"datasets": [], "error": "Medical dataset manager not available"}

@app.get("/api/system/memory")
async def get_memory_info():
    """Get current memory information"""
    if memory_manager:
        try:
            memory_info = memory_manager.get_memory_info()
            return memory_info
        except Exception as e:
            logger.error(f"Error getting memory info: {e}")
            raise HTTPException(status_code=500, detail=str(e))
    else:
        return {"error": "Memory manager not available"}

@app.get("/debug")
async def debug_info():
    """Debug information"""
    import psutil
    
    return {
        "python_version": sys.version,
        "platform": sys.platform,
        "memory_gb": psutil.virtual_memory().total / (1024**3),
        "cpu_cores": os.cpu_count(),
        "working_directory": str(Path.cwd()),
        "python_path": sys.path[:3],  # First 3 entries
        "environment_variables": {
            "OMP_NUM_THREADS": os.getenv("OMP_NUM_THREADS"),
            "MKL_NUM_THREADS": os.getenv("MKL_NUM_THREADS"),
            "HF_TOKEN": "***" if os.getenv("HF_TOKEN") else None
        },
        "components_status": {
            "memory_manager": memory_manager is not None,
            "token_manager": token_manager is not None,
            "medical_datasets": medical_dataset_manager is not None,
            "templates": templates is not None
        }
    }

if __name__ == "__main__":
    print("🚀 Starting AI Knowledge Distillation Platform (Minimal)")
    print("🌐 Access at: http://localhost:8000")
    print("🔍 Debug info: http://localhost:8000/debug")
    print("💊 Health check: http://localhost:8000/health")
    
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info"
    )