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
import re
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, 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
from document_parser import (
    is_supported,
    parse_file_text,
    supported_extensions_label,
)
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
LLM_PROVIDER = config.LLM_PROVIDER
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
MARKS_REFUSAL = (
    "I cannot discuss, infer, estimate, or predict marks, grades, scores, "
    "percentages, or whether changes would gain more marks. I can help you "
    "understand the tutor feedback and identify qualitative next steps grounded "
    "in the feedback, rubric, and coursework."
)
MAX_PARSED_TEXT_CHARS = 120_000
MAX_FEEDBACK_CONTEXT_CHARS = 12_000
MAX_COURSEWORK_CONTEXT_CHARS = 8_000
MAX_RUBRIC_CONTEXT_CHARS = 8_000

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_rubrics(), 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]]:
    if not hf_sync.is_configured():
        return False, []

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


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_rubrics(sync_dir)
    if paths:
        logger.info("Cold-start: indexing %d rubric file(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 rubric files 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_rubric(tmp_path, safe_name)
    nodes_added = rag.index_file(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_rubric(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_rubrics: List[str] = []
    if hf_sync.is_configured():
        try:
            cache_restored, remote_rubrics = 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_kwargs: dict[str, Any] = {"model": LLM_MODEL, "token": HF_TOKEN}
    if LLM_PROVIDER:
        llm_client_kwargs["provider"] = LLM_PROVIDER
    llm_client = AsyncInferenceClient(**llm_client_kwargs)

    personas_store.load()

    try:
        if cache_restored:
            logger.info("Cold-start: restored vector cache from dataset for %d rubric file(s)", len(remote_rubrics))
        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="Feedback Chatbot RAG", lifespan=lifespan)


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)
    tutor_feedback_text: str = ""
    coursework_text: str = ""


class LoginRequest(BaseModel):
    passcode: str


class DeleteRequest(BaseModel):
    filename: str


@app.get("/health")
async def health() -> dict:
    return {
        "ok": True,
        "model": LLM_MODEL,
        "provider": LLM_PROVIDER or "auto",
        "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_rubrics()}


@app.get("/documents")
async def list_documents() -> dict:
    if not hf_sync.is_configured():
        return {"files": []}
    return {"files": hf_sync.list_remote_rubrics()}


@app.get("/rubrics")
async def list_rubrics() -> dict:
    if not hf_sync.is_configured():
        return {"files": []}
    return {"files": hf_sync.list_remote_rubrics()}


@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,
        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}


class PersonaItem(BaseModel):
    id: str
    name: str
    description: str
    prompt: str


@app.get("/personas")
async def get_personas() -> dict:
    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}


@app.post("/upload")
async def upload(
    file: UploadFile = File(...),
    _: None = Depends(auth.require_admin),
) -> dict:
    if not file.filename or not is_supported(file.filename):
        raise HTTPException(
            status_code=400,
            detail=f"Supported rubric files: {supported_extensions_label()}.",
        )

    safe_name = Path(file.filename).name
    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("/student/parse-file")
async def parse_student_file(file: UploadFile = File(...)) -> dict:
    if not file.filename or not is_supported(file.filename, for_student=True):
        raise HTTPException(
            status_code=400,
            detail=f"Supported student files: {supported_extensions_label(for_student=True)}.",
        )

    safe_name = Path(file.filename).name
    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:
            text = await asyncio.to_thread(parse_file_text, tmp_path, safe_name)
        except Exception as exc:  # noqa: BLE001
            raise HTTPException(status_code=400, detail=f"Could not parse {safe_name}: {exc}") from exc
    finally:
        tmp_path.unlink(missing_ok=True)

    truncated = len(text) > MAX_PARSED_TEXT_CHARS
    if truncated:
        text = text[:MAX_PARSED_TEXT_CHARS].rstrip()

    return {
        "ok": True,
        "filename": safe_name,
        "text": text,
        "characters": len(text),
        "truncated": truncated,
    }


@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,
    }


def _words(text: str) -> set[str]:
    return {
        word
        for word in re.findall(r"[a-zA-Z0-9']{3,}", text.lower())
        if word not in {"the", "and", "for", "that", "this", "with", "you", "your", "about"}
    }


def _chunk_text(text: str, *, chunk_size: int = 1800, overlap: int = 220) -> List[str]:
    text = re.sub(r"\n{3,}", "\n\n", text.strip())
    if not text:
        return []
    chunks = []
    start = 0
    while start < len(text):
        end = min(len(text), start + chunk_size)
        if end < len(text):
            boundary = text.rfind("\n\n", start, end)
            if boundary > start + chunk_size // 2:
                end = boundary
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        if end >= len(text):
            break
        start = max(end - overlap, start + 1)
    return chunks


def _select_relevant_chunks(text: str, query: str, max_chars: int) -> List[str]:
    chunks = _chunk_text(text)
    if not chunks:
        return []
    query_words = _words(query)
    ranked = []
    for index, chunk in enumerate(chunks):
        score = len(_words(chunk) & query_words)
        ranked.append((score, -index, chunk))
    ranked.sort(reverse=True)
    selected = []
    total = 0
    for _, _, chunk in ranked:
        if total + len(chunk) > max_chars and selected:
            continue
        selected.append(chunk)
        total += len(chunk)
        if total >= max_chars:
            break
    return selected or [chunks[0][:max_chars]]


def _format_transient_context(label: str, text: str, query: str, max_chars: int) -> str:
    chunks = _select_relevant_chunks(text, query, max_chars)
    if not chunks:
        return ""
    return "\n\n---\n\n".join(
        f"[{label} excerpt {index}]\n{chunk}" for index, chunk in enumerate(chunks, 1)
    )


def _is_marks_query(text: str) -> bool:
    normalized = text.lower()
    if re.search(r"\b(question|quotation|punctuation)\s+mark\b", normalized):
        return False
    return bool(
        re.search(
            r"\b(marks?|grades?|scores?|percent(?:age)?s?|points?|pass(?:ed)?|fail(?:ed)?)\b",
            normalized,
        )
    )


def _fixed_chat_response(message: str) -> StreamingResponse:
    async def stream() -> AsyncIterator[bytes]:
        yield (json.dumps({"type": "sources", "sources": []}) + "\n").encode()
        yield (json.dumps({"type": "delta", "text": message}) + "\n").encode()
        yield (json.dumps({"type": "done"}) + "\n").encode()

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


SYSTEM_BASE = (
    "You are FeedbackChatbot, a feedback-support assistant for students. Your highest "
    "priority is the student's actual tutor feedback. Every answer must be grounded in "
    "the tutor feedback text when it is available. When a student asks about a feedback "
    "comment, locate the relevant comment in the tutor feedback excerpts, paraphrase "
    "that comment precisely, then explain it. Do not invent, infer, or generalise tutor "
    "feedback that is not explicitly present. Do not use phrases like \"the feedback may "
    "be pointing out\" or \"the feedback suggests\". If you cannot locate a relevant "
    "feedback comment, say that clearly and ask the student to point you to the exact "
    "comment. Use rubric excerpts only to explain marking criteria, and use coursework "
    "excerpts only to help interpret the student's submitted work. Do not state, infer, "
    "estimate, calculate, compare, or discuss marks, grades, scores, percentages, point "
    "totals, grade boundaries, pass/fail status, or likely mark changes. If asked about "
    "marks, say you cannot discuss marks and redirect to qualitative feedback. Do not "
    "say or imply that implementing enhancements will get, gain, recover, increase, or "
    "guarantee marks. Do not assert coursework results, metrics, charts, confusion "
    "matrices, outputs, conclusions, or findings unless they are explicitly present in "
    "the tutor feedback or coursework excerpts. If you cannot locate the evidence, say "
    "so clearly."
)


def _build_messages(
    req: ChatRequest,
    rubric_context: str,
    feedback_context: str,
    coursework_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, marks, and safety rules above take priority:\n"
            + persona
        )
    if feedback_context:
        sys_parts.append("Tutor feedback excerpts:\n\n" + feedback_context)
    else:
        sys_parts.append(
            "No tutor feedback text was supplied. If the student asks about feedback, "
            "ask them to paste or upload the relevant tutor feedback comment."
        )
    if rubric_context:
        sys_parts.append(
            "Rubric excerpts. These are secondary to tutor feedback:\n\n"
            + rubric_context
        )
    if coursework_context:
        sys_parts.append(
            "Student coursework excerpts. These are secondary to tutor feedback and rubric:\n\n"
            + coursework_context
        )
    msgs: List[dict] = [{"role": "system", "content": "\n\n".join(sys_parts)}]
    for m in req.messages:
        if m.role == "system":
            continue
        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"), "")
    if _is_marks_query(last_user):
        return _fixed_chat_response(MARKS_REFUSAL)

    nodes = await asyncio.to_thread(rag.retrieve, last_user, req.top_k) if last_user else []
    rubric_context = rag.format_context(nodes)
    if len(rubric_context) > MAX_RUBRIC_CONTEXT_CHARS:
        rubric_context = rubric_context[:MAX_RUBRIC_CONTEXT_CHARS].rstrip()
    feedback_context = _format_transient_context(
        "Tutor feedback",
        req.tutor_feedback_text,
        last_user,
        MAX_FEEDBACK_CONTEXT_CHARS,
    )
    coursework_context = _format_transient_context(
        "Coursework",
        req.coursework_text,
        last_user,
        MAX_COURSEWORK_CONTEXT_CHARS,
    )
    sources = sorted({n.node.metadata.get("source_filename", "unknown") for n in nodes})
    if feedback_context:
        sources.append("Tutor feedback upload/paste")
    if coursework_context:
        sources.append("Coursework upload")
    messages = _build_messages(req, rubric_context, feedback_context, coursework_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 tutor feedback, rubric, and coursework context. Do not repeat your thinking process. "
                "Do not produce a thinking process, analysis section, or hidden reasoning. "
                "Answer immediately and cite any rubric filenames in square brackets. "
                "Do not discuss marks or promise mark increases. Do not include coursework "
                "results or metrics unless they are explicitly present in the supplied context."
            ),
        },
    ]

    async def stream() -> AsyncIterator[bytes]:
        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",
        },
    )