"""Bounded monitoring agent: search -> fetch (fresh first) -> extract -> validate -> compare -> report.

Planner is deterministic by default (no LLM spend). All comparison, deltas, conflicts,
staleness and geometry are deterministic. Optional LLM extraction via Config(extractor=...); see providers.json and SQ_LLM_* env vars.
Usage: python agent.py --demo   (runs one offline fixture end to end)
"""
from __future__ import annotations
import argparse
import json
import math
import time
from dataclasses import dataclass

from bounds import IMPROVEMENT_MATERIAL, float_near_bound_margin, Bound, classify, find_conflicts, rel_score, load_bounds
from packing import parse_squares, validate
from tools import ToolError, eval_side
from baseline_regex import extract_naive


@dataclass
class Config:
    mode: str = "agent"        # "agent" | "regex_guards" | "baseline"
    retries: int = 1           # retry once on tool error, then degrade
    backoff_s: float = 0.0     # set >0 for real network
    metric_tol: float = 5e-4   # tolerance for reported relative score vs recomputed (rounding)
    extractor: object = None   # interfaces.Extractor; default RuleExtractor
    planner: object = None     # interfaces.Planner; default FixedPlanner


def _call(fn, arg, state, cfg):
    """Execute a tool with budget accounting and retry/backoff. Returns (ok, result)."""
    attempts = 1 + (cfg.retries if cfg.mode != "baseline" else 0)
    last = None
    for k in range(attempts):
        if state["tool_calls"] >= state["max_steps"]:
            state["budget_exhausted"] = True
            return False, last or "budget exhausted"
        state["tool_calls"] += 1
        try:
            return True, fn(arg)
        except ToolError as e:
            last = str(e)
            if cfg.backoff_s and k + 1 < attempts:
                time.sleep(cfg.backoff_s * (2 ** k))
    state["errors"].append(last)
    return False, last


def _validate_claim(c):
    d = c["packing"]
    try:
        side = eval_side(d["side"])
        sq = parse_squares(d["squares"])
        return validate(sq, side, expected_n=c["n"])
    except Exception as e:
        return {"valid": False, "mode": "unparsed", "proof": False, "issues": [f"parse error: {e}"]}


def run(prompt, tools, bounds, max_steps=8, cfg: Config | None = None):
    cfg = cfg or Config()
    wanted = set(prompt.get("n", [])) or None
    state = {"tool_calls": 0, "max_steps": max_steps, "errors": [], "budget_exhausted": False,
             "notes": [], "claims": [], "fetched": []}

    from interfaces import FixedPlanner, RegexGuardExtractor, RuleExtractor
    planner = cfg.planner or FixedPlanner()
    extractor = cfg.extractor or (RegexGuardExtractor() if cfg.mode == "regex_guards" else RuleExtractor())
    agentic = cfg.mode != "baseline"  # regex_guards gets every agent guard; only the extractor differs
    # plan step 1: search (one tool call per planned query; all count against the budget)
    results, any_ok = [], False
    for q in (planner.plan(prompt) if cfg.mode != "baseline" else FixedPlanner().plan(prompt)):
        ok, r = _call(tools.search, q, state, cfg)
        if ok:
            any_ok = True; results += r
    if any_ok and not results:
        state["errors"].append("search returned no sources")
    # plan step 2: fresh first; agent also skips duplicate URLs
    if cfg.mode != "baseline":
        uniq = {r["url"]: r for r in results}
        results = sorted(uniq.values(), key=lambda r: r.get("date", ""), reverse=True)
    for i, meta in enumerate(results):
        if state["tool_calls"] >= max_steps:
            state["budget_exhausted"] = True
            state["errors"].append(f"budget exhausted: {len(results) - i} source(s) not fetched")
            break
        ok, text = _call(tools.fetch, meta["url"], state, cfg)
        if not ok:
            continue
        state["fetched"].append(meta["url"])
        if cfg.mode == "baseline":
            cl, notes = extract_naive(text, meta)
        else:
            cl, notes = extractor.extract(text, meta)
        state["claims"] += cl
        state["notes"] += [f"{meta['url']}: {x}" for x in notes]
    return build_report(state, bounds, wanted, cfg)


def build_report(state, bounds, wanted, cfg):
    claims = [c for c in state["claims"] if wanted is None or c["n"] in wanted]
    rep = {"timestamp": "run-local",
           "improvements": [], "candidates": [], "conflicts": [], "partial": [], "unchanged": [], "no_claim": [],
           "worse": [], "invalid": [], "rejected": [], "stale": [], "no_baseline": [],
           "metric_mismatch": [], "errors": list(state["errors"]), "notes": state["notes"],
           "tool_calls": state["tool_calls"], "budget_exhausted": state["budget_exhausted"],
           "sources_cited": sorted(set(state["fetched"])), "summary": ""}

    agentic = cfg.mode != "baseline"
    if agentic:
        rep["rejected"] = [{k: c[k] for k in ("n", "s", "source_url", "rejected")}
                           for c in claims if c.get("rejected")]
        claims = [c for c in claims if not c.get("rejected")]
        # staleness: per (source_key, n) keep only the newest claim
        newest = {}
        for c in claims:
            k = (c["source_key"], c["n"])
            if k not in newest or c["date"] > newest[k]["date"]:
                newest[k] = c
        keep = {id(c) for c in newest.values()}
        stale = [c for c in claims if id(c) not in keep and
                 c["date"] < newest[(c["source_key"], c["n"])]["date"]]
        stale_ids = {id(c) for c in stale}
        rep["stale"] = [{k: c[k] for k in ("n", "s", "source_url", "date")} for c in stale]
        claims = [c for c in claims if id(c) not in stale_ids]

    for c in claims:
        b = bounds.get(c["n"])
        val = _validate_claim(c) if (agentic and c["coordinates_available"]) else None
        if cfg.mode == "baseline":  # baseline trusts any lower number
            val = {"valid": True, "mode": "none", "proof": False}
        status, delta = classify(c["s"], b.s if b else None, val)
        if agentic and status == "improvement" and not val.get("proof") and \
                -delta <= float_near_bound_margin(c["n"]):
            status = "unverified"  # float check can't separate this from the bound; only exact/interval promotes
            c = dict(c, _float_near_bound=True)
        if agentic and status in ("improvement", "unverified") and \
                c.get("s_uncertainty", 0) and c["s"] + c["s_uncertainty"] >= b.s - IMPROVEMENT_MATERIAL:
            status = "unchanged"  # printed precision cannot distinguish the claim from the bound
        out = {k: c[k] for k in ("n", "s", "source_url", "source_type", "date", "extracted_text",
                                 "coordinates_available")}
        out.update(status=status, delta_vs_bound=delta, bound_source=b.source if b else None,
                   validation=None if val is None else {k: val.get(k) for k in ("valid", "mode", "proof", "issues")})
        if c.get("score") is not None and c["s"] is not None:
            rec = rel_score(c["n"], c["s"])
            if abs(rec - c["score"]) > cfg.metric_tol:
                rep["metric_mismatch"].append({**out, "reported_score": c["score"], "recomputed": rec})
        bucket = {"improvement": "improvements", "unverified": "candidates", "invalid": "invalid",
                  "worse": "worse", "no_baseline": "no_baseline", "unchanged": "unchanged", "partial": None}[status]
        if bucket:
            rep[bucket].append(out)
        if c.get("_float_near_bound"):
            out["note"] = "float-validated but within accumulated tolerance of bound; needs exact/interval proof"
        if not c["coordinates_available"] and agentic:
            rep["partial"].append(out)

    rep["conflicts"] = find_conflicts(claims) if agentic else find_conflicts(
        [dict(c, source_url=c["source_url"] + f"#{i}") for i, c in enumerate(claims)])
    # unchanged = claim present and within +-material of the bound (list of claim dicts, like other buckets).
    # no_claim = wanted n with no surviving claim at all ("not checked", NOT "checked, no change").
    rep["no_claim"] = sorted((wanted or set()) - {c["n"] for c in claims})
    rep["summary"] = summarize(rep)
    return rep


def summarize(rep):
    lines = [f"{len(rep['improvements'])} validated improvement(s); {len(rep['candidates'])} unverified "
             f"candidate(s); {len(rep['unchanged'])} unchanged; {len(rep['worse'])} worse; "
             f"{len(rep['no_claim'])} n with no claim; {len(rep['conflicts'])} conflict(s); {len(rep['invalid'])} invalid; "
             f"{len(rep['errors'])} error(s); {rep['tool_calls']} tool call(s)."]
    for c in rep["improvements"]:
        proof = "exact proof" if c["validation"]["proof"] else "float check, not a proof"
        lines.append(f"- n={c['n']}: s={c['s']} (delta {c['delta_vs_bound']:+.3g}, {proof}) [{c['source_url']}]")
    for c in rep["candidates"]:
        lines.append(f"- CANDIDATE n={c['n']}: s={c['s']} below bound, unproven ({c.get('note') or 'no coordinates'}) [{c['source_url']}]")
    for c in rep["conflicts"]:
        lines.append(f"- CONFLICT n={c['n']}: spread {c['spread']:.3g} across {len(c['claims'])} claims")
    if not rep["sources_cited"]:
        lines.append("- No sources retrieved: refusing to assert progress (low confidence).")
    return "\n".join(lines)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--demo", action="store_true")
    ap.add_argument("--scenario", default="scenarios/s10_true_improvement_exact.json")
    a = ap.parse_args()
    from tools import MockTools
    sc = json.load(open(a.scenario))
    bounds = load_bounds("bounds.json")
    for n, s in sc.get("bounds_override", {}).items():
        bounds[int(n)] = Bound(int(n), s, "synthetic (scenario override)", "", False)
    rep = run(sc["prompt"], MockTools(sc["tools"]), bounds, sc.get("max_steps", 8))
    print(rep["summary"])
    print(json.dumps({k: rep[k] for k in ("improvements", "candidates", "errors", "tool_calls")}, indent=1, default=str)[:3000])