# filename: api_service.py (ENHANCED FOR HF SPACES)

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import Dict, Any, Optional
import logging
import time

# Import the core logic and model loading functions
from backend_pam import load_agent as load_backend_agent, PAM
from frontend_pam import load_frontend_agent, FrontendPAM

# --- Configure Logging ---
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("PAM_API")

# --- Global Agent Variables ---
backend_agent: Optional[PAM] = None
frontend_agent: Optional[FrontendPAM] = None

# --- Data Models for API Requests ---
class UserInput(BaseModel):
    user_input: str = Field(..., min_length=1, max_length=5000, description="User's message to PAM")
    user_id: Optional[str] = Field(None, description="Optional user identifier for tracking")
    backend_context: Optional[str] = Field(None, description="Optional backend context for frontend responses")

class HealthResponse(BaseModel):
    status: str
    backend_ready: bool
    frontend_ready: bool
    timestamp: str
    message: str

# --- FastAPI Initialization ---
app = FastAPI(
    title="PAM - Privacy-First AI Assistant",
    description="Unified inference service for Frontend (Chat) and Backend (Technical) PAM agents",
    version="2.1.0",
    docs_url="/docs",  # Enable docs for development
    redoc_url="/redoc"
)

# --- CORS Setup (Enhanced for HF Spaces) ---
origins = [
    "https://www.uminur.app",
    "https://api.uminur.app",
    "http://localhost:3000",  # Local development
    "http://localhost:7860",  # HF Spaces default port
    "https://*.hf.space",     # HF Spaces domain
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Open for HF Spaces deployment (restrict in production)
    allow_credentials=True,
    allow_methods=["POST", "GET", "OPTIONS"],
    allow_headers=["*"],
)

# --- Request Timing Middleware ---
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    logger.info(f"{request.method} {request.url.path} - {process_time:.3f}s")
    return response

# --- Global Exception Handler ---
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unhandled exception: {exc}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={
            "error": "Internal server error",
            "message": "PAM encountered an unexpected issue. Please try again.",
            "type": str(type(exc).__name__)
        }
    )

# --- Startup Event ---
@app.on_event("startup")
async def startup_event():
    global backend_agent, frontend_agent
    logger.info("🚀 Starting PAM agents initialization...")

    try:
        # Load Backend PAM (Technical Assistant)
        logger.info("🤓 Loading Backend PAM (Nerdy Lab Assistant)...")
        backend_agent = load_backend_agent()
        if backend_agent:
            logger.info("✅ Backend PAM loaded successfully")
        else:
            logger.error("❌ Backend PAM failed to initialize")

        # Load Frontend PAM (Chat Assistant)
        logger.info("💕 Loading Frontend PAM (Sweet Southern Receptionist)...")
        frontend_agent = load_frontend_agent()
        if frontend_agent:
            logger.info("✅ Frontend PAM loaded successfully")
        else:
            logger.error("❌ Frontend PAM failed to initialize")

        if backend_agent and frontend_agent:
            logger.info("🎉 Both PAM agents initialized successfully!")
        else:
            logger.warning("⚠️ One or both agents failed to initialize - service will run in degraded mode")

    except Exception as e:
        logger.error(f"❌ Critical error during startup: {e}", exc_info=True)

# --- Shutdown Event ---
@app.on_event("shutdown")
async def shutdown_event():
    logger.info("👋 Shutting down PAM service...")

# --- Root Endpoint ---
@app.get("/", tags=["Info"])
async def root():
    return {
        "service": "PAM - Privacy-First AI Assistant",
        "version": "2.1.0",
        "status": "operational",
        "endpoints": {
            "health": "/health",
            "technical": "/ai/technical/",
            "chat": "/ai/chat/",
            "docs": "/docs"
        },
        "message": "Welcome to PAM! Use /ai/chat/ for conversational support or /ai/technical/ for backend analysis."
    }

# --- Health Check ---
@app.get("/health", response_model=HealthResponse, tags=["Status"])
async def health_check():
    """Check the health status of both PAM agents"""
    backend_ok = backend_agent is not None
    frontend_ok = frontend_agent is not None
    
    status = "healthy" if (backend_ok and frontend_ok) else "degraded"
    
    if not backend_ok and not frontend_ok:
        status = "unavailable"
    
    response = HealthResponse(
        status=status,
        backend_ready=backend_ok,
        frontend_ready=frontend_ok,
        timestamp=time.strftime("%Y-%m-%d %H:%M:%S"),
        message=f"Backend PAM: {'✅' if backend_ok else '❌'} | Frontend PAM: {'✅' if frontend_ok else '❌'}"
    )
    
    if status == "unavailable":
        raise HTTPException(
            status_code=503,
            detail="Service Unavailable: Both agents failed to initialize"
        )
    
    return response

# --- Technical Endpoint (Backend PAM) ---
@app.post("/ai/technical/", tags=["Technical"])
async def technical_endpoint(input_data: UserInput) -> Dict[str, Any]:
    """
    Backend PAM - Technical Assistant Endpoint
    Handles: PHI detection, log parsing, compliance checks, SIEM analysis
    Personality: Nerdy, proactive lab assistant
    """
    if backend_agent is None:
        logger.error("Technical endpoint called but Backend PAM is not initialized")
        raise HTTPException(
            status_code=503, 
            detail="🤓 Backend PAM is still warming up... Give me a moment to get the lab equipment ready!"
        )

    try:
        logger.info(f"Technical request: {input_data.user_input[:100]}...")
        
        # Process through Backend PAM
        pam_reply = backend_agent.process_input(input_data.user_input)
        
        # Add metadata
        pam_reply["agent_type"] = "backend"
        pam_reply["personality"] = "nerdy_lab_assistant"
        pam_reply["timestamp"] = time.strftime("%Y-%m-%d %H:%M:%S")
        
        if input_data.user_id:
            pam_reply["user_id"] = input_data.user_id
        
        logger.info("Technical request processed successfully")
        return pam_reply
        
    except Exception as e:
        logger.error(f"Error during technical inference: {e}", exc_info=True)
        raise HTTPException(
            status_code=500, 
            detail="🤔 Oops, I hit a technical snag while processing that. Can you try rephrasing or breaking it into smaller parts?"
        )

# --- Chat Endpoint (Frontend PAM) ---
@app.post("/ai/chat/", tags=["Chat"])
async def chat_endpoint(input_data: UserInput) -> Dict[str, Any]:
    """
    Frontend PAM - Conversational Assistant Endpoint
    Handles: Appointments, resources, health inquiries, general chat
    Personality: Sweet southern receptionist
    """
    if frontend_agent is None:
        logger.error("Chat endpoint called but Frontend PAM is not initialized")
        raise HTTPException(
            status_code=503, 
            detail="💕 Frontend PAM is getting ready to help you, honey. Just a moment!"
        )

    try:
        logger.info(f"Chat request: {input_data.user_input[:100]}...")
        
        # Set user_id if provided
        if input_data.user_id:
            frontend_agent.user_id = input_data.user_id
        
        # Process through Frontend PAM with optional backend context
        pam_reply = frontend_agent.respond(
            user_text=input_data.user_input,
            backend_brief=input_data.backend_context
        )
        
        # Add metadata
        pam_reply["agent_type"] = "frontend"
        pam_reply["personality"] = "sweet_southern_receptionist"
        pam_reply["timestamp"] = time.strftime("%Y-%m-%d %H:%M:%S")
        
        if input_data.user_id:
            pam_reply["user_id"] = input_data.user_id
        
        logger.info("Chat request processed successfully")
        return pam_reply
        
    except Exception as e:
        logger.error(f"Error during chat inference: {e}", exc_info=True)
        raise HTTPException(
            status_code=500, 
            detail="Sorry dear, I'm having a little technical hiccup. Could you try that again for me?"
        )

# --- Unified Endpoint (Both Agents) ---
@app.post("/ai/unified/", tags=["Unified"])
async def unified_endpoint(input_data: UserInput) -> Dict[str, Any]:
    """
    Unified endpoint that intelligently routes to the appropriate PAM agent
    Based on intent detection or explicit routing
    """
    if not backend_agent or not frontend_agent:
        raise HTTPException(
            status_code=503,
            detail="One or both PAM agents are not ready. Please try again shortly."
        )
    
    try:
        user_text = input_data.user_input.lower()
        
        # Determine routing based on keywords
        backend_keywords = ["compliance", "logs", "phi", "parse", "scan", "analyze", "siem", "alert"]
        is_technical = any(keyword in user_text for keyword in backend_keywords)
        
        if is_technical:
            logger.info("Routing to Backend PAM (technical keywords detected)")
            return await technical_endpoint(input_data)
        else:
            logger.info("Routing to Frontend PAM (conversational/support)")
            return await chat_endpoint(input_data)
            
    except Exception as e:
        logger.error(f"Error in unified endpoint: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="Error processing request through unified endpoint")

# --- Metrics Endpoint ---
@app.get("/metrics", tags=["Status"])
async def get_metrics():
    """Basic metrics for monitoring"""
    return {
        "service": "PAM",
        "backend_status": "online" if backend_agent else "offline",
        "frontend_status": "online" if frontend_agent else "offline",
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "uptime": "tracking_not_implemented"  # Can add process start time tracking
    }

# --- Development/Debug Endpoint (Remove in production) ---
@app.get("/debug/test-agents", tags=["Debug"])
async def test_agents():
    """Quick test of both agents (for development only)"""
    results = {
        "backend_test": None,
        "frontend_test": None
    }
    
    if backend_agent:
        try:
            results["backend_test"] = backend_agent.process_input("check compliance")
        except Exception as e:
            results["backend_test"] = {"error": str(e)}
    
    if frontend_agent:
        try:
            results["frontend_test"] = frontend_agent.respond("Hey PAM, how are you?")
        except Exception as e:
            results["frontend_test"] = {"error": str(e)}
    
    return results