# app/main.py – FastAPI + Aida AI Agent (Production with ML Integration)
# ============================================================

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from contextlib import asynccontextmanager
import logging
import os

# ---------- core imports ----------
try:
    from app.config import settings
    from app.database import connect_db, disconnect_db, ensure_indexes as ensure_auth_indexes
    from app.routes import auth
    from app.utils.logger import setup_logger
    from app.core.exceptions import AuthException
    setup_logger()
except ImportError as e:
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)
    logger.error(f"Import error: {e}")

    class AuthException(Exception):
        def __init__(self, status_code=500, detail="Error", error_code="ERROR", message="Error"):
            self.status_code = status_code
            self.detail = detail
            self.error_code = error_code
            self.message = message
            self.data = {}

logger = logging.getLogger(__name__)

# ---------- AI imports ----------
from app.ai.routes.chat import router as ai_chat_router
from app.models.listing import ensure_listing_indexes
from app.ai.config import redis_client, qdrant_client
from app.ml.models.ml_listing_extractor import get_ml_extractor  # ✅ NEW

# ====================================================================
# ML Startup Validation - NEW
# ====================================================================
async def validate_ml_startup():
    """Validate ML extractor and models at startup"""
    try:
        ml = get_ml_extractor()
        logger.info("✅ ML Extractor initialized")
        
        # Check if models are trained
        if ml.field_models and ml.field_models.get("location_classifier") is not None:
            logger.info("✅ ML field models loaded")
        else:
            logger.warning("⚠️ ML field models not trained - limited accuracy")
            logger.warning("   Run: python app/ml/training/train_complete_model.py")
        
        # Test currency manager (non-blocking)
        try:
            currency, country, city, conf = await ml.currency_mgr.get_currency_for_location("Lagos")
            if currency:
                logger.info(f"✅ Currency Manager working (Lagos → {currency})")
        except Exception as e:
            logger.warning(f"⚠️ Currency Manager test failed: {e}")
        
        # Check embedder
        if ml.embedder is not None:
            logger.info("✅ Sentence embedder ready")
        else:
            logger.warning("⚠️ Sentence embedder not available")
        
        logger.info("✅ All ML checks passed")
        return True
    
    except Exception as e:
        logger.error("❌ ML Extractor initialization failed", exc_info=e)
        logger.warning("⚠️ Continuing without ML features (degraded mode)")
        return False

# ====================================================================
# Lifespan – non-blocking external services
# ====================================================================
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("🚀 Starting Lojiz Platform + Aida AI with ML Integration")

    # 1. MongoDB – critical, must succeed
    try:
        await connect_db()
        await ensure_auth_indexes()
        await ensure_listing_indexes()
        logger.info("✅ MongoDB connected & indexed")
    except Exception as e:
        logger.critical("❌ MongoDB unavailable – aborting start", exc_info=e)
        raise

    # 2. Redis – optional at boot
    try:
        await redis_client.ping()
        logger.info("✅ Redis connected")
    except Exception as e:
        logger.warning("⚠️ Redis unreachable at start-up (ok for now)", exc_info=e)

    # 3. Qdrant – optional at boot
    try:
        await qdrant_client.get_collections()
        logger.info("✅ Qdrant connected")
    except Exception as e:
        logger.warning("⚠️ Qdrant unreachable at start-up (ok for now)", exc_info=e)
    
    # 4. ML Extractor – optional but recommended
    try:
        ml_ready = await validate_ml_startup()
        if not ml_ready:
            logger.warning("⚠️ Running in degraded mode without ML features")
    except Exception as e:
        logger.error("❌ ML validation failed", exc_info=e)
        logger.warning("⚠️ Continuing without ML features")

    yield

    logger.info("🛑 Shutting down Lojiz Platform")
    try:
        # Clear ML caches
        try:
            ml = get_ml_extractor()
            ml.currency_mgr.clear_cache()
            logger.info("✅ ML caches cleared")
        except:
            pass
        
        await disconnect_db()
        await redis_client.close()
        logger.info("✅ Cleanup complete")
    except Exception as e:
        logger.warning("⚠️ Shutdown warning", exc_info=e)

# ====================================================================
# FastAPI instance
# ====================================================================
app = FastAPI(
    title="Lojiz Platform + Aida AI",
    description="Conversational real-estate agent with secure auth & ML features",
    version="1.0.0",
    lifespan=lifespan,
)

# ====================================================================
# CORS
# ====================================================================
environment = os.getenv("ENVIRONMENT", "development")
is_production = environment == "production"

cors_origins = [
    "https://lojiz.onrender.com",
    "https://lojiz.com",
    "https://www.lojiz.com",
] if is_production else [
    "http://localhost",
    "http://localhost:3000",
    "http://localhost:55211",
    "http://127.0.0.1",
    "http://127.0.0.1:3000",
    "http://127.0.0.1:5000",
    "http://127.0.0.1:8080",
    "http://127.0.0.1:56205",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
    max_age=600,
)

# ====================================================================
# Exception handlers
# ====================================================================
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    logger.error(f"Validation error: {exc}")
    errors = []
    for error in exc.errors():
        field = ".".join(str(loc) for loc in error["loc"][1:])
        errors.append({"field": field, "message": error["msg"]})
    return JSONResponse(
        status_code=400,
        content={
            "success": False,
            "message": "Validation error. Please check your input.",
            "error_code": "VALIDATION_ERROR",
            "errors": errors,
        },
    )

@app.exception_handler(AuthException)
async def auth_exception_handler(request: Request, exc: AuthException):
    logger.warning(f"Auth error [{exc.error_code}]: {exc.message}")
    response = {"success": False, "message": exc.message, "error_code": exc.error_code}
    if exc.data:
        response["data"] = exc.data
    return JSONResponse(status_code=exc.status_code, content=response)

@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unexpected error: {str(exc)}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={
            "success": False,
            "message": "An unexpected error occurred. Please try again later.",
            "error_code": "INTERNAL_SERVER_ERROR",
            "error": str(exc) if not is_production else None,
        },
    )

# ====================================================================
# Routers
# ====================================================================
app.include_router(auth.router, prefix="/api/auth", tags=["Authentication"])
app.include_router(ai_chat_router, prefix="/ai", tags=["Aida AI"])

# ====================================================================
# Health
# ====================================================================
@app.get("/health", tags=["Health"])
async def health_check():
    """Health check endpoint with ML status"""
    try:
        ml = get_ml_extractor()
        ml_ready = ml.field_models.get("location_classifier") is not None if ml.field_models else False
    except:
        ml_ready = False
    
    return {
        "status": "ok",
        "service": "Lojiz Platform + Aida AI",
        "version": "1.0.0",
        "environment": environment,
        "ml_ready": ml_ready,  # ✅ NEW
    }

@app.get("/", tags=["Root"])
async def root():
    return {
        "message": "Welcome to Lojiz Platform + Aida AI",
        "docs": "/docs",
        "health": "/health",
        "environment": environment,
    }

@app.options("/{full_path:path}", include_in_schema=False)
async def options_handler(full_path: str):
    return JSONResponse(status_code=200, content={})

# ====================================================================
# Run: uvicorn app.main:app --reload
# ====================================================================