backend/main.py

"""FastAPI app — paper review system."""
from __future__ import annotations

import secrets
from datetime import datetime, timezone
from time import perf_counter
from pathlib import Path

from fastapi import BackgroundTasks, Depends, FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles

from backend.config import settings
from backend.auth import CurrentUser, get_current_user, get_optional_user, require_admin
from backend.storage import jobs as job_store
from backend.observability import current_step_id, record_api_call, tracked_step

app = FastAPI(title="PaperMate API", version="0.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve frontend static files at /app/ (keeps /api/ routes reachable)
_frontend = Path(__file__).parent.parent / "frontend"
if _frontend.exists():
    app.mount("/app", StaticFiles(directory=str(_frontend), html=True), name="frontend")


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/")
async def root():
    return RedirectResponse(url="/app/home.html")


@app.get("/health")
async def health():
    return {"status": "ok", "reviewer": "multi-agent-v1"}


@app.post("/api/submit")
async def submit(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    email: str = Form(default=""),
    user: CurrentUser | None = Depends(get_optional_user),
):
    if user is None:
        raise HTTPException(status_code=401, detail="Please sign in to submit a paper.")

    email = user.email

    # Validate file type
    if not file.filename or not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted.")

    # Validate file size
    pdf_bytes = await file.read()
    max_bytes = settings.max_file_size_mb * 1024 * 1024
    if len(pdf_bytes) > max_bytes:
        raise HTTPException(
            status_code=400,
            detail=f"File too large. Maximum size is {settings.max_file_size_mb} MB.",
        )

    # Per-user monthly quota (admins are never limited).
    if not user.is_admin:
        await _enforce_quota(user)

    access_key = "sk-pm-" + secrets.token_urlsafe(24)
    job = await job_store.create_job(access_key, email, file.filename, pdf_bytes, user.profile_id)

    background_tasks.add_task(_run_pipeline, access_key, pdf_bytes, email, file.filename)

    return JSONResponse({
        "access_key": job.get("access_key", access_key),
        "message": "Your paper has been submitted for review.",
    })


@app.get("/api/review/{access_key}")
async def get_review(access_key: str):
    job = await job_store.get_job(access_key)
    if job is None:
        raise HTTPException(status_code=404, detail="Review not found. Check your access key.")
    # Retrieval links for the original PDF and the parsed markdown text.
    job["pdf_url"] = f"/api/review/{access_key}/pdf"
    if job.get("has_parsed_markdown"):
        job["markdown_url"] = f"/api/review/{access_key}/markdown"
    return JSONResponse(job)


@app.get("/api/review/{access_key}/pdf")
async def get_review_pdf(access_key: str):
    """Return the ORIGINAL uploaded PDF (application/pdf)."""
    result = await job_store.get_pdf_bytes(access_key)
    if result is None:
        raise HTTPException(status_code=404, detail="PDF not found for this access key.")
    data, filename = result
    return Response(
        content=data,
        media_type="application/pdf",
        headers={"Content-Disposition": f'inline; filename="{filename}"'},
    )


@app.get("/api/review/{access_key}/markdown")
async def get_review_markdown(access_key: str):
    """Return the parsed markdown text stored in the DB."""
    markdown = await job_store.get_parsed_markdown(access_key)
    if not markdown:
        raise HTTPException(status_code=404, detail="Parsed markdown not found for this access key.")
    return Response(content=markdown, media_type="text/markdown; charset=utf-8")


@app.post("/api/feedback/{access_key}")
async def submit_feedback(access_key: str, request: Request):
    job = await job_store.get_job(access_key)
    if job is None:
        raise HTTPException(status_code=404, detail="Review not found.")
    if job.get("status") != "completed":
        raise HTTPException(status_code=400, detail="Feedback can only be submitted for completed reviews.")

    body = await request.json()
    feedback = {
        "helpfulness": body.get("helpfulness"),
        "has_critical_error": body.get("has_critical_error"),
        "has_suggestions": body.get("has_suggestions"),
        "comments": str(body.get("comments", "")).strip()[:500],
        "submitted_at": datetime.now(timezone.utc).isoformat(),
    }
    await job_store.save_feedback(access_key, feedback)
    return JSONResponse({"message": "Thank you for your feedback!"})


# ---------------------------------------------------------------------------
# Auth / quota helpers
# ---------------------------------------------------------------------------

def _resolve_limit(profile_value, default_value):
    """Effective limit: per-profile override wins, else config default. 0/None = unlimited."""
    if profile_value:
        return profile_value
    return default_value or None


async def _enforce_quota(user: CurrentUser) -> None:
    review_limit = _resolve_limit(
        getattr(user, "monthly_review_limit", None), settings.default_monthly_review_limit
    )
    cost_limit = _resolve_limit(
        getattr(user, "monthly_cost_limit_usd", None), settings.default_monthly_cost_limit_usd
    )
    if not review_limit and not cost_limit:
        return  # unlimited

    usage = await job_store.count_user_usage_this_month(user.profile_id)
    if review_limit and usage.get("review_count", 0) >= review_limit:
        raise HTTPException(
            status_code=429,
            detail=f"Monthly review limit reached ({review_limit}). Try again next month.",
        )
    if cost_limit and usage.get("cost_usd", 0) >= cost_limit:
        raise HTTPException(
            status_code=429,
            detail=f"Monthly cost limit reached (${cost_limit}).",
        )


# ---------------------------------------------------------------------------
# Public config + user ("me") routes
# ---------------------------------------------------------------------------

@app.get("/api/public-config")
async def public_config():
    """Public values the frontend needs to initialise Supabase Auth.

    The anon key is designed to be public (it only grants RLS-gated access)."""
    return {
        "supabase_url": settings.supabase_url,
        "supabase_anon_key": settings.supabase_anon_key,
        "app_base_url": settings.app_base_url,
        "auth_required_for_submit": settings.auth_required_for_submit,
    }


@app.get("/api/me")
async def get_me(user: CurrentUser = Depends(get_current_user)):
    return {
        "profile_id": user.profile_id,
        "email": user.email,
        "role": user.role,
        "is_admin": user.is_admin,
    }


@app.get("/api/me/submissions")
async def my_submissions(user: CurrentUser = Depends(get_current_user)):
    return JSONResponse(await job_store.list_my_submissions(user.profile_id))


# ---------------------------------------------------------------------------
# Admin routes (all require role = 'admin')
# ---------------------------------------------------------------------------

@app.get("/api/admin/overview")
async def admin_overview(_admin: CurrentUser = Depends(require_admin)):
    overview = await job_store.get_admin_overview()
    cost_by_day = await job_store.list_cost_by_day()
    cost_by_provider = await job_store.list_cost_by_provider()
    alert_threshold = settings.cost_alert_daily_usd or 0
    over_budget_days = [
        d for d in cost_by_day
        if alert_threshold and float(d.get("cost_usd") or 0) > alert_threshold
    ]
    return JSONResponse({
        "overview": overview,
        "cost_by_day": cost_by_day,
        "cost_by_provider": cost_by_provider,
        "cost_alert_daily_usd": alert_threshold,
        "over_budget_days": over_budget_days,
    })


@app.get("/api/admin/users")
async def admin_users(
    limit: int = 100,
    offset: int = 0,
    _admin: CurrentUser = Depends(require_admin),
):
    return JSONResponse(await job_store.list_user_usage(limit, offset))


@app.patch("/api/admin/users/{profile_id}")
async def admin_update_user(
    profile_id: str,
    request: Request,
    admin: CurrentUser = Depends(require_admin),
):
    body = await request.json()
    if body.get("role") not in (None, "user", "admin"):
        raise HTTPException(status_code=400, detail="role must be 'user' or 'admin'.")
    # Guard: an admin cannot strip their own admin role (avoid self-lockout).
    if profile_id == admin.profile_id and body.get("role") == "user":
        raise HTTPException(status_code=400, detail="You cannot remove your own admin role.")
    # Guard: never demote/disable the last active admin (avoid locking everyone out).
    demoting = body.get("role") == "user"
    disabling = body.get("is_active") is False
    if demoting or disabling:
        target = await job_store.get_profile_by_id(profile_id)
        if target and target.get("role") == "admin" and target.get("is_active", True):
            if await job_store.count_active_admins() <= 1:
                raise HTTPException(
                    status_code=400,
                    detail="Cannot demote or disable the last active admin.",
                )
    updated = await job_store.update_profile_admin_fields(profile_id, body)
    if updated is None:
        raise HTTPException(status_code=404, detail="User not found.")
    return JSONResponse(updated)


@app.get("/api/admin/submissions")
async def admin_submissions(
    status: str | None = None,
    profile_id: str | None = None,
    q: str | None = None,
    limit: int = 50,
    offset: int = 0,
    _admin: CurrentUser = Depends(require_admin),
):
    rows = await job_store.list_submissions(status, profile_id, q, limit, offset)
    return JSONResponse(rows)


@app.get("/api/admin/submissions/{submission_id}")
async def admin_submission_detail(
    submission_id: int,
    _admin: CurrentUser = Depends(require_admin),
):
    detail = await job_store.get_submission_detail(submission_id)
    if detail is None:
        raise HTTPException(status_code=404, detail="Submission not found.")
    return JSONResponse(detail)


# ---------------------------------------------------------------------------
# Background pipeline
# ---------------------------------------------------------------------------

def _elapsed_ms(started: float) -> int:
    return max(0, round((perf_counter() - started) * 1000))


async def _run_tracked_step(
    access_key: str,
    step_name: str,
    operation,
    *,
    step_order: float | None = None,
    metadata: dict | None = None,
    api_provider: str | None = None,
    api_operation: str | None = None,
    request_count: int = 1,
    output_fn=None,
):
    """Run one pipeline step inside a tracked context.

    - ``api_provider``/``api_operation``: record a (cost-free) external API call
      for this step. The PDF-parse step deliberately omits these so it creates
      NO provider_calls row. The Tavily step records its own cost-bearing call
      inside ``operation`` (so it can price the actual successful searches).
    - ``output_fn(result) -> (output_json, output_summary)``: stash the step's
      output so it is persisted on the pipeline_steps row.
    """
    started = perf_counter()
    async with tracked_step(
        access_key,
        step_name,
        step_order=step_order,
        metadata=metadata,
    ) as step:
        try:
            result = await operation()
        except Exception as exc:
            if api_provider and api_operation:
                await record_api_call(
                    provider=api_provider,
                    operation=api_operation,
                    latency_ms=_elapsed_ms(started),
                    request_count=request_count,
                    status="failed",
                    error=f"{type(exc).__name__}: {exc}",
                )
            raise

        if api_provider and api_operation:
            await record_api_call(
                provider=api_provider,
                operation=api_operation,
                status_code=200,
                latency_ms=_elapsed_ms(started),
                request_count=request_count,
                status="success",
                error=None,
            )
        if output_fn is not None:
            try:
                output_json, output_summary = output_fn(result)
                step["_output_json"] = output_json
                step["_output_summary"] = output_summary
            except Exception:
                pass
        return result


async def _run_pipeline(access_key: str, pdf_bytes: bytes, email: str, filename: str):
    import logging as _logging
    import traceback as _tb
    _log = _logging.getLogger("uvicorn")

    # Catch-all safety net: BackgroundTasks silently swallows exceptions.
    # Log through uvicorn so output always reaches the terminal.
    try:
        await _run_pipeline_inner(access_key, pdf_bytes, email, filename)
    except Exception as _e:
        _log.error(f"[PIPELINE CRASH] {access_key[-8:]} — {type(_e).__name__}: {_e}")
        _log.error(_tb.format_exc())


async def _run_pipeline_inner(access_key: str, pdf_bytes: bytes, email: str, filename: str):
    from backend.pipeline.pdf2md import pdf2md
    from backend.pipeline.extract import extract_paper_title, extract_contributions, extract_research_topic
    from backend.pipeline.search import generate_scientific_search_queries, search_related_papers, tavily_cost
    from backend.pipeline.paper_info import get_paper_info
    from backend.pipeline.summarize import summarize_related_research
    from backend.pipeline.review_agent import run_review_agent
    from backend.pipeline.guardrails import is_nlp_cl_paper, is_paper_follow_page_limit, PaperRejected
    from backend.logger import JobLogger
    import asyncio

    log = JobLogger(access_key)
    paper_title = None

    try:
        await job_store.set_status(access_key, "processing")
        await job_store.set_progress_step(access_key, 1)  # Submission Screening
        log.separator(f"JOB START — {filename}")

        # Guardrail 1: NLP/CL scope check (parse page 1 only — fast).
        log.info("Guardrail 1 — NLP/CL scope check")
        try:
            await is_nlp_cl_paper(pdf_bytes, logger=log, access_key=access_key)
            log.success("Guardrail 1 passed")
        except PaperRejected as rej:
            log.warning("Guardrail 1 rejected submission", reason=rej.reason)
            await job_store.set_rejected(access_key, rej.reason)
            await _upload_log_artifact(access_key)
            return
        log.info("Job created", key=access_key, file=filename, email=email)

        # Guardrail 2: page limit check — parse pages 1-9 (no VLM), before expensive Step 1.
        log.info("Guardrail 2 — page limit check (LLM section header classification)")
        try:
            await is_paper_follow_page_limit(pdf_bytes, logger=log, access_key=access_key)
            log.success("Guardrail 2 passed")
        except PaperRejected as rej:
            log.warning("Guardrail 2 rejected submission", reason=rej.reason)
            await job_store.set_rejected(access_key, rej.reason)
            await _upload_log_artifact(access_key)
            return

        await job_store.set_progress_step(access_key, 2)  # Reading your paper

        # Step 1: PDF → ParsedPaper (markdown + optional structured JSON), tracked.
        log.info("Step 1/8 — PDF → ParsedPaper", provider=settings.pdf_parser)

        async def _parse_pdf():
            from backend.pipeline.postprocess import clean_parsed_paper
            result = await pdf2md(pdf_bytes)
            result = clean_parsed_paper(result)
            removed = (result.get("_postprocess") or {}).get("acl_line_numbers_removed", 0)
            if removed:
                log.info("Postprocess — ACL line numbers removed", count=removed)
            md = result.get("markdown") or ""
            # Guard: a near-empty parse means the parser failed to read the file.
            # Without this the pipeline would feed empty text to the LLM and
            # fabricate a review about nothing — fail fast instead.
            if len(md.strip()) < settings.min_parsed_markdown_chars:
                raise RuntimeError(
                    f"PDF parsing produced too little text "
                    f"({len(md.strip())} chars) via '{settings.pdf_parser}'. "
                    "The file may be image-only/scanned or the parser failed."
                )
            return result

        parsed = await _run_tracked_step(
            access_key,
            "pdf_to_markdown",
            _parse_pdf,
            step_order=1,
            metadata={"provider": settings.pdf_parser},
            output_fn=lambda p: (
                {
                    "char_count": len(p.get("markdown") or ""),
                    "structured": bool(p.get("structured")),
                    "pages_total": p.get("pages_total"),
                    "provider": settings.pdf_parser,
                },
                f"{len(p.get('markdown') or '')} chars",
            ),
        )
        paper_md = parsed.get("markdown") or ""
        has_structured = bool(parsed.get("structured"))

        await job_store.save_parsed_markdown(
            access_key,
            paper_md,
            provider=settings.pdf_parser,
            parser_version=settings.pdf_parser,
            raw_json={"char_count": len(paper_md), "pages_total": parsed.get("pages_total")},
        )
        # Local debug dump of the full ParsedPaper (markdown already saved to DB).
        _save_parsed(access_key, parsed, log)
        log.success("Step 1 done", chars=len(paper_md), structured=has_structured,
                    pages=parsed.get("pages_total"))
        await job_store.set_progress_step(access_key, 3)  # Researching the literature

        # Step 1b: Extract paper title
        log.info("Step 1b — extracting paper title")
        paper_title = await _run_tracked_step(
            access_key,
            "extract_paper_title",
            lambda: extract_paper_title(parsed, logger=log),
            step_order=1.1,
            output_fn=lambda t: ({"paper_title": t}, t or ""),
        )
        if paper_title:
            await job_store.set_paper_title(access_key, paper_title)
            log.success("Step 1b done", title=paper_title)
        else:
            log.warning("Could not extract paper title, will use filename")

        # Steps 2 & 3: contributions + research topic (parallel, tracked).
        # Uses abstract + introduction only when structured content is available.
        log.info("Step 2+3 — extracting contributions & research topic",
                 source="structured" if has_structured else "raw_markdown")
        contributions, research_topic = await asyncio.gather(
            _run_tracked_step(
                access_key,
                "extract_contributions",
                lambda: extract_contributions(parsed, logger=log),
                step_order=2,
                output_fn=lambda c: ({"contributions": c}, f"{len(c)} contributions"),
            ),
            _run_tracked_step(
                access_key,
                "extract_research_topic",
                lambda: extract_research_topic(parsed, logger=log),
                step_order=3,
                output_fn=lambda t: ({"research_topic": t}, (t or "")[:120]),
            ),
        )
        log.success("Step 2+3 done", contributions=len(contributions), topic=research_topic[:80])

        # Step 4: Generate search queries (linked to this step via evidence_items)
        log.info("Step 4/8 — generating search queries")

        async def _queries_step():
            queries = await generate_scientific_search_queries(contributions, research_topic, logger=log)
            await job_store.record_search_queries(access_key, current_step_id(), queries)
            return queries

        queries = await _run_tracked_step(
            access_key,
            "generate_search_queries",
            _queries_step,
            step_order=4,
            output_fn=lambda q: ({"queries": q}, f"{len(q)} queries"),
        )
        log.success("Step 4 done", queries=len(queries))

        # Step 5: Search related papers (Tavily) — record actual cost
        log.info("Step 5/8 — searching related papers via Tavily")

        async def _search_step():
            papers, search_count = await search_related_papers(queries)
            cost = tavily_cost(search_count)
            await record_api_call(
                provider="tavily",
                operation="search",
                status_code=200,
                request_count=search_count,
                status="success",
                error=None,
                cost_usd=cost,
                metadata={
                    "credits": search_count,
                    "price_per_credit": settings.tavily_price_per_credit,
                    "credits_per_search": settings.tavily_credits_per_search,
                    "search_depth": "basic",
                },
            )
            return papers, search_count

        raw_papers, search_count = await _run_tracked_step(
            access_key,
            "search_related_papers",
            _search_step,
            step_order=5,
            output_fn=lambda r: (
                {"raw_papers_found": r[1], "sample": r[0][:10]},
                f"{r[1]} searches, {len(r[0])} papers",
            ),
        )
        log.success("Step 5 done", papers_found=len(raw_papers), searches=search_count)

        # Step 6: Fetch paper metadata (arXiv) — linked evidence_items
        log.info("Step 6/8 — fetching paper metadata")

        async def _metadata_step():
            papers_with_info = await get_paper_info(raw_papers)
            await job_store.record_related_papers(access_key, current_step_id(), papers_with_info)
            return papers_with_info

        papers_with_info = await _run_tracked_step(
            access_key,
            "fetch_paper_metadata",
            _metadata_step,
            step_order=6,
            api_provider="arxiv",
            api_operation="metadata",
            request_count=max(1, len(raw_papers)),
            output_fn=lambda p: ({"count": len(p)}, f"{len(p)} papers with metadata"),
        )
        log.success("Step 6 done", papers_with_info=len(papers_with_info))

        # Step 7: Summarize related research — linked evidence_items
        log.info("Step 7/8 — summarizing related research")

        async def _summarize_step():
            related_summaries = await summarize_related_research(papers_with_info, logger=log)
            await job_store.record_related_paper_summaries(access_key, current_step_id(), related_summaries)
            return related_summaries

        related_summaries = await _run_tracked_step(
            access_key,
            "summarize_related_research",
            _summarize_step,
            step_order=7,
            output_fn=lambda s: ({"count": len(s)}, f"{len(s)} summaries"),
        )
        log.success("Step 7 done", summaries=len(related_summaries))

        # Step 8: Generate review
        await job_store.set_progress_step(access_key, 4)  # Drafting review
        # Uses structured content (excludes related work / references) when available.
        log.info("Step 8/8 — multi-agent ARR review",
                 source="structured" if has_structured else "raw_markdown")
        review = await _run_tracked_step(
            access_key,
            "generate_review",
            lambda: run_review_agent(
                parsed_paper=parsed,
                related_papers=papers_with_info,
                related_summaries=related_summaries,
                access_key=access_key,
                logger=log,
            ),
            step_order=8,
            output_fn=lambda r: (
                {"overall_assessment": r.get("overall_assessment")},
                r.get("overall_assessment_label", ""),
            ),
        )
        log.success("Step 8 done",
                    overall_assessment=review.get("overall_assessment"),
                    label=review.get("overall_assessment_label"))

        await job_store.set_completed(access_key, review)
        log.separator("JOB COMPLETED")
        await _upload_log_artifact(access_key)

    except Exception as e:
        error_msg = f"{type(e).__name__}: {e}"
        log.error("Pipeline failed", exc=e)
        await job_store.set_failed(access_key, error_msg)
        await _upload_log_artifact(access_key)


async def _upload_log_artifact(access_key: str) -> None:
    log_path = Path(__file__).parent.parent / settings.logs_dir / f"{access_key}.log"
    if not log_path.exists():
        return
    try:
        await job_store.save_job_log(access_key, log_path.read_text(encoding="utf-8"))
    except Exception:
        pass


def _save_parsed(access_key: str, parsed: dict, log) -> None:
    """Dump the full ParsedPaper to data/jobs/<key>_parsed.json for debugging.

    The markdown is already persisted to the DB (save_parsed_markdown); this
    keeps the richer structured JSON (VLM-enriched tables/formulas/figures)
    on local disk for inspection. Best-effort — never breaks the pipeline.
    """
    import json as _json
    try:
        jobs_dir = Path(__file__).parent.parent / settings.jobs_dir
        jobs_dir.mkdir(parents=True, exist_ok=True)
        path = jobs_dir / f"{access_key}_parsed.json"
        with path.open("w", encoding="utf-8") as f:
            _json.dump(parsed, f, ensure_ascii=False, indent=2)
        log.info("ParsedPaper saved for debug", path=str(path))
    except Exception as e:
        log.warning("Could not save ParsedPaper", error=str(e))