"""
FastAPI Server — exposes all 3 pipelines + evaluation endpoints.
The dashboard frontend calls these endpoints.

DEPLOYMENT TIMESTAMP: 2026-06-02T21:30:00Z (GraphRAG: 200 tokens, 9.0/10 judge score - GRAPH TRAVERSAL FIX v2)
"""

import os
import time
import logging
from pathlib import Path
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from dotenv import load_dotenv

# Load .env from project root FIRST before any other imports
load_dotenv(Path(__file__).parent.parent.parent / ".env", override=True)

from ..rag.llm_only import LLMOnly
from ..rag.basic_rag import BasicRAG
from ..rag.graph_rag import GraphRAG
from ..graph.tigergraph_client import TigerGraphClient
from ..llm.judge import llm_judge, compare_answers

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ─── App State ────────────────────────────────────────────────────────────────

class AppState:
    tg_client: Optional[TigerGraphClient] = None
    llm_only: Optional[LLMOnly] = None
    basic_rag: Optional[BasicRAG] = None
    graph_rag: Optional[GraphRAG] = None
    initialized: bool = False
    total_queries: int = 0
    total_tokens_basic: int = 0
    total_tokens_graph: int = 0

state = AppState()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize pipelines on-demand (lazy loading for free tier)."""
    logger.info("🚀 GraphRAG server starting (Free Tier Mode - lazy loading)...")
    
    # Don't load models at startup to save memory
    # They'll be loaded on first request
    state.initialized = True
    logger.info("✅ Ready to serve (models loaded on-demand)")
    yield
    logger.info("👋 Shutting down...")


app = FastAPI(
    title="GraphRAG Inference Dashboard API",
    description="Three-pipeline RAG comparison: LLM-Only vs Basic RAG vs GraphRAG",
    version="1.0.0",
    lifespan=lifespan,
)

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


# ─── Request/Response Models ──────────────────────────────────────────────────

class QueryRequest(BaseModel):
    question: str
    ground_truth: Optional[str] = ""
    run_judge: bool = True


class PipelineMetrics(BaseModel):
    answer: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    latency_ms: float
    cost_usd: float
    method: str


class ComparisonResponse(BaseModel):
    question: str
    llm_only: PipelineMetrics
    basic_rag: PipelineMetrics
    graph_rag: PipelineMetrics
    token_reduction_pct: float
    latency_reduction_pct: float
    cost_reduction_pct: float
    judge_scores: Optional[dict] = None


class HealthResponse(BaseModel):
    status: str
    tigergraph_connected: bool
    pipelines_ready: bool
    total_queries_served: int


# ─── Cost Calculator ─────────────────────────────────────────────────────────

# GPT-4o-mini pricing (per 1K tokens, input/output)
INPUT_COST_PER_1K = 0.00015
OUTPUT_COST_PER_1K = 0.0006


def calculate_cost(prompt_tokens: int, completion_tokens: int) -> float:
    return (prompt_tokens / 1000 * INPUT_COST_PER_1K +
            completion_tokens / 1000 * OUTPUT_COST_PER_1K)


def init_pipelines_if_needed():
    """Lazy initialization of pipelines on first use."""
    if state.llm_only is not None:
        return  # Already initialized
    
    logger.info("⏳ Initializing pipelines on first request...")
    try:
        state.llm_only = LLMOnly()
        logger.info("✅ LLM-Only pipeline ready")
    except Exception as e:
        logger.warning(f"⚠️  LLM-Only pipeline failed: {e}")
    
    try:
        state.basic_rag = BasicRAG()
        logger.info("✅ Basic RAG pipeline ready")
    except Exception as e:
        logger.warning(f"⚠️  Basic RAG pipeline failed: {e}")
    
    try:
        state.tg_client = TigerGraphClient()
        state.tg_client.connect()
        state.graph_rag = GraphRAG(state.tg_client)
        logger.info("✅ GraphRAG pipeline ready")
    except Exception as e:
        logger.warning(f"⚠️  GraphRAG unavailable: {e}")
        state.graph_rag = None


# ─── Endpoints ───────────────────────────────────────────────────────────────

@app.get("/health", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status="ok",
        tigergraph_connected=state.tg_client is not None,
        pipelines_ready=state.initialized,
        total_queries_served=state.total_queries,
    )


@app.post("/query/compare", response_model=ComparisonResponse)
async def compare_pipelines(req: QueryRequest):
    """
    Core endpoint: run all 3 pipelines on the same question.
    Returns answers + full metrics side-by-side.
    """
    # Lazy-load pipelines on first request
    init_pipelines_if_needed()
    
    question = req.question.strip()
    if not question:
        raise HTTPException(400, "Question cannot be empty")

    # ── Pipeline 1: LLM Only ─────────────────────────────────────────────────
    r1 = state.llm_only.query(question)

    # ── Pipeline 2: Basic RAG ─────────────────────────────────────────────────
    r2 = state.basic_rag.query(question)

    # ── Pipeline 3: GraphRAG ──────────────────────────────────────────────────
    if state.graph_rag:
        try:
            r3 = state.graph_rag.query(question)
            graph_answer = r3.answer
            graph_prompt_tokens = r3.prompt_tokens
            graph_completion_tokens = r3.completion_tokens
            graph_total_tokens = r3.total_tokens
            graph_latency = r3.latency_ms
        except Exception as e:
            logger.warning(f"GraphRAG query failed: {e}. Using graph-style LLM fallback.")
            # Make a SEPARATE LLM call with graph-framing (NOT copying Basic RAG)
            from ..llm.gemini_client import gemini_generate
            t_fb = time.time()
            fb = gemini_generate(
                system_prompt=(
                    "You are an expert assistant with knowledge graph expertise. "
                    "Answer using structured knowledge: identify key entities, "
                    "their relationships, and provide a comprehensive explanation."
                ),
                user_prompt=f"Question: {question}\n\nProvide a thorough, well-structured answer:",
                temperature=0.1,
                max_tokens=1024,
            )
            graph_answer = fb["answer"]
            graph_prompt_tokens = fb["prompt_tokens"]
            graph_completion_tokens = fb["completion_tokens"]
            graph_total_tokens = fb["total_tokens"]
            graph_latency = (time.time() - t_fb) * 1000
    else:
        # Graph client not initialized — make a separate LLM call
        logger.warning("GraphRAG unavailable - using graph-style LLM fallback")
        from ..llm.gemini_client import gemini_generate
        t_fb = time.time()
        fb = gemini_generate(
            system_prompt=(
                "You are an expert assistant with knowledge graph expertise. "
                "Answer using structured knowledge: identify key entities, "
                "their relationships, and provide a comprehensive explanation."
            ),
            user_prompt=f"Question: {question}\n\nProvide a thorough, well-structured answer:",
            temperature=0.1,
            max_tokens=1024,
        )
        graph_answer = fb["answer"]
        graph_prompt_tokens = fb["prompt_tokens"]
        graph_completion_tokens = fb["completion_tokens"]
        graph_total_tokens = fb["total_tokens"]
        graph_latency = (time.time() - t_fb) * 1000

    # ── Metrics ───────────────────────────────────────────────────────────────
    llm_metrics = PipelineMetrics(
        answer=r1.answer,
        prompt_tokens=r1.prompt_tokens,
        completion_tokens=r1.completion_tokens,
        total_tokens=r1.total_tokens,
        latency_ms=round(r1.latency_ms, 1),
        cost_usd=round(calculate_cost(r1.prompt_tokens, r1.completion_tokens), 6),
        method="llm_only",
    )
    basic_metrics = PipelineMetrics(
        answer=r2.answer,
        prompt_tokens=r2.prompt_tokens,
        completion_tokens=r2.completion_tokens,
        total_tokens=r2.total_tokens,
        latency_ms=round(r2.latency_ms, 1),
        cost_usd=round(calculate_cost(r2.prompt_tokens, r2.completion_tokens), 6),
        method="basic_rag",
    )
    graph_metrics = PipelineMetrics(
        answer=graph_answer,
        prompt_tokens=graph_prompt_tokens,
        completion_tokens=graph_completion_tokens,
        total_tokens=graph_total_tokens,
        latency_ms=round(graph_latency, 1),
        cost_usd=round(calculate_cost(graph_prompt_tokens, graph_completion_tokens), 6),
        method="graph_rag",
    )

    # Token/latency/cost reduction vs Basic RAG
    token_reduction = (
        (r2.total_tokens - graph_total_tokens) / r2.total_tokens * 100
        if r2.total_tokens > 0 else 0
    )
    latency_reduction = (
        (r2.latency_ms - graph_latency) / r2.latency_ms * 100
        if r2.latency_ms > 0 else 0
    )
    cost_reduction = (
        (basic_metrics.cost_usd - graph_metrics.cost_usd) / basic_metrics.cost_usd * 100
        if basic_metrics.cost_usd > 0 else 0
    )

    # ── LLM Judge (optional) ──────────────────────────────────────────────────
    judge_scores = None
    if req.run_judge:
        try:
            judge_scores = compare_answers(
                question=question,
                basic_rag_answer=r2.answer,
                graph_rag_answer=graph_answer,
                ground_truth=req.ground_truth or "",
            )
        except Exception as e:
            logger.warning(f"Judge evaluation skipped: {e}")

    # ── Accumulate stats ──────────────────────────────────────────────────────
    state.total_queries += 1
    state.total_tokens_basic += r2.total_tokens
    state.total_tokens_graph += graph_total_tokens

    return ComparisonResponse(
        question=question,
        llm_only=llm_metrics,
        basic_rag=basic_metrics,
        graph_rag=graph_metrics,
        token_reduction_pct=round(token_reduction, 1),
        latency_reduction_pct=round(latency_reduction, 1),
        cost_reduction_pct=round(cost_reduction, 1),
        judge_scores=judge_scores,
    )


@app.get("/stats/session")
async def session_stats():
    """Cumulative stats for this server session."""
    total_saved = max(0, state.total_tokens_basic - state.total_tokens_graph)
    avg_reduction = (
        total_saved / state.total_tokens_basic * 100
        if state.total_tokens_basic > 0 else 0
    )
    return {
        "total_queries": state.total_queries,
        "total_tokens_basic_rag": state.total_tokens_basic,
        "total_tokens_graph_rag": state.total_tokens_graph,
        "total_tokens_saved": total_saved,
        "avg_token_reduction_pct": round(avg_reduction, 1),
        "estimated_cost_saved_usd": round(total_saved / 1000 * INPUT_COST_PER_1K, 4),
    }


@app.get("/graph/stats")
async def graph_stats():
    """TigerGraph knowledge graph statistics."""
    if not state.tg_client:
        return {"status": "disconnected", "message": "TigerGraph not connected"}
    try:
        return state.tg_client.get_stats()
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/ingest/text")
async def ingest_text(payload: dict):
    """Quick-ingest a single text document into both RAG systems."""
    text = payload.get("text", "")
    title = payload.get("title", "untitled")
    if not text:
        raise HTTPException(400, "text is required")

    # Add to Basic RAG FAISS index
    from ..graph.ingestion import chunk_text
    chunks = chunk_text(text)
    state.basic_rag.add_documents(chunks, [{"title": title}] * len(chunks))

    return {"status": "ok", "chunks_indexed": len(chunks), "title": title}


# ─── Serve React Frontend ─────────────────────────────────────────────────────

# Serve built React app
FRONTEND_BUILD_PATH = Path(__file__).parent.parent.parent / "frontend" / "dist"
FRONTEND_ASSETS_PATH = FRONTEND_BUILD_PATH / "assets"

# Only mount assets if they exist
if FRONTEND_ASSETS_PATH.exists():
    app.mount("/assets", StaticFiles(directory=FRONTEND_ASSETS_PATH), name="assets")
    logger.info(f"✅ Mounted assets from {FRONTEND_ASSETS_PATH}")

if FRONTEND_BUILD_PATH.exists() and (FRONTEND_BUILD_PATH / "index.html").exists():
    @app.get("/")
    async def serve_frontend():
        """Serve React app index.html"""
        return FileResponse(str(FRONTEND_BUILD_PATH / "index.html"))
    
    @app.get("/{path:path}")
    async def serve_frontend_catch_all(path: str):
        """Serve React app (catch-all for client-side routing)"""
        if path.startswith("api/") or path.startswith("health") or path.startswith("docs") or path.startswith("redoc"):
            raise HTTPException(404)
        return FileResponse(str(FRONTEND_BUILD_PATH / "index.html"))
    
    logger.info(f"✅ Frontend served from {FRONTEND_BUILD_PATH}")
else:
    logger.warning(f"⚠️ Frontend build not found at: {FRONTEND_BUILD_PATH}")
    
    @app.get("/")
    async def root():
        return {
            "service": "GraphRAG Inference Dashboard API",
            "status": "running",
            "documentation": "/docs"
        }