"""Placement Policy Advisor — FastAPI backend + Gradio chat UI (single app).

Flow for every question:
    query -> FastEmbed embedding -> Qdrant top-K search -> prompt with injected
    context -> Hugging Face chat.completions (streaming) with automatic fallback
    to a second model on error/timeout.

The Gradio ChatInterface is mounted on the same FastAPI app at "/", and a
token-protected POST /api/query exposes the same pipeline programmatically.
Everything runs in one process on APP_PORT. All config comes from the
environment (loaded from .env when present); secrets are never defaulted.
"""

from __future__ import annotations

import atexit
import inspect
import logging
import os
import secrets as _secrets

import gradio as gr
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, Header, HTTPException
from fastapi.responses import StreamingResponse
from huggingface_hub import AsyncInferenceClient
from pydantic import BaseModel
from qdrant_client import QdrantClient

from system_prompt import INSUFFICIENT_CONTEXT_MESSAGE, SYSTEM_PROMPT

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | main | %(message)s",
)
log = logging.getLogger("main")

# --- ZeroGPU compatibility shim ------------------------------------------------
# This app does no local GPU work (FastEmbed runs ONNX on CPU; the LLM call goes
# to HF's hosted Inference API over HTTP). But on a Hugging Face "ZeroGPU"
# Space, the platform refuses to serve unless at least one function is
# decorated with @spaces.GPU, so it knows when to allocate/release the shared
# GPU. The `_zerogpu_probe` function below exists solely to satisfy that
# startup scan; it is never called.
#
# IMPORTANT: `spaces` must NOT be listed in requirements.txt. On a real
# ZeroGPU Space, the platform injects its own pinned, instrumented build of
# `spaces` that's wired into its GPU scheduler; if we also pip-install a
# generic version ourselves, it shadows that build with an inert stub that
# never actually registers — the decorator applies but the platform still
# reports "No @spaces.GPU function detected". Elsewhere (local, Docker, any
# non-HF host) the package simply isn't installed, so the import fails and we
# fall back to a no-op decorator.
try:
    import spaces as _hf_spaces

    _gpu_decorator = _hf_spaces.GPU
    log.info("ZeroGPU 'spaces' package detected; @spaces.GPU shim active.")
except Exception as _spaces_import_error:  # noqa: BLE001
    log.info(
        "ZeroGPU 'spaces' package not available (%s) — using no-op GPU decorator "
        "(expected outside a HF ZeroGPU Space).",
        _spaces_import_error,
    )

    def _gpu_decorator(fn):
        return fn


@_gpu_decorator
def _zerogpu_probe() -> None:
    return None


def require_env(name: str) -> str:
    """Return a required secret/config value or fail fast with a clear message."""
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(
            f"Required environment variable '{name}' is not set. "
            f"Add it to .env (local) or your platform secrets (deployment)."
        )
    return value


# --- Secrets (REQUIRED — no defaults) -----------------------------------------
HF_TOKEN = require_env("HF_TOKEN")
BACKEND_API_AUTH_TOKEN = require_env("BACKEND_API_AUTH_TOKEN")

# --- Non-secret operational config (env with build-safe defaults) -------------
LLM_MODEL_ID = os.environ.get("LLM_MODEL_ID", "meta-llama/Llama-3.1-8B-Instruct")
LLM_FALLBACK_MODEL_ID = os.environ.get("LLM_FALLBACK_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
EMBEDDING_MODEL_ID = os.environ.get("EMBEDDING_MODEL_ID", "BAAI/bge-small-en-v1.5")
QDRANT_STORAGE_PATH = os.environ.get("QDRANT_STORAGE_PATH", "./qdrant_storage")
QDRANT_COLLECTION_NAME = os.environ.get("QDRANT_COLLECTION_NAME", "placement_policy")
TOP_K_RESULTS = int(os.environ.get("TOP_K_RESULTS", "5"))
APP_PORT = int(os.environ.get("APP_PORT", "7860"))
FASTEMBED_CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", ".fastembed_cache")

# Generation params (env-overridable; not secrets).
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "1024"))
TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.2"))
LLM_REQUEST_TIMEOUT = float(os.environ.get("LLM_REQUEST_TIMEOUT", "60"))


def _bootstrap_if_needed() -> None:
    """Ensure the vector index exists before serving.

    On a Docker deployment (see Dockerfile), convert_docs/download_models/
    data_indexer already ran at image-build time, so this is a no-op. On a
    plain Python launch (e.g. a Hugging Face Gradio-SDK Space, which has no
    build-time hook), this runs those same steps once, live, at startup.
    """
    probe = QdrantClient(path=QDRANT_STORAGE_PATH)
    try:
        indexed = probe.collection_exists(QDRANT_COLLECTION_NAME) and probe.count(
            QDRANT_COLLECTION_NAME
        ).count > 0
    finally:
        probe.close()

    if indexed:
        log.info("Vector index already present; skipping bootstrap.")
        return

    log.info("No vector index found — running one-time bootstrap ...")
    import convert_docs
    import data_indexer
    import download_models

    for step_name, module in (
        ("convert_docs", convert_docs),
        ("download_models", download_models),
        ("data_indexer", data_indexer),
    ):
        exit_code = module.main()
        if exit_code != 0:
            raise RuntimeError(f"Bootstrap step '{step_name}' failed (exit {exit_code})")
    log.info("Bootstrap complete.")


_bootstrap_if_needed()

# --- Heavy singletons: load once at import ------------------------------------
log.info("Loading FastEmbed model '%s' ...", EMBEDDING_MODEL_ID)
from fastembed import TextEmbedding  # noqa: E402 — after config so cache_dir is set

_embedder = TextEmbedding(model_name=EMBEDDING_MODEL_ID, cache_dir=FASTEMBED_CACHE_DIR)

log.info("Opening Qdrant local storage at '%s' ...", QDRANT_STORAGE_PATH)
_qdrant = QdrantClient(path=QDRANT_STORAGE_PATH)
# Close the embedded store cleanly on exit (avoids a noisy __del__ during
# interpreter shutdown when local mode holds the storage lock).
atexit.register(_qdrant.close)
if not _qdrant.collection_exists(QDRANT_COLLECTION_NAME):
    log.warning(
        "Qdrant collection '%s' does not exist. Run data_indexer.py before "
        "serving, or the assistant will have no context to answer from.",
        QDRANT_COLLECTION_NAME,
    )

_hf_client = AsyncInferenceClient(token=HF_TOKEN, timeout=LLM_REQUEST_TIMEOUT)


# --- Retrieval ----------------------------------------------------------------
def embed_query(query: str) -> list[float]:
    return next(iter(_embedder.embed([query]))).tolist()


def retrieve(query: str):
    """Return the top-K scored points for a query (empty list if none/unavailable)."""
    if not _qdrant.collection_exists(QDRANT_COLLECTION_NAME):
        return []
    vector = embed_query(query)
    response = _qdrant.query_points(
        collection_name=QDRANT_COLLECTION_NAME,
        query=vector,
        limit=TOP_K_RESULTS,
        with_payload=True,
    )
    return response.points


def build_messages(query: str, points) -> list[dict]:
    """Assemble the chat messages with retrieved policy context injected."""
    blocks = []
    for i, point in enumerate(points, start=1):
        payload = point.payload or {}
        label = payload.get("section_path") or payload.get("source") or "policy"
        blocks.append(f"[Source {i} — {label}]\n{payload.get('text', '').strip()}")
    context = "\n\n---\n\n".join(blocks)
    user_content = (
        "POLICY CONTEXT (use only this to answer):\n"
        f"{context}\n\n"
        f"QUESTION: {query}"
    )
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]


# --- LLM streaming with fallback ----------------------------------------------
async def _stream_model(model_id: str, messages: list[dict]):
    """Yield content deltas from one model. Raises on failure."""
    result = _hf_client.chat.completions.create(
        model=model_id,
        messages=messages,
        stream=True,
        max_tokens=MAX_NEW_TOKENS,
        temperature=TEMPERATURE,
    )
    # AsyncInferenceClient returns a coroutine (await -> async iterator) across
    # versions; guard so we work whether or not it must be awaited.
    if inspect.isawaitable(result):
        result = await result
    async for chunk in result:
        delta = chunk.choices[0].delta.content
        if delta:
            yield delta


async def stream_answer(query: str):
    """Full RAG pipeline as an async token stream."""
    query = (query or "").strip()
    if not query:
        yield "Please enter a question about the RV University Placement Policy."
        return

    points = retrieve(query)
    if not points:
        yield INSUFFICIENT_CONTEXT_MESSAGE
        return

    messages = build_messages(query, points)

    emitted = False
    for idx, model_id in enumerate((LLM_MODEL_ID, LLM_FALLBACK_MODEL_ID)):
        try:
            async for delta in _stream_model(model_id, messages):
                emitted = True
                yield delta
            return  # completed successfully
        except Exception as exc:  # noqa: BLE001 — any error triggers fallback
            log.warning("Model '%s' failed (%s): %s", model_id, type(exc).__name__, exc)
            if emitted:
                # Already streamed partial output; can't cleanly restart.
                yield "\n\n_(The response was interrupted. Please ask again.)_"
                return
            # else: fall through to the fallback model
    yield (
        "The assistant is temporarily unavailable because the language model "
        "could not be reached. Please try again shortly."
    )


# --- FastAPI app --------------------------------------------------------------
app = FastAPI(title="Placement Policy Advisor")


class QueryRequest(BaseModel):
    query: str


def _authorize(authorization: str | None, x_api_token: str | None) -> None:
    """Constant-time check of the backend auth token from either header."""
    provided = ""
    if authorization:
        provided = authorization[7:].strip() if authorization.lower().startswith("bearer ") else authorization.strip()
    elif x_api_token:
        provided = x_api_token.strip()

    if not provided:
        raise HTTPException(status_code=401, detail="Missing authentication token.")
    if not _secrets.compare_digest(provided, BACKEND_API_AUTH_TOKEN):
        raise HTTPException(status_code=403, detail="Invalid authentication token.")


@app.get("/api/health")
async def health():
    return {
        "status": "ok",
        "collection": QDRANT_COLLECTION_NAME,
        "indexed": _qdrant.collection_exists(QDRANT_COLLECTION_NAME),
    }


@app.post("/api/query")
async def api_query(
    payload: QueryRequest,
    authorization: str | None = Header(default=None),
    x_api_token: str | None = Header(default=None),
):
    _authorize(authorization, x_api_token)
    return StreamingResponse(stream_answer(payload.query), media_type="text/plain")


# --- Gradio UI (mounted at "/") -----------------------------------------------
async def _chat_fn(message: str, history):
    """Gradio streaming callback — accumulates deltas into the growing answer."""
    answer = ""
    async for delta in stream_answer(message):
        answer += delta
        yield answer


_demo = gr.ChatInterface(
    fn=_chat_fn,
    title="RV University — Placement Policy Advisor",
    description=(
        "Ask about eligibility, pre-placement training, the offer rules "
        "(One Offer 1×, Offer Progression 1.5×, Dream Offer), attempt limits, "
        "and placement procedures. Answers come only from the official policy. "
        "Disciplinary matters are handled by the Student Disciplinary Committee (STDC)."
    ),
    examples=[
        "When does the 1.5× Offer Progression Rule get activated?",
        "How many placement opportunities do I get in a cycle?",
        "What are the eligibility criteria to participate in placements?",
        "What counts as a Dream Offer?",
    ],
)

# Running on a Hugging Face Space? The platform sets SPACE_ID.
ON_HF_SPACE = bool(os.environ.get("SPACE_ID"))

# Only build the combined FastAPI app (Gradio UI + /api/query) for the self-host
# / Docker path. On a HF Gradio-SDK Space we hand serving to Gradio's own
# launcher instead (see __main__), so mounting here is unnecessary.
if not ON_HF_SPACE:
    app = gr.mount_gradio_app(app, _demo, path="/")


if __name__ == "__main__":
    log.info("Starting Placement Policy Advisor on port %d", APP_PORT)
    if ON_HF_SPACE:
        # Hugging Face Gradio-SDK Space: the platform expects the app to be
        # served by Gradio's own launcher and already manages port 7860. Calling
        # our own uvicorn.run() here double-binds that port (our Uvicorn finishes
        # lifespan startup, then fails to bind because HF's server already holds
        # it). ssr_mode=False avoids Gradio 6's separate Node.js SSR server,
        # which is unreliable on a constrained free Space.
        # NOTE: the token-protected REST API (/api/query) is exposed only in the
        # Docker / self-hosted deployment below, not on the Space UI.
        _demo.launch(server_name="0.0.0.0", server_port=APP_PORT, ssr_mode=False)
    else:
        # Self-host / Docker: single-process Uvicorn serving both the Gradio UI
        # (at "/") and the REST API. workers=1 is explicit so a platform-injected
        # WEB_CONCURRENCY can't spawn extra workers that fight over the port.
        uvicorn.run(app, host="0.0.0.0", port=APP_PORT, workers=1)