"""
Helix API Server: Unified FastAPI application.

Combines LightRAG retrieval, Graphiti temporal KG, multi-hop reasoning,
and hallucination detection into a single production-ready REST API.

Usage:
    helix serve                          # starts on 0.0.0.0:9621
    helix serve --port 8000 --reload     # dev mode
    uvicorn helix.api.server:app         # direct uvicorn
"""

from __future__ import annotations

import os
import asyncio
import logging
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any

from dotenv import load_dotenv

load_dotenv(dotenv_path=".env", override=False)

logger = logging.getLogger("helix.api")


# ============================================================
# Global Helix instance (initialized on startup, torn down on shutdown)
# ============================================================
_helix_instance = None


def _get_helix():
    """Return the running Helix instance or raise."""
    if _helix_instance is None:
        raise RuntimeError("Helix not initialized. Server is still starting up.")
    return _helix_instance


# ============================================================
# FastAPI app factory
# ============================================================

def create_app() -> "FastAPI":
    """
    Create and return the configured FastAPI application.

    Uses a lifespan context manager to initialize Helix on startup
    and tear it down on shutdown — this is the modern FastAPI pattern
    (replaces the deprecated on_event decorators).
    """
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        """Startup: init Helix.  Shutdown: close connections."""
        global _helix_instance
        from helix.core import Helix, HelixConfig

        logger.info("Helix API starting — initializing pipeline...")
        config = HelixConfig()
        _helix_instance = Helix(config=config)
        await _helix_instance.initialize()
        logger.info("Helix pipeline ready.")
        yield
        logger.info("Helix API shutting down...")
        if _helix_instance:
            await _helix_instance.close()
        _helix_instance = None
        logger.info("Helix shutdown complete.")

    app = FastAPI(
        title="Helix API",
        description=(
            "Unified REST API for the Helix Temporal GraphRAG pipeline. "
            "Combines LightRAG retrieval, Graphiti temporal knowledge graphs, "
            "multi-hop reasoning, and hallucination detection."
        ),
        version=_get_version(),
        lifespan=lifespan,
    )

    # --- CORS (permissive for dev; tighten for production) ---
    origins = os.getenv("HELIX_CORS_ORIGINS", "*").split(",")
    app.add_middleware(
        CORSMiddleware,
        allow_origins=origins,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # --- Register route groups ---
    _register_routes(app)

    return app


def _get_version() -> str:
    try:
        from helix import __version__
        return __version__
    except Exception:
        return "0.0.0"


# ============================================================
# Route registration
# ============================================================

def _register_routes(app: "FastAPI"):
    """Register all API routes on the FastAPI app."""
    from fastapi import HTTPException
    from fastapi.responses import StreamingResponse
    from helix.api.models import (
        QueryRequest,
        QueryResponse,
        TemporalQueryRequest,
        MultiHopQueryRequest,
        MultiHopQueryResponse,
        InsertRequest,
        InsertWithMetadataRequest,
        InsertBatchRequest,
        InsertResponse,
        ValidateRequest,
        ValidateResponse,
        HealthResponse,
        KnowledgeGraphRequest,
        KnowledgeGraphResponse,
    )

    # ----------------------------------------------------------
    # Health & Info
    # ----------------------------------------------------------

    @app.get("/health", response_model=HealthResponse, tags=["System"])
    async def health_check():
        """Health check — confirms the pipeline is up and connected."""
        helix = _get_helix()
        components = {
            "lightrag": "connected" if helix.lightrag else "disconnected",
            "graphiti": "connected" if helix.graphiti else "disconnected",
        }
        return HealthResponse(
            status="ok",
            version=_get_version(),
            components=components,
        )

    # ----------------------------------------------------------
    # Standard Query
    # ----------------------------------------------------------

    @app.post("/query", response_model=QueryResponse, tags=["Query"])
    async def query(req: QueryRequest):
        """
        Standard RAG query.

        Retrieves context from both the Knowledge Graph (Neo4j) and
        the Vector DB (Supabase) and generates an answer using the
        configured LLM.
        """
        helix = _get_helix()
        try:
            result = await helix.query(
                query_text=req.query,
                mode=req.mode,
                top_k=req.top_k,
            )
            return QueryResponse(
                answer=result.get("answer", ""),
                query=result.get("query", req.query),
                mode=result.get("mode", req.mode),
                temporal_context=result.get("temporal_context"),
                valid_at=result.get("valid_at"),
            )
        except Exception as e:
            logger.exception("Query failed")
            raise HTTPException(status_code=500, detail=str(e))

    # ----------------------------------------------------------
    # Temporal Query
    # ----------------------------------------------------------

    @app.post("/query/temporal", response_model=QueryResponse, tags=["Query"])
    async def query_temporal(req: TemporalQueryRequest):
        """
        Time-aware query.

        Searches the Graphiti temporal knowledge graph filtered to
        facts valid at a specific point in time, or within a time range.
        Useful for questions like *"What was our pricing in Q3 2024?"*.
        """
        helix = _get_helix()
        try:
            # Use the TemporalQueryHandler for specialized filtering
            from helix.temporal import TemporalQueryHandler
            from helix.utils.temporal_utils import TemporalParams

            handler = TemporalQueryHandler(graphiti=helix.graphiti)

            # Build temporal params from the request
            params = TemporalParams(
                has_temporal_intent=True,
                valid_at=req.valid_at,
                valid_from=req.valid_from,
                valid_to=req.valid_to,
                temporal_keywords=[],
            )

            # Search temporal graph
            graph_results = await handler.search_temporal_graph(
                query=req.query,
                valid_at=req.valid_at,
                num_results=req.top_k,
            )

            # Apply temporal filtering
            filtered = await handler.apply_temporal_filter(graph_results, params)

            # Also run LightRAG retrieval for vector context
            from lightrag.base import QueryParam
            lightrag_answer = await helix.lightrag.aquery(
                req.query,
                param=QueryParam(mode=req.mode, top_k=req.top_k),
            )

            return QueryResponse(
                answer=lightrag_answer,
                query=req.query,
                mode=req.mode,
                temporal_context=filtered,
                valid_at=req.valid_at.isoformat() if req.valid_at else "current",
            )
        except Exception as e:
            logger.exception("Temporal query failed")
            raise HTTPException(status_code=500, detail=str(e))

    # ----------------------------------------------------------
    # Multi-Hop Query
    # ----------------------------------------------------------

    @app.post("/query/multihop", response_model=MultiHopQueryResponse, tags=["Query"])
    async def query_multihop(req: MultiHopQueryRequest):
        """
        Multi-hop reasoning query.

        Explores the knowledge graph using BFS-based path discovery
        to find hidden connections across multiple entity hops.
        Ideal for *"How is X connected to Y?"* style questions.
        """
        helix = _get_helix()
        try:
            from helix.multihop import MultiHopRetriever

            retriever = MultiHopRetriever(
                graphiti=helix.graphiti,
                max_hops=req.max_hops,
                max_paths=req.max_paths,
            )

            paths = await retriever.find_paths(
                query=req.query,
                seed_entities=req.seed_entities,
            )

            # Build a combined context from discovered paths
            path_dicts = []
            reasoning_context = []
            for p in paths:
                path_dicts.append({
                    "nodes": p.nodes,
                    "edges": p.edges,
                    "score": p.total_score,
                    "hops": p.hop_count,
                })
                edge_facts = [e.get("fact", "") for e in p.edges if e.get("fact")]
                reasoning_context.extend(edge_facts)

            # Also run standard retrieval for a combined answer
            from lightrag.base import QueryParam
            lightrag_answer = await helix.lightrag.aquery(
                req.query,
                param=QueryParam(mode=req.mode, top_k=req.top_k),
            )

            # Collect seed entities used
            used_seeds = list({n for p in paths for n in p.nodes[:1]}) if paths else []

            return MultiHopQueryResponse(
                answer=lightrag_answer,
                query=req.query,
                mode=req.mode,
                paths=path_dicts,
                seed_entities=used_seeds,
            )
        except Exception as e:
            logger.exception("Multi-hop query failed")
            raise HTTPException(status_code=500, detail=str(e))

    # ----------------------------------------------------------
    # Insert
    # ----------------------------------------------------------

    @app.post("/insert", response_model=InsertResponse, tags=["Ingest"])
    async def insert(req: InsertRequest):
        """
        Insert a document into the RAG pipeline.

        Runs the document through LightRAG's chunking + embedding pipeline
        and writes the extracted entities/relations to Neo4j.
        """
        helix = _get_helix()
        try:
            await helix.lightrag.ainsert(req.document)
            return InsertResponse(
                doc_id=req.doc_id or "auto",
                status="success",
                message="Document inserted into LightRAG pipeline.",
            )
        except Exception as e:
            logger.exception("Insert failed")
            raise HTTPException(status_code=500, detail=str(e))

    @app.post("/insert/with-metadata", response_model=InsertResponse, tags=["Ingest"])
    async def insert_with_metadata(req: InsertWithMetadataRequest):
        """
        Insert a document with temporal metadata.

        In addition to standard RAG ingestion, this endpoint records
        `reference_time` and `group_id` so that the knowledge graph
        retains *when* the facts were true and *which project* they
        belong to.
        """
        helix = _get_helix()
        try:
            result = await helix.insert(
                document=req.document,
                doc_id=req.doc_id,
                source_description=req.source_description,
                reference_time=req.reference_time,
                group_id=req.group_id,
            )
            return InsertResponse(
                doc_id=result.get("doc_id", ""),
                status="success",
                entities_extracted=result.get("entities_extracted", 0),
                relations_extracted=result.get("relations_extracted", 0),
                entity_names=result.get("entity_names", []),
                message="Document inserted with temporal metadata.",
            )
        except Exception as e:
            logger.exception("Insert with metadata failed")
            raise HTTPException(status_code=500, detail=str(e))

    @app.post("/insert/batch", response_model=list[InsertResponse], tags=["Ingest"])
    async def insert_batch(req: InsertBatchRequest):
        """Insert multiple documents in a single request."""
        helix = _get_helix()
        responses: list[InsertResponse] = []
        for doc_req in req.documents:
            try:
                result = await helix.insert(
                    document=doc_req.document,
                    doc_id=doc_req.doc_id,
                    source_description=doc_req.source_description,
                    reference_time=doc_req.reference_time,
                    group_id=doc_req.group_id,
                )
                responses.append(InsertResponse(
                    doc_id=result.get("doc_id", ""),
                    status="success",
                    entities_extracted=result.get("entities_extracted", 0),
                    relations_extracted=result.get("relations_extracted", 0),
                    entity_names=result.get("entity_names", []),
                    message="Inserted.",
                ))
            except Exception as e:
                responses.append(InsertResponse(
                    doc_id=doc_req.doc_id or "unknown",
                    status="error",
                    message=str(e),
                ))
        return responses

    # ----------------------------------------------------------
    # Validation (Hallucination Detection)
    # ----------------------------------------------------------

    @app.post("/validate", response_model=ValidateResponse, tags=["Validation"])
    async def validate(req: ValidateRequest):
        """
        Validate an LLM response against the knowledge graph.

        Uses the Composite Fidelity Index (CFI) to check entity coverage,
        relation preservation, and temporal consistency of the answer.
        Returns a grounding score and a list of ungrounded claims.
        """
        helix = _get_helix()
        try:
            from helix.hallucination import HallucinationDetector

            detector = HallucinationDetector(
                graphiti=helix.graphiti,
                grounding_threshold=req.grounding_threshold,
            )

            result = await detector.verify_response(
                response=req.response,
                query=req.query,
                context=req.context,
            )

            return ValidateResponse(
                is_grounded=result.is_grounded,
                confidence_score=result.confidence_score,
                entity_coverage=result.entity_coverage,
                relation_coverage=result.relation_coverage,
                temporal_consistency=result.temporal_consistency,
                ungrounded_claims=list(result.ungrounded_claims),
                verification_details=result.verification_details,
            )
        except Exception as e:
            logger.exception("Validation failed")
            raise HTTPException(status_code=500, detail=str(e))

    # ----------------------------------------------------------
    # Knowledge Graph Exploration
    # ----------------------------------------------------------

    @app.post("/graph/explore", response_model=KnowledgeGraphResponse, tags=["Graph"])
    async def graph_explore(req: KnowledgeGraphRequest):
        """
        Explore the knowledge graph around an entity.

        Returns nodes and edges for visualization or debugging.
        """
        helix = _get_helix()
        try:
            result = await helix.get_knowledge_graph(
                entity_name=req.entity_name,
                max_depth=req.max_depth,
                max_nodes=req.max_nodes,
            )
            return KnowledgeGraphResponse(
                nodes=result.get("nodes", []),
                edges=result.get("edges", []),
                is_truncated=result.get("is_truncated", False),
            )
        except Exception as e:
            logger.exception("Graph exploration failed")
            raise HTTPException(status_code=500, detail=str(e))


# ============================================================
# Module-level app instance for `uvicorn helix.api.server:app`
# ============================================================
app = create_app()
