backend/main.py

"""FastAPI entrypoint — runs internally on :8000, fronted by Next.js on :7860."""
from __future__ import annotations

# ── Load .env FIRST so all module-level os.environ.get() calls in local
#    modules (auth, rag, hf_sync) pick up the values when they are imported.
import asyncio
import json
import os
from pathlib import Path
from dotenv import load_dotenv

_env_file = Path(__file__).parent / ".env"
if not _env_file.exists():
    _env_file = Path(__file__).parent.parent / ".env"
load_dotenv(_env_file)

import logging
import shutil
import uuid
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, List, Literal, Optional

from fastapi import Depends, FastAPI, File, Form, HTTPException, Response, UploadFile
from fastapi.responses import StreamingResponse
from huggingface_hub import AsyncInferenceClient
from pydantic import BaseModel, Field

# Local (imported AFTER load_dotenv so their module-level env reads are correct)
import auth
import config
import hf_sync
import personas_store
from rag import RagEngine

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("backend")

TMP_UPLOAD_DIR = config.TMP_UPLOAD_DIR
LANCEDB_DIR = config.LANCEDB_DIR
LLM_MODEL = config.LLM_MODEL
LLM_MODEL_LOWER = config.LLM_MODEL_LOWER
EMBED_MODEL = config.EMBED_MODEL
HF_TOKEN = config.HF_TOKEN
LLM_EXTRA_BODY_JSON = config.LLM_EXTRA_BODY_JSON
LLM_FINAL_ANSWER_EXTRA_BODY_JSON = config.LLM_FINAL_ANSWER_EXTRA_BODY_JSON
SHOW_LLM_REASONING = config.SHOW_LLM_REASONING
LLM_MAX_TOKENS = config.LLM_MAX_TOKENS
LLM_FINAL_ANSWER_MAX_TOKENS = config.LLM_FINAL_ANSWER_MAX_TOKENS

# Globals populated in lifespan.
rag: Optional[RagEngine] = None
llm_client: Optional[AsyncInferenceClient] = None


def _load_json_object_env(raw: str, env_name: str) -> Optional[dict[str, Any]]:
    if not raw:
        return None
    try:
        loaded = json.loads(raw)
    except json.JSONDecodeError as exc:
        logger.warning("Ignoring invalid %s: %s", env_name, exc)
        return None
    if not isinstance(loaded, dict):
        logger.warning("Ignoring %s because it is not a JSON object", env_name)
        return None
    return loaded


def _is_qwen_thinking_model() -> bool:
    return "qwen/qwen3" in LLM_MODEL_LOWER or "qwen3" in LLM_MODEL_LOWER


def _non_thinking_extra_body() -> dict[str, Any]:
    return {
        "chat_template_kwargs": {"enable_thinking": False},
        # Some HF router/provider combinations look for this at the top level.
        "enable_thinking": False,
    }


def _chat_extra_body() -> Optional[dict[str, Any]]:
    configured = _load_json_object_env(LLM_EXTRA_BODY_JSON, "LLM_EXTRA_BODY_JSON")
    if configured is not None:
        return configured

    if _is_qwen_thinking_model():
        return _non_thinking_extra_body()

    return None


def _final_answer_extra_body() -> Optional[dict[str, Any]]:
    configured = _load_json_object_env(
        LLM_FINAL_ANSWER_EXTRA_BODY_JSON,
        "LLM_FINAL_ANSWER_EXTRA_BODY_JSON",
    )
    if configured is not None:
        return configured

    if _is_qwen_thinking_model():
        return _non_thinking_extra_body()

    return None


CHAT_EXTRA_BODY = _chat_extra_body()
FINAL_ANSWER_EXTRA_BODY = _final_answer_extra_body()


def _sync_vector_cache_sync() -> bool:
    if not hf_sync.is_configured():
        return False
    try:
        return hf_sync.upload_vector_store(LANCEDB_DIR, hf_sync.list_remote_pdfs(), EMBED_MODEL)
    except Exception as exc:  # noqa: BLE001
        logger.warning("Vector cache sync failed: %s", exc)
        return False


def _prepare_local_vector_cache_sync() -> tuple[bool, List[str]]:
    """Restore the local LanceDB folder from the dataset when a valid cache exists."""
    if not hf_sync.is_configured():
        return False, []

    remote_pdfs = hf_sync.list_remote_pdfs()
    cache_restored = hf_sync.download_vector_store(LANCEDB_DIR, remote_pdfs, EMBED_MODEL)
    if not cache_restored:
        shutil.rmtree(LANCEDB_DIR, ignore_errors=True)
    return cache_restored, remote_pdfs


def _cold_start_index_sync() -> None:
    if rag is None:
        return

    sync_dir = TMP_UPLOAD_DIR / "_sync"
    sync_dir.mkdir(parents=True, exist_ok=True)
    paths = hf_sync.download_all_pdfs(sync_dir)
    if paths:
        logger.info("Cold-start: indexing %d PDF(s) from dataset", len(paths))
        added = rag.index_many((p, p.name) for p in paths)
        logger.info("Cold-start: %d nodes indexed", added)
        _sync_vector_cache_sync()
    else:
        logger.info("Cold-start: dataset has no PDFs yet")
        _sync_vector_cache_sync()


def _upload_and_index_sync(tmp_path: Path, safe_name: str) -> tuple[int, bool]:
    if rag is None:
        return 0, False

    hf_sync.upload_pdf(tmp_path, safe_name)
    nodes_added = rag.index_pdf(tmp_path, safe_name)
    return nodes_added, _sync_vector_cache_sync()


def _delete_and_sync_sync(filename: str) -> tuple[bool, int, bool]:
    removed = hf_sync.delete_pdf(filename)
    nodes_removed = rag.delete_by_filename(filename) if rag else 0
    cache_synced = _sync_vector_cache_sync() if removed or nodes_removed > 0 else False
    return removed, nodes_removed, cache_synced


@asynccontextmanager
async def lifespan(app: FastAPI):  # noqa: ARG001
    global rag, llm_client
    TMP_UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

    if not HF_TOKEN:
        logger.warning("HF_TOKEN is not set — chat & embeddings will fail until configured.")

    cache_restored = False
    remote_pdfs: List[str] = []
    if hf_sync.is_configured():
        try:
            cache_restored, remote_pdfs = await asyncio.to_thread(_prepare_local_vector_cache_sync)
        except Exception as exc:  # noqa: BLE001
            logger.exception("Vector cache restore failed: %s", exc)

    rag = RagEngine()
    llm_client = AsyncInferenceClient(model=LLM_MODEL, token=HF_TOKEN)

    # Load personas (HF Dataset → local file → bundled defaults).
    personas_store.load()

    # Cold-start sync: pull all PDFs from the linked dataset and rebuild the index.
    try:
        if cache_restored:
            logger.info("Cold-start: restored vector cache from dataset for %d PDF(s)", len(remote_pdfs))
        elif hf_sync.is_configured():
            await asyncio.to_thread(_cold_start_index_sync)
    except Exception as exc:  # noqa: BLE001
        logger.exception("Cold-start sync failed: %s", exc)

    yield


app = FastAPI(title="IamEarthDev RAG", lifespan=lifespan)


# ---------------- Models ----------------
class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str


class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    persona_prompt: str = Field(default="", description="Persona-specific system prompt")
    temperature: float = Field(default=0.4, ge=0.0, le=1.5)
    top_k: int = Field(default=4, ge=1, le=10)


class LoginRequest(BaseModel):
    passcode: str


class DeleteRequest(BaseModel):
    filename: str


# ---------------- Health & files ----------------
@app.get("/health")
async def health() -> dict:
    return {
        "ok": True,
        "model": LLM_MODEL,
        "embed_model": EMBED_MODEL,
        "dataset_id": config.DATASET_ID,
        "admin_passcode_defaulted": config.ADMIN_PASSCODE == "password" and "ADMIN_PASSCODE" not in os.environ,
        "hf_token_configured": bool(HF_TOKEN),
        "chat_extra_body": CHAT_EXTRA_BODY,
        "final_answer_extra_body": FINAL_ANSWER_EXTRA_BODY,
        "show_reasoning": SHOW_LLM_REASONING,
    }


@app.get("/files")
async def list_files(_: None = Depends(auth.require_admin)) -> dict:
    return {"files": hf_sync.list_remote_pdfs()}


@app.get("/documents")
async def list_documents() -> dict:
    """Public read-only list of uploaded course materials."""
    if not hf_sync.is_configured():
        return {"files": []}
    return {"files": hf_sync.list_remote_pdfs()}


# ---------------- Admin auth ----------------
@app.post("/admin/login")
async def admin_login(body: LoginRequest, response: Response) -> dict:
    if not auth.verify_passcode(body.passcode):
        raise HTTPException(status_code=401, detail="Invalid passcode")
    response.set_cookie(
        key=auth.COOKIE_NAME,
        value=auth.issue_token(),
        httponly=True,
        samesite="lax",
        secure=False,  # HF Spaces serves over https at the edge; cookie still works.
        max_age=60 * 60 * 12,
        path="/",
    )
    return {"ok": True}


@app.post("/admin/logout")
async def admin_logout(response: Response) -> dict:
    response.delete_cookie(auth.COOKIE_NAME, path="/")
    return {"ok": True}


@app.get("/admin/me")
async def admin_me(_: None = Depends(auth.require_admin)) -> dict:
    return {"authenticated": True}


# ---------------- Personas ----------------
class PersonaItem(BaseModel):
    id: str
    name: str
    description: str
    prompt: str


@app.get("/personas")
async def get_personas() -> dict:
    """Public — students and the chat UI need this."""
    return {"personas": personas_store.get_all()}


@app.put("/personas")
async def put_personas(
    personas: List[PersonaItem],
    _: None = Depends(auth.require_admin),
) -> dict:
    saved = personas_store.save([p.model_dump() for p in personas])
    return {"personas": saved}


@app.post("/personas/reset")
async def reset_personas(_: None = Depends(auth.require_admin)) -> dict:
    restored = personas_store.reset()
    return {"personas": restored}


# ---------------- Upload / Delete ----------------
@app.post("/upload")
async def upload(
    file: UploadFile = File(...),
    _: None = Depends(auth.require_admin),
) -> dict:
    if not file.filename or not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only .pdf files are accepted.")

    safe_name = Path(file.filename).name  # strip any path components
    tmp_path = TMP_UPLOAD_DIR / f"{uuid.uuid4().hex}_{safe_name}"
    try:
        with tmp_path.open("wb") as out:
            shutil.copyfileobj(file.file, out)

        try:
            nodes_added, vector_cache_synced = await asyncio.to_thread(
                _upload_and_index_sync, tmp_path, safe_name
            )
        except RuntimeError as exc:
            raise HTTPException(status_code=502, detail=str(exc)) from exc
        except Exception as exc:  # noqa: BLE001
            logger.exception("Indexing failed for %s", safe_name)
            raise HTTPException(
                status_code=500,
                detail=f"Uploaded {safe_name} to the dataset, but indexing failed: {exc}",
            ) from exc
    finally:
        tmp_path.unlink(missing_ok=True)

    return {
        "ok": True,
        "filename": safe_name,
        "nodes_added": nodes_added,
        "vector_cache_synced": vector_cache_synced,
    }


@app.post("/delete")
async def delete(body: DeleteRequest, _: None = Depends(auth.require_admin)) -> dict:
    safe_name = Path(body.filename).name
    try:
        removed, nodes_removed, vector_cache_synced = await asyncio.to_thread(
            _delete_and_sync_sync, safe_name
        )
    except RuntimeError as exc:
        raise HTTPException(status_code=502, detail=str(exc)) from exc
    if not removed and nodes_removed == 0:
        raise HTTPException(status_code=404, detail=f"{safe_name} not found")
    return {
        "ok": True,
        "filename": safe_name,
        "nodes_removed": nodes_removed,
        "vector_cache_synced": vector_cache_synced,
    }


# ---------------- Chat (streaming) ----------------
SYSTEM_BASE = (
    "You are a helpful course assistant for students. Use the selected persona "
    "instructions to decide the tone, teaching style, level of directness, and "
    "structure of your reply. Use the provided course material excerpts as your "
    "primary source of truth. If the material does not contain the answer, say so "
    "honestly. Always cite source filenames in square brackets when you use them."
)


def _build_messages(req: ChatRequest, context: str) -> List[dict]:
    persona = req.persona_prompt.strip()
    sys_parts = [SYSTEM_BASE]
    if persona:
        sys_parts.append(
            "Selected persona instructions. Follow these for the response style; "
            "only the grounding and citation rules above take priority:\n"
            + persona
        )
    if context:
        sys_parts.append(
            "Course material excerpts (use these as your primary source of truth):\n\n"
            + context
        )
    msgs: List[dict] = [{"role": "system", "content": "\n\n".join(sys_parts)}]
    for m in req.messages:
        if m.role == "system":
            continue  # we control the system prompt server-side
        msgs.append({"role": m.role, "content": m.content})
    return msgs


@app.post("/chat")
async def chat(req: ChatRequest) -> StreamingResponse:
    if rag is None or llm_client is None:
        raise HTTPException(status_code=503, detail="Backend not ready.")
    if not req.messages:
        raise HTTPException(status_code=400, detail="messages must not be empty.")

    last_user = next((m.content for m in reversed(req.messages) if m.role == "user"), "")
    nodes = await asyncio.to_thread(rag.retrieve, last_user, req.top_k) if last_user else []
    context = rag.format_context(nodes)
    sources = sorted({n.node.metadata.get("source_filename", "unknown") for n in nodes})
    messages = _build_messages(req, context)
    final_answer_messages = [
        *messages,
        {
            "role": "user",
            "content": (
                "You have already reasoned about this request. Now produce the "
                "student-facing reply in the selected persona style, based on the "
                "same course material excerpts. Do not repeat your thinking process. "
                "Do not produce a thinking process, analysis section, or hidden reasoning. "
                "Answer immediately and cite source filenames in square brackets."
            ),
        },
    ]

    async def stream() -> AsyncIterator[bytes]:
        # Emit sources first as a single SSE-style event so the UI can render them.
        yield (json.dumps({"type": "sources", "sources": sources}) + "\n").encode()
        saw_content = False
        saw_reasoning = False
        finish_reason = None
        try:
            primary_stream = await llm_client.chat_completion(
                messages=messages,
                max_tokens=LLM_MAX_TOKENS,
                temperature=req.temperature,
                stream=True,
                extra_body=CHAT_EXTRA_BODY,
            )
            async for chunk in primary_stream:
                delta = ""
                reasoning = ""
                try:
                    finish_reason = chunk.choices[0].finish_reason or finish_reason
                    delta = chunk.choices[0].delta.content or ""
                    reasoning = getattr(chunk.choices[0].delta, "reasoning", "") or ""
                except (AttributeError, IndexError):
                    pass
                if delta:
                    saw_content = True
                    yield (json.dumps({"type": "delta", "text": delta}) + "\n").encode()
                if reasoning:
                    saw_reasoning = True
                    if SHOW_LLM_REASONING:
                        yield (json.dumps({"type": "reasoning", "text": reasoning}) + "\n").encode()
                    elif not saw_content:
                        logger.info(
                            "Model emitted hidden reasoning before visible content; "
                            "switching to final-answer pass early."
                        )
                        close_stream = getattr(primary_stream, "close", None)
                        if callable(close_stream):
                            close_stream()
                        break

            if not saw_content and saw_reasoning:
                yield (
                    json.dumps(
                        {
                            "type": "phase",
                            "value": "finalizing",
                            "message": "Thinking completed. Generating the final answer.",
                        }
                    )
                    + "\n"
                ).encode()
                final_stream = await llm_client.chat_completion(
                    messages=final_answer_messages,
                    max_tokens=LLM_FINAL_ANSWER_MAX_TOKENS,
                    temperature=req.temperature,
                    stream=True,
                    extra_body=FINAL_ANSWER_EXTRA_BODY,
                )
                async for chunk in final_stream:
                    delta = ""
                    try:
                        delta = chunk.choices[0].delta.content or ""
                    except (AttributeError, IndexError):
                        pass
                    if delta:
                        saw_content = True
                        yield (json.dumps({"type": "delta", "text": delta}) + "\n").encode()

            if not saw_content and saw_reasoning:
                detail = "The model finished its reasoning but did not emit a final answer."
                if finish_reason == "length":
                    detail += (
                        " It likely exhausted the completion budget while thinking. "
                        "Increase LLM_MAX_TOKENS if you want a longer first-pass reasoning budget."
                    )
                yield (json.dumps({"type": "error", "message": detail}) + "\n").encode()
            yield (json.dumps({"type": "done"}) + "\n").encode()
        except Exception as exc:  # noqa: BLE001
            logger.exception("LLM stream failed")
            yield (json.dumps({"type": "error", "message": str(exc)}) + "\n").encode()

    return StreamingResponse(
        stream(),
        media_type="application/x-ndjson",
        headers={
            "X-Accel-Buffering": "no",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        },
    )