backend/main.py

# backend/main.py
import sys
sys.path.insert(0, ".")

import os
import time
import shutil
from pathlib import Path

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from backend.parser import parse_document
from backend.graph import FinancialGraph
from backend.llm import ask, generate_report, compare_companies as llm_compare
from backend.red_flags import evaluate_red_flags, get_value
from backend.recommendations import evaluate_recommendation
from backend.entity_resolver import format_money
from backend.verifier import verify_answer

app = FastAPI(title="FinSight API")

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

DATA_DIR = Path("data")
DATA_DIR.mkdir(exist_ok=True)

GRAPH_PATH = str(DATA_DIR / "graph.json")

UPLOADS_DIR = DATA_DIR / "uploads"
UPLOADS_DIR.mkdir(exist_ok=True)

fg = FinancialGraph()

if Path(GRAPH_PATH).exists():
    fg.load(GRAPH_PATH)


class QueryRequest(BaseModel):
    question: str
    company: str | None = None


class CompareRequest(BaseModel):
    companies: list[str]
    metric: str = "revenue"
    question: str | None = None


def try_direct_metric_answer(
    question: str,
    company: str,
    metrics_by_year: dict
):
    q_lower = question.lower()

    for raw_year, metrics in metrics_by_year.items():
        year_str = str(raw_year).strip()

        year_matches = year_str in q_lower
        if not year_matches:
            try:
                year_matches = str(int(year_str)) in q_lower
            except ValueError:
                year_matches = False

        if not year_matches:
            continue

        for key, raw in metrics.items():
            metric_words = key.replace("_", " ").lower().split()

            if all(w in q_lower for w in metric_words):
                value, confidence = get_value(raw)

                if value is None:
                    continue

                label = key.replace("_", " ")

                answer = f"{company} {year_str} {label}: {value}"

                if confidence in ("low", "medium"):
                    answer += (
                        f" (confidence: {confidence} — "
                        f"verify before relying on this figure)"
                    )

                return answer

    return None


@app.get("/health")
def health():
    return {
        "status": "ok",
        "service": "FinSight API"
    }


@app.post("/upload")
async def upload_document(
    file: UploadFile = File(...),
    company: str = Form(...),
    year: str = Form(...)
):
    file_path = UPLOADS_DIR / file.filename

    with open(file_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    try:
        parsed = parse_document(
            str(file_path),
            company,
            year
        )
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=str(e)
        )

    # Relation extraction's LLM fallback fires one local-model call per
    # ambiguous org pair — on a 300-page filing that's hours of CPU.
    # Pattern-based relations always run; the LLM pass is opt-in.
    use_llm_relations = os.getenv(
        "RELATION_LLM_FALLBACK", "false"
    ).strip().lower() == "true"

    fg.add_document(parsed, use_llm_fallback=use_llm_relations)
    fg.save(GRAPH_PATH)

    metrics_summary = {}
    low_confidence_flags = []

    for k, v in parsed["metrics"].items():
        if isinstance(v, dict):
            metrics_summary[k] = v.get("value")

            if v.get("needs_clarification"):
                low_confidence_flags.append(k)
        else:
            metrics_summary[k] = v

    return {
        "status": "ok",
        "company": company,
        "year": year,
        "sector": parsed.get("sector", "GENERAL"),
        "chunks": parsed["chunk_count"],
        "metrics": parsed["metrics"],
        "metrics_summary": metrics_summary,
        "low_confidence_metrics": low_confidence_flags,
        "entities": len(parsed["entities"])
    }


@app.post("/query")
def query(req: QueryRequest):
    t0 = time.time()

    if req.company:
        all_metrics = fg.get_company_metrics(req.company)

        direct = try_direct_metric_answer(
            req.question,
            req.company,
            all_metrics
        )

        if direct:
            return {
                "answer": direct,
                "latency": round(time.time() - t0, 2),
                "chunks_used": 0,
                "source": "direct_lookup"
            }

    chunks = fg.get_relevant_chunks(
        req.question,
        req.company
    )

    metrics_context = ""

    if req.company:
        all_metrics = fg.get_company_metrics(req.company)

        if all_metrics:
            lines = []

            for year, metrics in all_metrics.items():
                for key, raw in metrics.items():

                    if isinstance(raw, dict):
                        val = raw.get("value")
                        confidence = raw.get(
                            "confidence",
                            "unknown"
                        )
                        currency = raw.get("currency", "USD")
                    else:
                        val = raw
                        confidence = "n/a"
                        currency = "USD"

                    if val is None:
                        continue

                    formatted = format_money(
                        val,
                        currency
                    )

                    line = (
                        f"{req.company} {year} "
                        f"{key}: {formatted}"
                    )

                    if confidence in ("low", "medium"):
                        line += (
                            f" (confidence: {confidence} — "
                            f"verify before stating as certain)"
                        )

                    lines.append(line)

            metrics_context = (
                "Key Metrics:\n"
                + "\n".join(lines)
                + "\n\n"
            )

    context = metrics_context + "\n\n".join(c["text"] for c in chunks)

    answer = ask(
        req.question,
        context
    )

    # every number in the generated answer must be traceable to the
    # source material; unverified numbers are flagged, not hidden
    all_metrics_flat = {}
    if req.company:
        for _year, m in fg.get_company_metrics(req.company).items():
            for k, v in m.items():
                all_metrics_flat[f"{_year}_{k}"] = v

    verification = verify_answer(answer, context, all_metrics_flat)

    citations = [
        {
            "company": c.get("company"),
            "year": c.get("year"),
            "page": c.get("page")
        }
        for c in chunks
    ]

    return {
        "answer": answer,
        "latency": round(time.time() - t0, 2),
        "chunks_used": len(chunks),
        "citations": citations,
        "verification": verification
    }


@app.get("/report/{company}/{year}")
def get_report(company: str, year: str):
    t0 = time.time()

    metrics = fg.get_company_metrics(company)

    if year not in metrics:
        raise HTTPException(
            status_code=404,
            detail=f"No data for {company} {year}"
        )

    year_metrics = metrics[year]

    # llm.py's generate_report() formats metrics assuming bare floats
    # (its internal fmt() does `if val >= 1_000_000_000`, which raises
    # TypeError on a dict). Most metrics here are dict-shaped
    # ({value, confidence, alternatives, ...}) from find_metric_in_text;
    # ratios from find_ratio_in_text are still bare floats. Flatten to
    # bare values before calling generate_report — same pattern /compare
    # already uses for llm_compare(), applied here too.
    flat_year_metrics = {
        k: (v.get("value") if isinstance(v, dict) else v)
        for k, v in year_metrics.items()
    }

    chunks = fg.get_relevant_chunks(
        f"{company} financial performance {year}",
        company
    )

    context = "\n\n".join(c["text"] for c in chunks)

    report = generate_report(
        company,
        year,
        flat_year_metrics,
        context
    )

    verification = verify_answer(report, context, year_metrics)

    citations = [
        {
            "company": c.get("company"),
            "year": c.get("year"),
            "page": c.get("page")
        }
        for c in chunks
    ]

    return {
        "company": company,
        "year": year,
        "metrics": year_metrics,
        "report": report,
        "latency": round(time.time() - t0, 2),
        "citations": citations,
        "verification": verification
    }


@app.get("/red_flags/{company}/{year}")
def get_red_flags(company: str, year: str):
    sector = fg.get_filing_sector(
        company,
        year
    )

    if sector is None:
        raise HTTPException(
            status_code=404,
            detail=f"No data for {company} {year}"
        )

    return evaluate_red_flags(
        fg,
        company,
        year,
        sector=sector
    )


@app.get("/recommendation/{company}/{year}")
def get_recommendation(company: str, year: str):
    sector = fg.get_filing_sector(
        company,
        year
    )

    if sector is None:
        raise HTTPException(
            status_code=404,
            detail=f"No data for {company} {year}"
        )

    return evaluate_recommendation(
        fg,
        company,
        year,
        sector=sector
    )


@app.get("/metrics/{company}")
def get_metrics(company: str):
    metrics = fg.get_company_metrics(company)

    if not metrics:
        raise HTTPException(
            status_code=404,
            detail=f"No data for {company}"
        )

    return {
        "company": company,
        "metrics": metrics
    }


@app.post("/compare")
def compare(req: CompareRequest):
    t0 = time.time()

    metric_data = fg.compare_companies(
        req.companies,
        req.metric
    )

    flat_metric_data = {
        company: {
            year: (
                value.get("value")
                if isinstance(value, dict)
                else value
            )
            for year, value in years.items()
        }
        for company, years in metric_data.items()
    }

    analysis = llm_compare(
        req.companies,
        flat_metric_data,
        req.question
    )

    return {
        "metric": req.metric,
        "data": metric_data,
        "analysis": analysis,
        "latency": round(time.time() - t0, 2)
    }


@app.get("/companies")
def list_companies():
    companies = set()

    for node, data in fg.G.nodes(data=True):
        if data.get("type") == "company":
            companies.add(node)

    return {
        "companies": list(companies)
    }


# ── static frontend (single-container deploys) ───────────────
# When a built React bundle exists (Docker image / production), serve it
# from the same origin. Mounted LAST so every API route above wins first.
# In dev this directory doesn't exist and Vite serves the frontend instead.
_FRONTEND_DIST = Path(__file__).resolve().parent.parent / "frontend" / "dist"
if _FRONTEND_DIST.is_dir():
    from fastapi.staticfiles import StaticFiles
    app.mount("/", StaticFiles(directory=str(_FRONTEND_DIST), html=True), name="frontend")


@app.get("/graph/{company}")
def get_graph(company: str):
    return fg.get_company_graph(company)