from dotenv import load_dotenv
import os
import sys

# Add project root to Python path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# Load environment variables first
load_dotenv()

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel, Field
from src.graph import graph
from src.state import init_state
from src.utils.semantic_cache import get_cached_response, store_in_cache, get_cache_stats
import uuid
from collections import defaultdict
from datetime import datetime, timedelta
import threading
import asyncio
from functools import partial
from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup: Pre-warm LLM and models to avoid first-request timeout"""
    from src.utils.llm_factory import get_llm
    from langchain_core.messages import SystemMessage, HumanMessage
    
    try:
        print("Warming up LLM connection...")
        llm = get_llm(temperature=0.0)
        loop = asyncio.get_running_loop()
        # Use proper message format for llama-3.1
        await loop.run_in_executor(
            None, 
            lambda: llm.invoke([
                SystemMessage(content="You are a helpful assistant."),
                HumanMessage(content="Hello")
            ])
        )
        print("LLM warmed up successfully.")
    except Exception as e:
        print(f"LLM warm-up failed: {e}")
        print("LLM will initialize on first request.")
    
    yield

app = FastAPI(title="Olist Intelligence Layer", version="1.0", lifespan=lifespan)

# Simple in-memory rate limiter (for production, use Redis)
rate_limit_store = defaultdict(list)
rate_limit_lock = threading.Lock()
RATE_LIMIT_REQUESTS = 20  # requests per window
RATE_LIMIT_WINDOW = 60  # seconds

def check_rate_limit(client_ip: str) -> bool:
    """Check if client has exceeded rate limit"""
    with rate_limit_lock:
        now = datetime.now()
        cutoff = now - timedelta(seconds=RATE_LIMIT_WINDOW)
        
        # Clean old entries
        rate_limit_store[client_ip] = [
            timestamp for timestamp in rate_limit_store[client_ip]
            if timestamp > cutoff
        ]
        
        # Check limit
        if len(rate_limit_store[client_ip]) >= RATE_LIMIT_REQUESTS:
            return False
        
        # Add current request
        rate_limit_store[client_ip].append(now)
        return True

class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000)
    session_id: str = None

class ChatResponse(BaseModel):
    response: str
    intent: str = "unknown"  # Default value if intent not set
    session_id: str
    debug: dict = {}

@app.post("/chat", response_model=ChatResponse)
async def chat(req: ChatRequest, request: Request):
    """Process customer query through the intelligence layer"""
    # Rate limiting
    client_ip = request.client.host
    if not check_rate_limit(client_ip):
        raise HTTPException(
            status_code=429,
            detail=f"Rate limit exceeded. Maximum {RATE_LIMIT_REQUESTS} requests per {RATE_LIMIT_WINDOW} seconds."
        )
    
    session_id = req.session_id or str(uuid.uuid4())
    config = {"configurable": {"thread_id": session_id}}

    # Restore prior state from checkpoint to maintain conversation history
    try:
        prior_state = graph.get_state(config).values
        prior_messages = prior_state.get("messages", [])
        prior_intent = prior_state.get("intent")  # Get prior intent for cache check
        prior_context = {
            "last_order_id": prior_state.get("last_order_id"),
            "last_product_id": prior_state.get("last_product_id"),
            "last_seller_id": prior_state.get("last_seller_id"),
            "last_category": prior_state.get("last_category"),
            "session_context": prior_state.get("session_context", {}),
            "compensation_offered": prior_state.get("compensation_offered"),
            "compensation_tier": prior_state.get("compensation_tier"),
            "case_id": prior_state.get("case_id"),
            "escalation_required": prior_state.get("escalation_required"),
        }
    except Exception as e:
        print(f"[DEBUG] No prior state found for session {session_id[:8]}: {e}")
        prior_messages = []
        prior_intent = None
        prior_context = {}

    # Check semantic cache BEFORE graph execution (only for cacheable intents)
    # Skip cache if user has order context (personalized queries)
    has_order_context = prior_context.get("last_order_id") is not None
    if not has_order_context:
        cached_response = get_cached_response(req.message, prior_intent)
        if cached_response:
            print(f"[API] Returning cached response for session {session_id[:8]}")
            return ChatResponse(
                response=cached_response,
                intent=prior_intent or "informational",
                session_id=session_id,
                debug={"cache_hit": True} if request.query_params.get("debug") == "true" else {}
            )

    # Initialize state with current message and restore history
    input_state = init_state(req.message)
    input_state["messages"] = prior_messages  # Restore conversation history
    input_state.update(prior_context)         # Restore entity context

    try:
        # Run blocking graph.invoke in executor with timeout to prevent hung threads
        loop = asyncio.get_running_loop()
        print(f"[DEBUG] Starting graph execution for session {session_id[:8]}: {req.message[:50]}...")
        result = await asyncio.wait_for(
            loop.run_in_executor(None, partial(graph.invoke, input_state, config=config)),
            timeout=45.0  # Circuit breaker: 3 retries × (1+2+4)s backoff + ~15s LLM = ~37s max
        )
        print(f"[DEBUG] Graph execution completed. Intent: {result.get('intent')}")
    except asyncio.TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="Request timed out after 45 seconds. This may be due to high load or an LLM provider issue. Please try again in a moment."
        )
    except Exception as e:
        # Log the full error for debugging
        import traceback
        error_trace = traceback.format_exc()
        print(f"Graph execution error: {error_trace}")
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")

    # Store in semantic cache AFTER successful execution (only for cacheable intents)
    final_response = result.get("final_response", "I'm sorry, I couldn't process that request.")
    result_intent = result.get("intent")
    if not has_order_context and result_intent:
        store_in_cache(req.message, final_response, result_intent)

    # Only include debug info if explicitly requested (security)
    debug_info = {}
    if request.query_params.get("debug") == "true":
        debug_info = {
            "sql_query": result.get("sql_query"),
            "rag_score": result.get("rag_score"),
            "frustration_score": result.get("frustration_score"),
            "is_late_delivery": result.get("is_late_delivery"),
            "compensation_offered": result.get("compensation_offered"),
            "escalation_required": result.get("escalation_required"),
            "escalation_summary": result.get("escalation_summary"),
            "error_log": result.get("error_log", []),
            "cache_hit": False,
        }

    return ChatResponse(
        response=final_response,
        intent=result_intent or "unknown",
        session_id=session_id,
        debug=debug_info
    )

@app.get("/health")
def health():
    """Health check endpoint"""
    return {"status": "ok", "service": "olist-intelligence-layer"}

@app.get("/cache/stats")
def cache_stats():
    """Return semantic cache statistics for monitoring"""
    return get_cache_stats()

@app.get("/escalations/{case_id}")
def get_escalation_summary(case_id: str):
    """Retrieve escalation summary by case ID"""
    import os
    import json
    
    json_path = f"logs/escalations/{case_id}.json"
    
    if not os.path.exists(json_path):
        raise HTTPException(status_code=404, detail=f"Case {case_id} not found")
    
    try:
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error reading case file: {str(e)}")

@app.get("/escalations")
def list_escalations():
    """List all escalation case IDs"""
    import os
    import json
    from pathlib import Path
    
    log_dir = "logs/escalations"
    
    if not os.path.exists(log_dir):
        return {"cases": [], "total": 0}
    
    cases = []
    for filename in os.listdir(log_dir):
        if filename.endswith(".json"):
            try:
                with open(os.path.join(log_dir, filename), "r", encoding="utf-8") as f:
                    data = json.load(f)
                    cases.append({
                        "case_id": data["case_id"],
                        "timestamp": data["timestamp"],
                        "urgency": data["summary"].get("urgency", "unknown"),
                        "customer_issue": data["summary"].get("customer_issue", "")[:100]
                    })
            except Exception as e:
                print(f"Error reading {filename}: {e}")
                continue
    
    # Sort by timestamp descending (newest first)
    cases.sort(key=lambda x: x["timestamp"], reverse=True)
    
    return {"cases": cases, "total": len(cases)}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)