src/analysis/forensic.py

"""Forensic accounting engine (premium/owner tier).

Multi-layer automated investigation over the uploaded document set:

  1. Structured extraction — company profile + multi-year key figures (LLM)
  2. Red-flag rules      — deterministic Schilit-style manipulation checks
                           (revenue, expenses, assets, liabilities, cash flow,
                           acquisitions, non-GAAP, policy changes)
  3. Risk scoring        — per-category scores + overall fraud-risk score
  4. Fraud timeline      — notable year-over-year signals
  5. External validation — SEC EDGAR / Companies House lookups (best effort)
  6. Explainability      — LLM narrative: numbered reasons + overall assessment

Rules are deterministic and auditable; the LLM is used only for extraction and
for explaining findings, never for deciding them.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field

from langchain_core.messages import HumanMessage, SystemMessage

from src.agents.qa_agent import format_evidence
from src.llm import get_llm

CATEGORIES = ("revenue", "expenses", "assets", "liabilities", "cash_flow", "governance")

METRICS = ["revenue", "receivables", "inventory", "net_income",
           "operating_cash_flow", "goodwill", "cash", "total_debt",
           "current_assets", "current_liabilities", "total_assets",
           "shareholders_equity", "provisions"]


@dataclass
class Flag:
    category: str
    severity: str           # low | medium | high
    title: str
    explanation: str
    confidence: int         # 0-100

    def as_dict(self) -> dict:
        return self.__dict__.copy()


@dataclass
class ForensicResult:
    profile: dict
    figures: dict                      # {year: {metric: float}}
    kpis: dict                         # latest-year KPIs incl. ratios
    flags: list[Flag]
    category_scores: dict[str, int]
    overall_risk: int
    risk_label: str
    timeline: list[tuple[str, list[str]]]
    external: list[dict]
    narrative: str
    activity: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Layer 1: structured extraction
# ---------------------------------------------------------------------------

EXTRACT_SYSTEM = """You are a forensic data extractor. From the evidence excerpts,
extract the company profile and key figures for EVERY fiscal year present.
Respond with ONLY a raw JSON object (no markdown fences, no prose) shaped as:
{
  "company": str|null, "industry": str|null, "auditor": str|null,
  "exchange": str|null, "country": str|null, "market_cap": str|null,
  "currency": str|null,
  "years": { "<fiscal year>": {
      "revenue": number|null, "receivables": number|null, "inventory": number|null,
      "net_income": number|null, "operating_cash_flow": number|null,
      "goodwill": number|null, "cash": number|null, "total_debt": number|null,
      "current_assets": number|null, "current_liabilities": number|null,
      "total_assets": number|null, "shareholders_equity": number|null,
      "provisions": number|null } },
  "non_gaap_metrics": [str], "policy_changes": [str],
  "capitalised_costs_mentioned": bool, "acquisitions_mentioned": bool
}
All figures as plain numbers in the SAME unit (e.g. millions) — no currency
symbols, no thousands separators. Use null when a figure is not in the evidence.
Never invent numbers."""


def _extract_json_object(text: str) -> dict:
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    fence = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
    candidates = [fence.group(1)] if fence else []
    start = text.find("{")
    if start != -1:
        depth, in_str, esc = 0, False, False
        for i in range(start, len(text)):
            ch = text[i]
            if in_str:
                if esc:
                    esc = False
                elif ch == "\\":
                    esc = True
                elif ch == '"':
                    in_str = False
                continue
            if ch == '"':
                in_str = True
            elif ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    candidates.append(text[start:i + 1])
                    break
    for cand in candidates:
        try:
            data = json.loads(cand)
            if isinstance(data, dict):
                return data
        except json.JSONDecodeError:
            continue
    return {}


def _num(v) -> float | None:
    if isinstance(v, (int, float)):
        return float(v)
    if isinstance(v, str):
        cleaned = re.sub(r"[^\d.\-]", "", v)
        try:
            return float(cleaned) if cleaned not in ("", "-", ".") else None
        except ValueError:
            return None
    return None


def extract_figures(evidence) -> dict:
    llm = get_llm("analyst")
    resp = llm.invoke([
        SystemMessage(content=EXTRACT_SYSTEM),
        HumanMessage(content=f"Evidence excerpts:\n\n{format_evidence(evidence)}"),
    ])
    data = _extract_json_object(resp.content)
    years = {}
    for year, metrics in (data.get("years") or {}).items():
        if isinstance(metrics, dict):
            years[str(year)] = {m: _num(metrics.get(m)) for m in METRICS}
    data["years"] = years
    return data


# ---------------------------------------------------------------------------
# Layer 2: deterministic red-flag rules (Schilit-style)
# ---------------------------------------------------------------------------

def _g(prev: float | None, cur: float | None) -> float | None:
    if prev in (None, 0) or cur is None:
        return None
    return (cur - prev) / abs(prev)


def run_rules(data: dict) -> list[Flag]:
    flags: list[Flag] = []
    years = sorted(data.get("years", {}).keys())
    if len(years) >= 2:
        prev, cur = data["years"][years[-2]], data["years"][years[-1]]

        rev_g = _g(prev.get("revenue"), cur.get("revenue"))
        rec_g = _g(prev.get("receivables"), cur.get("receivables"))
        inv_g = _g(prev.get("inventory"), cur.get("inventory"))
        ni_g = _g(prev.get("net_income"), cur.get("net_income"))
        ocf_g = _g(prev.get("operating_cash_flow"), cur.get("operating_cash_flow"))
        gw_g = _g(prev.get("goodwill"), cur.get("goodwill"))
        prov_g = _g(prev.get("provisions"), cur.get("provisions"))

        if rev_g is not None and rec_g is not None and rev_g > 0 and rec_g > rev_g * 1.5 and rec_g > 0.2:
            flags.append(Flag("revenue", "high", "Receivables outpacing revenue",
                              f"Revenue grew {rev_g:.0%} while receivables grew {rec_g:.0%} — "
                              "possible aggressive recognition or channel stuffing.", 80))
        if rev_g is not None and rev_g > 0.4:
            flags.append(Flag("revenue", "medium", "Unusually rapid revenue growth",
                              f"Revenue grew {rev_g:.0%} year-over-year — verify sustainability "
                              "and recognition policy.", 65))
        if ni_g is not None and ocf_g is not None and ni_g > 0 and ocf_g < 0:
            flags.append(Flag("cash_flow", "high", "Earnings up, operating cash flow down",
                              f"Net income rose {ni_g:.0%} while operating cash flow fell "
                              f"{abs(ocf_g):.0%} — a classic earnings-quality warning.", 85))
        ni, ocf = cur.get("net_income"), cur.get("operating_cash_flow")
        if ni and ocf and ni > 0 and 0 < ocf < 0.6 * ni:
            flags.append(Flag("cash_flow", "medium", "Weak cash conversion",
                              f"Operating cash flow ({ocf:,.0f}) is only {ocf / ni:.0%} of net "
                              f"income ({ni:,.0f}).", 75))
        if inv_g is not None and rev_g is not None and inv_g > max(rev_g * 1.5, 0.2):
            flags.append(Flag("assets", "medium", "Inventory building faster than sales",
                              f"Inventory grew {inv_g:.0%} vs revenue {rev_g:.0%} — possible "
                              "obsolescence or overproduction to absorb overheads.", 70))
        if gw_g is not None and gw_g > 0.3:
            flags.append(Flag("assets", "medium", "Goodwill spike",
                              f"Goodwill grew {gw_g:.0%} — review acquisition accounting and "
                              "purchase-price allocation (IFRS 3 / IAS 36).", 70))
        if prov_g is not None and rev_g is not None and prov_g < -0.2 and rev_g > 0:
            flags.append(Flag("liabilities", "medium", "Declining provisions while growing",
                              f"Provisions fell {abs(prov_g):.0%} while revenue rose — possible "
                              "liability understatement or cookie-jar release.", 65))

    # latest-year point checks
    if years:
        cur = data["years"][years[-1]]
        rev, rec = cur.get("revenue"), cur.get("receivables")
        if rev and rec and rev > 0:
            days = rec / rev * 365
            if days > 75:
                flags.append(Flag("revenue", "medium", "Slow receivable collection",
                                  f"Receivable days ≈ {days:.0f} (typical range 30–60) — "
                                  "potential collection problem.", 70))
        ca, cl = cur.get("current_assets"), cur.get("current_liabilities")
        if ca and cl and cl > 0 and ca / cl < 1.0:
            flags.append(Flag("liabilities", "medium", "Liquidity strain",
                              f"Current ratio {ca / cl:.2f} — current liabilities exceed "
                              "current assets.", 80))

    if data.get("non_gaap_metrics"):
        metrics = ", ".join(map(str, data["non_gaap_metrics"][:5]))
        flags.append(Flag("governance", "medium", "Non-GAAP metrics in use",
                          f"Adjusted measures reported ({metrics}) — verify each adjustment "
                          "is justified and reconciled to statutory figures.", 60))
    if data.get("policy_changes"):
        changes = "; ".join(map(str, data["policy_changes"][:3]))
        flags.append(Flag("governance", "high", "Accounting policy change",
                          f"Disclosed change(s): {changes} — assess earnings impact and "
                          "timing.", 75))
    if data.get("capitalised_costs_mentioned"):
        flags.append(Flag("expenses", "medium", "Cost capitalisation signals",
                          "Document mentions capitalised development or deferred costs — "
                          "check whether operating expenses are being parked on the "
                          "balance sheet.", 60))
    if data.get("acquisitions_mentioned"):
        flags.append(Flag("assets", "low", "Acquisition activity",
                          "Acquisitions mentioned — review purchase-price allocation, "
                          "goodwill and any bargain-purchase gains.", 55))
    return flags


# ---------------------------------------------------------------------------
# Layer 3: risk scoring
# ---------------------------------------------------------------------------

_SEVERITY_POINTS = {"low": 10, "medium": 20, "high": 35}


def score(flags: list[Flag]) -> tuple[dict[str, int], int, str]:
    scores = {c: 15 for c in CATEGORIES}
    for f in flags:
        scores[f.category] = min(100, scores[f.category] + _SEVERITY_POINTS[f.severity])
    vals = list(scores.values())
    overall = round(0.6 * max(vals) + 0.4 * (sum(vals) / len(vals)))
    label = "Low Risk" if overall < 40 else "Moderate Risk" if overall < 70 else "High Risk"
    return scores, overall, label


# ---------------------------------------------------------------------------
# Layer 4: fraud timeline
# ---------------------------------------------------------------------------

def build_timeline(data: dict) -> list[tuple[str, list[str]]]:
    years = sorted(data.get("years", {}).keys())
    timeline: list[tuple[str, list[str]]] = []
    for i, year in enumerate(years):
        cur = data["years"][year]
        signals: list[str] = []
        if i > 0:
            prev = data["years"][years[i - 1]]
            for metric, arrow_up, arrow_dn, bad_up in [
                ("receivables", "Receivables ↑", "Receivables ↓", True),
                ("goodwill", "Goodwill ↑", "Goodwill ↓", True),
                ("operating_cash_flow", "Operating cash flow ↑", "Operating cash flow ↓", False),
                ("cash", "Cash ↑", "Cash ↓", False),
            ]:
                g = _g(prev.get(metric), cur.get(metric))
                if g is None:
                    continue
                if g > 0.25 and bad_up:
                    signals.append(f"{arrow_up} {g:.0%}")
                elif g < -0.15 and not bad_up:
                    signals.append(f"{arrow_dn} {abs(g):.0%}")
        timeline.append((year, signals or ["No notable signals"]))
    return timeline


# ---------------------------------------------------------------------------
# Layer 5: external validation (best effort, never blocks)
# ---------------------------------------------------------------------------

def external_checks(company: str | None) -> list[dict]:
    if not company:
        return []
    from src.tools import external
    results = []
    for fn in (external.sec_edgar_search, external.companies_house_search):
        try:
            results.append(fn(company))
        except Exception as e:
            results.append({"error": str(e)})
    return results


# ---------------------------------------------------------------------------
# Layer 6: explainability narrative
# ---------------------------------------------------------------------------

NARRATIVE_SYSTEM = """You are a forensic accountant writing the explainability section
of an investigation report. Given deterministic red flags and figures (JSON),
write: numbered reasons ("Reason 1: ...", one per flag, quantified where the
data allows), then an "Overall Assessment" paragraph of 2-3 sentences in a
measured, audit-adjacent register. Findings are indicators for review, never
proof of fraud — say so."""


def explain(flags: list[Flag], data: dict) -> str:
    if not flags:
        return ("No red flags were triggered by the deterministic checks. This does not "
                "prove the absence of manipulation — extend the document set (multiple "
                "years, audit report, cash flow statement) for stronger coverage.")
    llm = get_llm("verifier")
    payload = {"flags": [f.as_dict() for f in flags], "figures": data.get("years", {})}
    resp = llm.invoke([
        SystemMessage(content=NARRATIVE_SYSTEM),
        HumanMessage(content=json.dumps(payload, indent=2)),
    ])
    return re.sub(r"<think>.*?</think>", "", resp.content, flags=re.DOTALL).strip()


# ---------------------------------------------------------------------------
# KPIs + pipeline entry point
# ---------------------------------------------------------------------------

def compute_kpis(data: dict) -> dict:
    years = sorted(data.get("years", {}).keys())
    if not years:
        return {}
    cur = data["years"][years[-1]]
    kpis: dict = {"Fiscal year": years[-1]}
    for label, key in [("Revenue", "revenue"), ("Net Income", "net_income"),
                       ("Cash", "cash"), ("Debt", "total_debt")]:
        v = cur.get(key)
        kpis[label] = f"{v:,.0f}" if v is not None else "n/a"

    def ratio(a, b):
        va, vb = cur.get(a), cur.get(b)
        return round(va / vb, 2) if va is not None and vb not in (None, 0) else None

    kpis["Current Ratio"] = ratio("current_assets", "current_liabilities") or "n/a"
    ca, inv, cl = cur.get("current_assets"), cur.get("inventory"), cur.get("current_liabilities")
    kpis["Quick Ratio"] = (round((ca - inv) / cl, 2)
                           if None not in (ca, inv, cl) and cl != 0 else "n/a")
    roe = ratio("net_income", "shareholders_equity")
    roa = ratio("net_income", "total_assets")
    kpis["ROE"] = f"{roe:.0%}" if isinstance(roe, float) else "n/a"
    kpis["ROA"] = f"{roa:.0%}" if isinstance(roa, float) else "n/a"
    return kpis


def run(retriever, doc_ids: list[str] | None = None) -> ForensicResult:
    from src.graph.workflow import _balanced_retrieve

    activity: list[str] = []
    query = ("revenue receivables inventory net income operating cash flow "
             "goodwill cash debt provisions equity assets auditor company")
    evidence = _balanced_retrieve(retriever, query, doc_ids, per_doc=6)[:18]
    activity.append(f"Evidence gathered — {len(evidence)} excerpts across "
                    f"{len({r.chunk.doc_id for r in evidence})} documents")

    data = extract_figures(evidence)
    n_years = len(data.get("years", {}))
    activity.append(f"Financial statements extracted — {n_years} fiscal year(s)")

    flags = run_rules(data)
    activity.append(f"Manipulation checks run — {len(flags)} red flag(s)")

    category_scores, overall, label = score(flags)
    activity.append(f"Fraud-risk scoring completed — overall {overall}/100 ({label})")

    timeline = build_timeline(data)
    activity.append("Fraud timeline assembled")

    external = external_checks(data.get("company"))
    activity.append("External filings checked (SEC EDGAR / Companies House)"
                    if external else "External check skipped — company name not identified")

    narrative = explain(flags, data)
    activity.append("Investigation narrative generated")

    return ForensicResult(
        profile={k: data.get(k) for k in ("company", "industry", "auditor",
                                          "exchange", "country", "market_cap", "currency")},
        figures=data.get("years", {}),
        kpis=compute_kpis(data),
        flags=flags,
        category_scores=category_scores,
        overall_risk=overall,
        risk_label=label,
        timeline=timeline,
        external=external,
        narrative=narrative,
        activity=activity,
    )