app.py

"""
FastAPI application — routes for chat (SSE streaming), auth, settings, and static files.
"""

import os
import json
import time
import io
import uuid
import logging
from typing import Optional
from fastapi import FastAPI, Request, UploadFile, File, Header, HTTPException, Depends
from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from langchain_core.messages import HumanMessage, AIMessage #type:ignore

import db
import guard
import retriever
import auth as auth_module
import feed as feed_module
import security
from config import GOOGLE_CLIENT_ID, APP_URL
from graph import chatbot, _classify
from tools import run_web_search, fetch_yt_transcript

logger = logging.getLogger(__name__)


app = FastAPI()

# Explicit allow-list instead of a wildcard. The app is normally same-origin
# (FastAPI serves the SPA), so CORS only matters for the deployed origin plus
# local dev. Auth uses Bearer tokens (not cookies), so credentials stay off.
_ALLOWED_ORIGINS = [
    o.strip() for o in os.environ.get(
        "ALLOWED_ORIGINS",
        f"{APP_URL},http://localhost:7860,http://127.0.0.1:7860",
    ).split(",") if o.strip()
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=_ALLOWED_ORIGINS,
    allow_credentials=False,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["Authorization", "Content-Type"],
)


# ── Authentication dependency ─────────────────────────────────

def require_user(authorization: str = Header(default="")) -> str:
    """
    Authenticate the caller from `Authorization: Bearer <token>` and return their
    google_id. Raises 401 if missing/invalid/expired. Every protected endpoint
    derives identity from this — a client-supplied user_id is never trusted.
    """
    if not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Not authenticated")
    uid = security.verify_session(authorization[7:].strip())
    if not uid:
        raise HTTPException(status_code=401, detail="Invalid or expired session")
    return uid

STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")


# --- Request / Response models ---

class GoogleAuthRequest(BaseModel):
    token: str


class ChatRequest(BaseModel):
    message: str
    thread_id: str
    persona: str = "nerd"
    language: str = "auto"
    username: str = ""
    user_id: str = ""  # ignored — identity comes from the session token
    search_mode: str = "none"  # "web" | "none" — structured replacement for the old sentinel
    image: str = ""  # base64 encoded image data
    doc_text: str = ""  # extracted text from attached document
    doc_name: str = ""  # exact name of the document
    doc_bytes: str = ""  # base64 encoded bytes of the document


class RenameRequest(BaseModel):
    thread_id: str
    title: str


class ApiKeyRequest(BaseModel):
    user_id: str
    key: str


class TavilyKeyRequest(BaseModel):
    user_id: str
    key: str


class ProfileRequest(BaseModel):
    user_id: str
    profile: str


class NoteCreateRequest(BaseModel):
    user_id: str
    title: str = ""
    content: str = ""


class NoteUpdateRequest(BaseModel):
    title: str = ""
    content: str = ""


# --- Feedback Endpoint ---
class FeedbackRequest(BaseModel):
    user_id: str
    user_email: Optional[str] = ""
    user_name: Optional[str] = ""
    category: str
    overall: int = 5
    ease: int = 4
    quality: int = 4
    message: str
    attachments: Optional[list[dict]] = None   # [{filename, content (b64), content_type}]

@app.post("/feedback")
async def submit_feedback(req: FeedbackRequest, uid: str = Depends(require_user)):
    ok = feed_module.send_feedback(
        user_id=uid,
        user_email=req.user_email or "",
        user_name=req.user_name or "",
        overall=req.overall,
        ease=req.ease,
        quality=req.quality,
        category=req.category,
        message=req.message,
        attachments=req.attachments,
    )
    return JSONResponse(content={"status": "success" if ok else "logged"})


# --- SSE helpers ---

def sse_token(token: str) -> str:
    return f"data: {json.dumps({'token': token})}\n\n"


def sse_error(message: str) -> str:
    return f"data: {json.dumps({'error': message})}\n\n"


def sse_done() -> str:
    return "data: [DONE]\n\n"


def sse_tool_event(event: str, tool_name: str) -> str:
    """Emit tool_start / tool_end events so the frontend can show a progress bar."""
    return f"data: {json.dumps({'tool_event': event, 'tool': tool_name})}\n\n"


# --- Auth routes ---

@app.post("/auth/google")
def google_login(req: GoogleAuthRequest):
    """Verify Google OAuth ID token and create/update user."""
    idinfo = auth_module.verify_google_token(req.token)
    if not idinfo:
        return JSONResponse({"error": "Invalid token"}, status_code=401)

    user = db.upsert_user(
        google_id=idinfo["sub"],
        email=idinfo.get("email", ""),
        name=idinfo.get("name", ""),
        picture=idinfo.get("picture", ""),
    )
    has_key = bool(db.get_user_api_key(idinfo["sub"]))
    has_tavily = bool(db.get_tavily_key(idinfo["sub"]))
    # Issue a signed session token; the client sends it on every later request.
    token = security.mint_session(idinfo["sub"])
    return {"user": user, "token": token, "has_api_key": has_key, "has_tavily_key": has_tavily}


@app.get("/auth/me")
def get_me(uid: str = Depends(require_user)):
    """Get the authenticated user's own data (identity from the session token)."""
    user = db.get_user(uid)
    if not user:
        return JSONResponse({"error": "User not found"}, status_code=404)
    return {
        "user": user,
        "has_api_key": bool(user.get("openrouter_key")),
        "has_tavily_key": bool(user.get("tavily_key")),
    }


@app.get("/auth/client_id")
def get_client_id():
    """Return the Google Client ID for frontend OAuth init."""
    return {"client_id": GOOGLE_CLIENT_ID}


# --- User settings routes (all scoped to the authenticated user) ---

@app.post("/user/apikey")
def save_api_key(req: ApiKeyRequest, uid: str = Depends(require_user)):
    db.save_user_api_key(uid, req.key)
    return {"ok": True}


@app.post("/user/tavilykey")
def save_tavily_key(req: TavilyKeyRequest, uid: str = Depends(require_user)):
    db.save_tavily_key(uid, req.key)
    return {"ok": True, "has_tavily_key": bool(req.key.strip())}


@app.post("/user/profile")
def save_profile(req: ProfileRequest, uid: str = Depends(require_user)):
    db.save_student_profile(uid, req.profile)
    return {"ok": True}


@app.post("/upload-doc")
async def upload_doc(file: UploadFile = File(...), uid: str = Depends(require_user)):
    """
    Accept a PDF, DOCX, or plain-text file and return its extracted text.
    The frontend sends this text back as `doc_text` in the chat request.
    """
    filename  = file.filename or ""
    ext       = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    try:
        raw_bytes = await file.read()
    except Exception as exc:
        logger.warning("upload read failed: %s", exc)
        return JSONResponse({"error": f"Failed to read file: {exc}"}, status_code=400)

    if len(raw_bytes) > 25 * 1024 * 1024:
        return JSONResponse({"error": "File too large (max 25 MB). Try splitting it into smaller files."}, status_code=413)

    text = ""

    try:
        if ext == "pdf":
            import pdfplumber #type:ignore
            parts: list[str] = []
            with pdfplumber.open(io.BytesIO(raw_bytes)) as pdf:
                for page in pdf.pages:
                    parts.append(page.extract_text() or "")
                    # Also pull tabular data (numericals, data tables) so the model
                    # sees structured rows, not just the flowing-text layer.
                    try:
                        for table in (page.extract_tables() or []):
                            for row in table:
                                cells = [c for c in row if c]
                                if cells:
                                    parts.append(" | ".join(cells))
                    except Exception:
                        pass  # table extraction is best-effort
            text = "\n".join(parts)
        elif ext in ("docx", "doc"):
            from docx import Document as DocxDocument
            doc  = DocxDocument(io.BytesIO(raw_bytes))
            parts = [p.text for p in doc.paragraphs]
            # Include table cells (common for chemistry/physics data sheets).
            for table in doc.tables:
                for row in table.rows:
                    cells = [c.text.strip() for c in row.cells if c.text.strip()]
                    if cells:
                        parts.append(" | ".join(cells))
            text = "\n".join(parts)
        else:
            # Plain text / markdown / code files
            text = raw_bytes.decode("utf-8", errors="replace")
    except Exception as exc:
        logger.warning("doc parse failed for %s: %s", filename, exc)
        return JSONResponse({"error": f"Could not parse file content: {exc}"}, status_code=422)

    text = text.strip()
    if not text:
        # Most common cause: a scanned / image-only PDF with no text layer.
        return JSONResponse(
            {"error": "No readable text found. If this is a scanned PDF or photo, "
                      "attach it as an image instead so it can be read visually."},
            status_code=422,
        )

    # Keep a generous slice so the model sees the whole document, not a fragment.
    return {"filename": filename, "text": text[:50_000]}


# --- Core routes ---

APP_DIR = os.path.dirname(os.path.abspath(__file__))

@app.get("/sw.js")
def serve_sw():
    """Service worker must be served from root for PWA scope."""
    return FileResponse(os.path.join(APP_DIR, "sw.js"), media_type="application/javascript")


@app.get("/manifest.json")
def serve_manifest():
    return FileResponse(os.path.join(STATIC_DIR, "manifest.json"), media_type="application/manifest+json")


@app.get("/")
def serve_index():
    return FileResponse(os.path.join(STATIC_DIR, "index.html"))


@app.get("/threads")
def get_threads(uid: str = Depends(require_user)):
    return {"threads": db.get_threads(uid)}


@app.get("/history/{thread_id}")
def get_history(thread_id: str, uid: str = Depends(require_user)):
    owner = db.get_thread_owner(thread_id)
    if owner is not None and owner != uid:
        raise HTTPException(status_code=403, detail="Not your thread")

    config = {"configurable": {"thread_id": thread_id}}
    state = chatbot.get_state(config)
    messages = state.values.get("messages", [])

    result = []
    for msg in messages:
        role = "user" if isinstance(msg, HumanMessage) else "assistant"
        result.append({"role": role, "content": msg.content})

    return {"messages": result}


@app.post("/chat")
def chat(request: ChatRequest, req: Request, uid: str = Depends(require_user)):
    now = time.time()

    # Ownership: a user may only post into a thread that's new or already theirs.
    owner = db.get_thread_owner(request.thread_id)
    if owner is not None and owner != uid:
        raise HTTPException(status_code=403, detail="Not your thread")
    db.upsert_thread(request.thread_id, request.message, now, uid)

    # Resolve user-specific data from the authenticated identity (never the body).
    user_api_key = db.get_user_api_key(uid)
    tavily_key = db.get_tavily_key(uid)
    user_data = db.get_user(uid)
    student_profile = user_data.get("student_profile", "") if user_data else ""

    def stream():
        # Step 1: Guard
        is_ok, rejection = guard.check_input(request.message)

        if not is_ok:
            yield sse_token(rejection)
            yield sse_done()
            return

        # Step 2: Search mode comes from a structured, trusted field — not a
        # spoofable string prefix in the message body.
        clean_message = request.message
        web_search_enabled = (request.search_mode == "web")
        yt_search_enabled = False

        # Step 3: Retrieve NCERT context from Pinecone. Skip it for casual chit-chat
        # and when a document is attached (the document itself is the context) —
        # saves Pinecone reads and avoids injecting irrelevant noise tokens.
        category = _classify(clean_message)
        if category == "casual" or request.doc_text:
            context = ""
        else:
            try:
                results = retriever.search(clean_message)
                context = retriever.format_context(results)
            except Exception as exc:
                logger.warning("pinecone retrieval failed: %s", exc)
                context = "No relevant context found due to a temporary search error."

        # Step 3.5: Execute search/transcript tools programmatically
        if web_search_enabled:
            yield sse_tool_event("tool_start", "web_search")
            try:
                web_results = run_web_search(clean_message, api_key=tavily_key)
                logger.info("web_search fetched %d chars", len(web_results))
                context += f"\n\n[Web Search Results]\n{web_results}\n[End Web Search Results]\n"
            except Exception as e:
                logger.warning("web_search failed: %s", e)
                context += "\n\nWeb search failed to complete.\n"
            yield sse_tool_event("tool_end", "web_search")

        elif yt_search_enabled:
            yield sse_tool_event("tool_start", "yt_transcript")
            try:
                transcript = fetch_yt_transcript(clean_message)
                if transcript.startswith("TRANSCRIPT_UNAVAILABLE"):
                    logger.info("yt transcript unavailable: %s", transcript[:120])
                    context += f"\n\n[YouTube Transcript Status]\n{transcript}\n"
                else:
                    logger.info("yt transcript fetched %d chars", len(transcript))
                    context += f"\n\n[YouTube Video Transcript]\n{transcript[:9_000]}\n[End Transcript]\n"
            except Exception as e:
                logger.warning("yt transcript failed: %s", e)
                context += "\n\n[YouTube Transcript Status]\nTRANSCRIPT_UNAVAILABLE: extraction failed unexpectedly.\n"
            yield sse_tool_event("tool_end", "yt_transcript")

        # Step 4: Handle document context
        doc_context = ""
        if request.doc_text:
            doc_label = f" (Name: {request.doc_name})" if request.doc_name else ""
            doc_context = (
                f"\n\n[ATTACHED DOCUMENT{doc_label}]\n"
                + request.doc_text[:50_000]
                + "\n[END DOCUMENT]\n"
            )

        # Step 5: Build message (text or multimodal with image)
        if request.image:
            msg_content = [
                {"type": "text", "text": clean_message + doc_context},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{request.image}"},
                },
            ]
        else:
            msg_content = clean_message + doc_context

        # Step 6: Stream LLM response via LangGraph
        has_image = bool(request.image)

        # Resolve target model dynamically to trace and pass to LangGraph config
        try:
            from graph import _pick
            target_model, _ = _pick(category, has_image=has_image)
        except Exception as e:
            logger.warning("model pick failed: %s", e)
            target_model = ""

        logger.info("routing query to model: %s", target_model)

        config = {
            "configurable": {
                "thread_id":       request.thread_id,
                "persona":         request.persona,
                "context":         context,
                "language":        request.language,
                "username":        request.username,
                "student_profile": student_profile,
                "user_api_key":    user_api_key,
                "has_image":       has_image,
                "search_enabled":  (web_search_enabled or yt_search_enabled),
                "model":           target_model,
                "doc_name":        request.doc_name,
                "doc_bytes":       request.doc_bytes,
            }
        }

        try:
            for chunk, _metadata in chatbot.stream(
                {"messages": [HumanMessage(content=msg_content)]},
                config=config,
                stream_mode="messages",
            ):
                if isinstance(chunk, AIMessage) and chunk.content:
                    yield sse_token(chunk.content)
        except Exception as e:
            error_str = str(e).lower()
            if "429" in str(e) or "rate" in error_str or "toomany" in error_str:
                yield sse_error("This model is busy right now (rate limited). Tap regenerate to retry — it'll route to another model.")
            elif "402" in str(e) or "payment" in error_str or "credits" in error_str:
                yield sse_error("Free credits exhausted on OpenRouter. Add credits or try again later.")
            # Check auth BEFORE the generic 404/"not found" branch — OpenRouter's
            # invalid-key error reads "User not found." which contains "not found".
            elif ("401" in str(e) or "unauthorized" in error_str
                  or "user not found" in error_str or "invalid api key" in error_str):
                yield sse_error("Your OpenRouter API key is invalid or wasn't saved. Re-enter it in Settings → API Key.")
            elif "404" in str(e) or "not found" in error_str or "no endpoints" in error_str:
                yield sse_error("That model is temporarily unavailable. Tap regenerate to retry on a different model.")
            elif "timeout" in error_str or "timed out" in error_str:
                yield sse_error("The model took too long to respond. Tap regenerate to retry.")
            else:
                yield sse_error("The request couldn't be completed — it may have routed to a model that's down. Tap regenerate to retry.")
            yield sse_done()
            return

        yield sse_done()

    return StreamingResponse(
        stream(),
        media_type="text/event-stream",
        headers={"X-Accel-Buffering": "no"},
    )


@app.post("/rename")
def rename_thread(request: RenameRequest, uid: str = Depends(require_user)):
    owner = db.get_thread_owner(request.thread_id)
    if owner is not None and owner != uid:
        raise HTTPException(status_code=403, detail="Not your thread")
    db.rename_thread(request.thread_id, request.title)
    return {"ok": True}


@app.delete("/thread/{thread_id}")
def delete_thread(thread_id: str, uid: str = Depends(require_user)):
    owner = db.get_thread_owner(thread_id)
    if owner is not None and owner != uid:
        raise HTTPException(status_code=403, detail="Not your thread")
    db.delete_thread(thread_id)
    return {"ok": True}


# --- Notes (Canvas) — all scoped to the authenticated user ---

@app.get("/notes")
def list_notes(uid: str = Depends(require_user)):
    return {"notes": db.list_notes(uid)}


@app.get("/notes/{note_id}")
def get_note(note_id: str, uid: str = Depends(require_user)):
    if db.get_note_owner(note_id) != uid:
        raise HTTPException(status_code=403, detail="Not your note")
    note = db.get_note(note_id)
    if not note:
        return JSONResponse({"error": "not found"}, status_code=404)
    return note


@app.post("/notes")
def create_note(request: NoteCreateRequest, uid: str = Depends(require_user)):
    now = time.time()
    note_id = str(uuid.uuid4())
    db.create_note(note_id, uid, request.title, request.content, now)
    return {"note_id": note_id, "updated_at": now}


@app.put("/notes/{note_id}")
def update_note(note_id: str, request: NoteUpdateRequest, uid: str = Depends(require_user)):
    if db.get_note_owner(note_id) != uid:
        raise HTTPException(status_code=403, detail="Not your note")
    now = time.time()
    db.update_note(note_id, request.title, request.content, now)
    return {"ok": True, "updated_at": now}


@app.delete("/notes/{note_id}")
def delete_note(note_id: str, uid: str = Depends(require_user)):
    if db.get_note_owner(note_id) != uid:
        raise HTTPException(status_code=403, detail="Not your note")
    db.delete_note(note_id)
    return {"ok": True}


# --- Health check ---

@app.get("/health")
def health():
    checks = {"db": False, "pinecone": False}
    try:
        db.conn.execute("SELECT 1")
        checks["db"] = True
    except Exception:
        pass
    try:
        retriever._index.describe_index_stats()
        checks["pinecone"] = True
    except Exception:
        pass
    ok = all(checks.values())
    return JSONResponse(
        {"status": "ok" if ok else "degraded", "checks": checks},
        status_code=200 if ok else 503,
    )


# --- Static file serving ---
if os.path.isdir(ASSETS_DIR):
    app.mount("/assets", StaticFiles(directory=ASSETS_DIR), name="assets")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


_STATIC_ROOT = os.path.realpath(STATIC_DIR)


@app.get("/{filepath:path}")
def serve_static(filepath: str):
    # Resolve symlinks/`..` and require the result to stay inside STATIC_DIR,
    # so crafted paths can't escape to read arbitrary files on the host.
    full_path = os.path.realpath(os.path.join(STATIC_DIR, filepath))
    if full_path != _STATIC_ROOT and not full_path.startswith(_STATIC_ROOT + os.sep):
        return JSONResponse({"error": "not found"}, status_code=404)
    if os.path.isfile(full_path):
        return FileResponse(full_path)
    return JSONResponse({"error": "not found"}, status_code=404)