# app/app.py
import asyncio
import logging
import os
import litellm
from fastapi import FastAPI
from contextlib import asynccontextmanager
from typing import List
logger = logging.getLogger(__name__)

# Import centralized settings and other components
from app.config import settings
from app.core.vector_store.faiss_store import FaissVectorStore
from app.core.vector_store.embedder.factory import get_embedder_from_config
from app.core.providers.factory import get_tts_provider, get_stt_provider
from app.core.retrievers.faiss_db_retriever import FaissDBRetriever
from app.core.retrievers.base_retriever import Retriever
from app.db.session import create_db_and_tables
from app.db.migrate import run_migrations
from app.api.routes.api import create_api_router
from app.utils import print_config 
from app.api.dependencies import ServiceContainer, get_db
from app.core.services.session import SessionService
from app.core.services.tts import TTSService
from app.core.services.stt import STTService # NEW: Added the missing import for STTService
from app.core.services.user import UserService
from app.core.services.prompt import PromptService
from app.core.services.tool import ToolService
from app.core.services.node_registry import NodeRegistryService
from app.core.services.browser_client import BrowserServiceClient
from app.core.orchestration.scheduler import AgentScheduler # Area 4: Background Scheduler
# Note: The llm_clients import and initialization are removed as they
# are not used in RAGService's constructor based on your services.py
# from app.core.llm_clients import DeepSeekClient, GeminiClient
from fastapi.middleware.cors import CORSMiddleware


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Manages application startup and shutdown events.
    - On startup, it creates database tables.
    - On shutdown, it saves the FAISS index to disk.
    """
    print("Application startup...")
    print_config(settings)
    create_db_and_tables()
    run_migrations()
    
    # --- Reset Node Statuses (Zombie Fix) ---
    try:
        app.state.services.node_registry_service.reset_all_statuses()
    except Exception as e:
        logger.warning(f"Failed to reset node statuses: {e}")
    
    # --- Start gRPC Orchestrator (M6) ---
    try:
        from app.core.grpc.services.grpc_server import serve_grpc
        registry = app.state.services.node_registry_service
        server, orchestrator = serve_grpc(registry, port=settings.GRPC_PORT)
        app.state.grpc_server = server
        app.state.orchestrator = orchestrator
        app.state.services.with_service("orchestrator", orchestrator)
        logger.info(f"[M6] Agent Orchestrator gRPC server started on port {settings.GRPC_PORT}.")
        
        # Launch periodic Ghost Mirror cleanup
        asyncio.create_task(_periodic_mirror_cleanup(orchestrator))
        
        # Area 4: Launch Agent Background Scheduler (Zombie Sweeper & CRON)
        try:
            scheduler = app.state.services.agent_scheduler
            await scheduler.start()
        except Exception as se:
            logger.error(f"[AgentScheduler] Start fail: {se}")

        # Launch periodic LLM provider health check
        asyncio.create_task(_periodic_provider_health_check())
    except Exception as e:
        logger.error(f"[M6] Failed to start gRPC server: {e}")

    # --- Bootstrap System Skills and Local Admin ---
    try:
        from app.core.skills.bootstrap import bootstrap_system_skills
        # Use the context manager to ensure session is closed
        from app.db.session import get_db_session
        with get_db_session() as db:
            # 1. Day 1 Local Admin Seeding
            app.state.services.user_service.bootstrap_local_admin(db)
            # 2. Base System Skills setup
            bootstrap_system_skills(db)
    except Exception as e:
        logger.error(f"Failed to bootstrap database elements: {e}")

    yield
    print("Application shutdown...")
    # --- Stop gRPC Orchestrator ---
    if hasattr(app.state, 'grpc_server'):
        logger.info("[M6] Stopping gRPC server...")
        app.state.grpc_server.stop(5)
    # Access the vector_store from the application state to save it
    if hasattr(app.state, 'vector_store'):
        app.state.vector_store.save_index()

async def _periodic_mirror_cleanup(orchestrator):
    """Periodically purges orphaned ghost mirror directories from the server."""
    await asyncio.sleep(10) # Initial delay to let DB settle
    while True:
        try:
            from app.db.session import get_db_session
            from app.db import models
            with get_db_session() as db:
                # Fetch all unique sync_workspace_ids currently in DB
                sessions = db.query(models.Session).filter(models.Session.sync_workspace_id != None).all()
                active_ids = [s.sync_workspace_id for s in sessions]
                
                if hasattr(orchestrator, 'mirror'):
                    orchestrator.mirror.purge_orphaned(active_ids)
                else:
                    logger.warning("[📁🧹] Orchestrator missing mirror manager during cleanup pass.")
        except Exception as e:
            logger.error(f"[📁🧹] Ghost Mirror periodic cleanup fail: {e}")
        
        await asyncio.sleep(600) # Run every 10 minutes to prevent CPU spin-lock (Assignment #5)

async def _periodic_provider_health_check():
    """Periodically tests actual connectivity to all system-level LLM providers."""
    await asyncio.sleep(60) # Initial delay to let the system boot fully
    from app.core.providers.factory import get_llm_provider
    from sqlalchemy.orm.attributes import flag_modified
    import copy

    while True:
        try:
            from app.db.session import get_db_session
            from app.db import models
            from app.config import settings

            with get_db_session() as db:
                admin_user = db.query(models.User).filter(models.User.role == "admin").first()
                if admin_user and admin_user.preferences:
                    prefs = copy.deepcopy(admin_user.preferences)
                    statuses = prefs.get("statuses", {})
                    
                    llm_providers = prefs.get("llm", {}).get("providers", {})
                        
                    changed = False
                    for p_name, p_data in list(llm_providers.items()):
                        api_key = p_data.get("api_key")
                        if not api_key or "***" in api_key or api_key.lower() in ["none", ""]:
                            continue
                            
                        status_key = f"llm_{p_name}"
                        new_status = "error"
                        try:
                            llm = get_llm_provider(
                                provider_name=p_name,
                                model_name=p_data.get("model") or "",
                                api_key_override=api_key
                            )
                            # Short timeout to avoid blocking the loop heavily
                            async def test_llm():
                                return await llm.acompletion(prompt="Hello")
                            await asyncio.wait_for(test_llm(), timeout=10.0)
                            new_status = "success"
                        except Exception as e:
                            logger.error(f"[Health Check] LLM {p_name} background test failed: {e}")
                            new_status = "error"
                            
                        if statuses.get(status_key) != new_status:
                            statuses[status_key] = new_status
                            changed = True
                            
                    if changed:
                        prefs["statuses"] = statuses
                        admin_user.preferences = prefs
                        flag_modified(admin_user, "preferences")
                        db.commit()
                        logger.info("[Health Check] System LLM statuses updated.")
                        
        except Exception as e:
            logger.error(f"[Health Check] Background task error: {e}")
        
        await asyncio.sleep(300) # Re-check every 5 minutes

def create_app() -> FastAPI:
    """
    Factory function to create and configure the FastAPI application.
    This encapsulates all setup logic, making the main entry point clean.
    """
    app = FastAPI(
        # Use metadata from the central settings
        title=settings.PROJECT_NAME,
        version=settings.VERSION,
        description="A modular API to route requests to various LLMs with RAG capabilities.",
        lifespan=lifespan
    )

    logging.basicConfig(level=settings.LOG_LEVEL, format='%(asctime)s - %(levelname)s - %(message)s')
    
    # Global settings for LiteLLM to handle provider-specific quirks
    litellm.drop_params = True

    # --- Initialize Core Services defensively ---
    
    # RAG Components are optional for now as requested
    embedder = None
    vector_store = None
    retrievers = []
    
    try:
        # Resolve from config/settings
        if settings.EMBEDDING_PROVIDER:
            embedder = get_embedder_from_config(
                provider=settings.EMBEDDING_PROVIDER,
                dimension=settings.EMBEDDING_DIMENSION,
                model_name=settings.EMBEDDING_MODEL_NAME,
                api_key=settings.EMBEDDING_API_KEY
            )
            
            vector_store = FaissVectorStore(
                index_file_path=settings.FAISS_INDEX_PATH,
                dimension=settings.EMBEDDING_DIMENSION,
                embedder=embedder
            )
            app.state.vector_store = vector_store
            retrievers.append(FaissDBRetriever(vector_store=vector_store))
    except Exception as e:
        logger.error(f"Failed to initialize Vector Store: {e}. RAG functionality might be restricted.")

    # Voice Providers (optional fallback)
    tts_provider = None
    stt_provider = None
    try:
        if settings.TTS_PROVIDER:
            tts_provider = get_tts_provider(
                provider_name=settings.TTS_PROVIDER,
                api_key=settings.TTS_API_KEY,
                model_name=settings.TTS_MODEL_NAME,
                voice_name=settings.TTS_VOICE_NAME
            )
        if settings.STT_PROVIDER:
            stt_provider = get_stt_provider(
                provider_name=settings.STT_PROVIDER,
                api_key=settings.STT_API_KEY,
                model_name=settings.STT_MODEL_NAME
            )
    except ValueError as e:
        logger.info(f"TTS/STT will be initialized later via UI: {e}")
    except Exception as e:
        logger.warning(f"Failed to initialize TTS/STT: {e}")

    prompt_service = PromptService()
    # 9. Initialize the Service Container with all initialized services
    services = ServiceContainer()
    services.with_service("settings", service=lambda: settings)
    services.with_document_service(vector_store=vector_store)
    
    node_registry_service = NodeRegistryService()
    services.with_service("node_registry_service", service=node_registry_service)
    
    # Initialize Browser Service Client
    browser_endpoint = settings.BROWSER_SERVICE_ENDPOINT
    browser_service = BrowserServiceClient(endpoint=browser_endpoint)
    services.with_service("browser_service", service=browser_service)

    tool_service = ToolService(services=services)
    services.with_service("tool_service", service=tool_service)

    services.with_rag_service(
        retrievers=retrievers, 
        prompt_service=prompt_service, 
        tool_service=tool_service,
        node_registry_service=node_registry_service
    )
    
    services.with_service("stt_service", service=STTService(stt_provider=stt_provider))
    services.with_service("tts_service", service=TTSService(tts_provider=tts_provider))
        
    services.with_service("prompt_service", service=prompt_service)
    services.with_service("session_service", service=SessionService(services=services))
    services.with_service("user_service", service=UserService())

    from app.core.services.mesh import MeshService
    services.with_service("mesh_service", service=MeshService(services=services))

    from app.core.services.auth import AuthService
    services.with_service("auth_service", service=AuthService(services=services))

    from app.core.services.preference import PreferenceService
    services.with_service("preference_service", service=PreferenceService(services=services))

    # Area 4: Initialize Agent Scheduler
    agent_scheduler = AgentScheduler(services=services)
    services.with_service("agent_scheduler", service=agent_scheduler)

    from app.core.services.agent import AgentService
    services.with_service("agent_service", service=AgentService(services=services))

    app.state.services = services

    # Create and include the API router, injecting the service
    # Assignment #5: Enforce uniform /api/v1 prefix for all operations
    api_router = create_api_router(services=services)
    app.include_router(api_router, prefix="/api/v1")

    # ── MCP Well-Known Discovery (root level, no /api/v1 prefix) ──────────────
    # External MCP clients (Claude Desktop, IDE plugins) look for this at the root.
    from fastapi.responses import JSONResponse as _JSONResponse
    @app.get("/.well-known/mcp/manifest.json", include_in_schema=False, tags=["MCP"])
    async def _mcp_manifest_root():
        from app.config import settings as _s
        base_url = _s.HUB_PUBLIC_URL or "https://ai.jerxie.com"
        return _JSONResponse({
            "schema_version": "1.0",
            "name": "Cortex Hub",
            "description": "AI-powered swarm orchestration platform with distributed agent nodes, skill management, and multi-provider LLM routing.",
            "version": "1.0.0",
            "server_url": f"{base_url}/api/v1/mcp/sse",
            "transport": "sse",
        }, headers={"Access-Control-Allow-Origin": "*"})

    cors_origins = settings.CORS_ORIGINS
    if settings.HUB_PUBLIC_URL and settings.HUB_PUBLIC_URL not in cors_origins:
        cors_origins.append(settings.HUB_PUBLIC_URL)

    app.add_middleware(
        CORSMiddleware,
        allow_origins=cors_origins,
        allow_credentials=True,
        allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
        allow_headers=["Content-Type", "Authorization"],
    )
    return app
