"""E4: repeated handoffs (relay chains) with an evaluator-only fact ledger.

Chain: agent_1 reads the document and writes m_1 (A4 prose or A5 structured, question-blind). agent_i (i>1) reads only
m_{i-1} (closed) or m_{i-1} plus metered source lookups (repairable) and writes m_i under the same budget. After hops
h in HOPS the receiver answers all questions from m_h. Retention/distortion is scored from the ledger (never shown to agents).
"""
from __future__ import annotations
import json, re, time, yaml
from pathlib import Path
import mlx.core as mx
from .util import ROOT, read_json, sha256_json, sha256_text, utc_now, read_jsonl
from .records import new_record, write_record, trial_key
from .data.prepare import load_doc, load_key
from .scoring import score_synthetic, normalize_answer, extract_numbers
from .prompts import SYSTEM_RECEIVER, receiver_user_text, SUMMARY_SYSTEM, EVIDENCE_SYSTEM
from .backends.mlx_backend import MLXModel, reset_peak
from .methods import text as TM

HARDWARE_ID = "laptop-m5max-48g"

RELAY_SUMMARY_USER = ("You received the following notes from a colleague about a document you cannot see. Rewrite them as a faithful summary of at most "
                      "{budget} tokens for the next colleague, preserving every concrete fact (names, identifiers, numbers, units, dates, statuses, negations) "
                      "exactly as written. Do not add facts that are not in the notes.{repair}\n\nNotes:\n{msg}\n\nSummary:")
RELAY_EVIDENCE_USER = ("You received the following structured evidence (JSON) from a colleague about a document you cannot see. Re-emit it as a JSON object with the same "
                       "schema (keys facts, optional notes, optional unknowns), under {budget} tokens, keeping every fact and value exactly; drop nothing unless the "
                       "budget forces you to, and never invent values.{repair}\n\nEvidence:\n{msg}\n\nJSON:")

def _ledger_metrics(message: str, facts: list[dict], doc_text: str) -> dict:
    """Retention: a fact counts as retained if its value (normalized) appears in the message. Unsupported additions:
    numbers with >=4 significant digits, identifiers (XX-00X-0000) and 4-digit-year dates in the message that do not
    occur in the document."""
    # structured messages carry character-offset spans ("span": [a, b]); strip them so offsets are not counted as numbers
    message = re.sub(r'"span":\s*\[[^\]]*\]', '"span": []', message)
    msg_n = normalize_answer(message); msg_nc = message.replace(",", "")
    retained = 0; by_kind = {}
    for f in facts:
        v = f["value"]
        ok = normalize_answer(v) in msg_n or v.replace(",", "") in msg_nc
        by_kind.setdefault(f["kind"], [0, 0]); by_kind[f["kind"]][1] += 1
        if ok:
            retained += 1; by_kind[f["kind"]][0] += 1
    doc_nums = {str(int(x)) for x in extract_numbers(doc_text) if abs(x) >= 1000}
    msg_nums = {str(int(x)) for x in extract_numbers(message) if abs(x) >= 1000}
    unsupported_nums = msg_nums - doc_nums
    doc_ids = set(re.findall(r"\b[A-Z]{2}-\d{2}[A-Z]-\d{4}\b", doc_text)); msg_ids = set(re.findall(r"\b[A-Z]{2}-\d{2}[A-Z]-\d{4}\b", message))
    unsupported_ids = msg_ids - doc_ids
    # distorted amounts: a project code co-occurring (same sentence) with an unsupported number
    distorted = 0
    for sent in re.split(r"(?<=[.!?])\s+|\n+", message):
        if re.search(r"\b[A-Z][a-z]+-\d{1,2}\b", sent) and any(n in sent.replace(",", "") for n in unsupported_nums):
            distorted += 1
    return {"facts_total": len(facts), "facts_retained": retained, "retention_rate": retained / max(1, len(facts)),
            "retention_by_kind": {k: v[0] / v[1] for k, v in by_kind.items()},
            "unsupported_numbers": len(unsupported_nums), "unsupported_identifiers": len(unsupported_ids), "distorted_sentences": distorted,
            "message_bytes": len(message.encode("utf-8"))}

def _repair_lookup(prev_msg: str, doc_text: str, max_bytes: int, max_lookups: int) -> tuple[str, dict]:
    """Repairable chain: the agent re-reads source sentences for entity codes mentioned in the notes (metered)."""
    codes = []
    for c in re.findall(r"\b[A-Z][a-z]+-\d{1,2}\b", prev_msg):
        if c not in codes: codes.append(c)
    sents = re.split(r"(?<=[.!?])\s+", doc_text)
    fetched, nb, n = [], 0, 0
    for c in codes:
        if n >= max_lookups: break
        hits = [s for s in sents if c in s]
        if not hits: continue
        chunk = " ".join(hits)
        b = len(chunk.encode("utf-8"))
        if nb + b > max_bytes: continue
        fetched.append(f"[source passages for {c}] {chunk}"); nb += b; n += 1
    return "\n".join(fetched), {"lookups": n, "fetched_bytes": nb, "codes_mentioned": len(codes)}

def run_chains(config_path: str, resume: bool = False) -> int:
    cfg = yaml.safe_load(open(config_path))
    mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]
    models = {}
    def get(a):
        if a not in models: models[a] = MLXModel(mcfg[a]["repo"])
        return models[a]
    man = read_json(ROOT / "data/manifests/splits.json")
    items = [it for it in man["items"] if it["family"] == "synthetic" and it["split"] in cfg["splits"] and it.get("target_tokens") in cfg["lengths"]]
    if cfg.get("limit_docs"): items = items[: cfg["limit_docs"]]
    out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"
    log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log"
    done = {r["trial_id"] for r in read_jsonl(out_path) if r["status"] == "success"} if (resume and out_path.exists()) else set()
    protocol_hash = cfg.get("protocol_hash", "unfrozen"); config_hash = sha256_json(cfg)
    def log(m):
        line = f"{utc_now()} {m}"; print(line, flush=True)
        with open(log_path, "a") as f: f.write(line + "\n")
    log(f"chains start {cfg['experiment_id']} items={len(items)} done={len(done)}")
    HOPS = cfg["hops"]; budget = cfg["budget_tokens"]
    for it in items:
        doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic")
        qkeys = {q["question_id"]: q for q in key["questions"]}
        for chain in cfg["chains"]:
            rep, mode, agents = chain["representation"], chain["mode"], chain["agents"]  # agents: list of aliases cycled
            cid = f"{rep}|{mode}|{'-'.join(agents)}"
            # skip if all hop trials done
            keys_needed = [trial_key(cfg["experiment_id"], it["document_id"], q["question_id"], rep, {"mode": mode, "agents": agents, "budget": budget}, agents[0], agents[-1], "blind", protocol_hash, 0, "chain", h) for h in HOPS for q in doc["questions"]]
            if all(k in done for k in keys_needed):
                continue
            msg = None; hop_msgs = {}; hop_meta = {}
            t_chain = 0.0; fetched_total = 0; lookups_total = 0
            for h in range(1, max(HOPS) + 1):
                agent = get(agents[(h - 1) % len(agents)])
                t0 = time.perf_counter()
                repair_txt, rinfo = "", {"lookups": 0, "fetched_bytes": 0}
                if h == 1:
                    m = TM.a4_summary(doc["text"], budget, agent, None) if rep == "A4" else TM.a5_structured(doc["text"], budget, agent, None)
                    text = m.text; usage = m.sender_usage
                else:
                    if mode == "repairable":
                        repair_txt, rinfo = _repair_lookup(msg, doc["text"], cfg.get("max_fetch_bytes", 1500), cfg.get("max_lookups", 4))
                    rep_clause = ("\n\nYou may also use these passages re-read from the original source:\n" + repair_txt) if repair_txt else ""
                    if rep == "A4":
                        user = RELAY_SUMMARY_USER.format(budget=budget, msg=msg, repair=rep_clause)
                        text, usage, _ = TM._sender_generate(agent, SUMMARY_SYSTEM, user, budget)
                    else:
                        user = RELAY_EVIDENCE_USER.format(budget=budget, msg=msg, repair=rep_clause)
                        raw, usage, _ = TM._sender_generate(agent, EVIDENCE_SYSTEM, user, budget)
                        obj = TM._parse_json_loose(raw)
                        text = json.dumps(obj, ensure_ascii=False) if isinstance(obj, dict) else raw
                hop_ms = (time.perf_counter() - t0) * 1000; t_chain += hop_ms
                fetched_total += rinfo["fetched_bytes"]; lookups_total += rinfo["lookups"]
                msg = text; hop_msgs[h] = text
                hop_meta[h] = {"hop_ms": hop_ms, "agent": agent.info.repo_id, "usage": usage, "repair": rinfo, "ledger": _ledger_metrics(text, key["facts"], doc["text"])}
                if h in HOPS:
                    receiver = get(agents[(h - 1) % len(agents)])  # the last agent in the chain answers
                    for q in doc["questions"]:
                        tk = trial_key(cfg["experiment_id"], it["document_id"], q["question_id"], rep, {"mode": mode, "agents": agents, "budget": budget}, agents[0], agents[-1], "blind", protocol_hash, 0, "chain", h)
                        if tk in done: continue
                        rec = new_record(experiment_id=cfg["experiment_id"], trial_id=tk, run_group_id=f"{cfg['experiment_id']}:{cid}:h{h}", protocol_hash=protocol_hash, config_hash=config_hash,
                                         task_family="synthetic", document_id=it["document_id"], document_hash=it["document_hash"], question_id=q["question_id"], split=it["split"], cluster_id=it["cluster_id"],
                                         sender_model_revision="|".join(f"{get(a).info.repo_id}@{get(a).info.revision}" for a in agents), receiver_model_revision=f"{receiver.info.repo_id}@{receiver.info.revision}",
                                         tokenizer_hashes={"receiver": receiver.info.tokenizer_hash}, backend_versions={"mlx": mx.__version__, "mlx_lm": receiver.info.mlx_lm_version}, hardware_id=HARDWARE_ID,
                                         device_state={}, weight_precision="bf16", method_id=rep, implementation_fidelity="custom", compression_config={"mode": mode, "agents": agents, "budget_tokens": budget},
                                         question_access="blind", topology=f"chain_{mode}", handoff_index=h, reuse_count=len(doc["questions"]), cache_hit=False, cache_compatibility_hash=None, random_seed=0,
                                         source_tokens_sender=None, source_tokens_receiver=len(receiver.encode(doc["text"])), message_utf8_bytes=len(text.encode("utf-8")), serialized_payload_bytes=len(text.encode("utf-8")),
                                         cache_tensor_bytes=None, side_information_bytes=0, fetched_bytes=fetched_total, timing_spans={"chain_ms": t_chain, "hops": {str(k): v["hop_ms"] for k, v in hop_meta.items()}},
                                         input_usage={"lookups": lookups_total}, output_usage={}, provider_cached_usage=None, process_memory=None, backend_memory=None, peak_prefill_memory=None, peak_decode_memory=None,
                                         mapper_memory=None, cost_with_source=None, evidence_artifact={"ledger": hop_meta[h]["ledger"], "ledger_by_hop": {str(k): v["ledger"] for k, v in hop_meta.items()}},
                                         reference_key_id=f"keys/synthetic/{it['document_id']}.json", status="implementation_error", error_code=None, log_artifact=str(log_path.relative_to(ROOT)), payload_hash=sha256_text(text), output_hash=None)
                        try:
                            ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text("Notes received after %d handoff(s)" % h, text, q["question"], "synthetic")}])
                            t1 = time.perf_counter(); g = receiver.generate_text(ids, max_tokens=48)
                            rec.update({"prediction_artifact": g["text"], "output_hash": sha256_text(g["text"]), "query_tokens": len(ids), "message_tokens": len(receiver.encode(text)),
                                        "ttft_ms": g["ttft_ms"], "total_wall_ms": (time.perf_counter() - t1) * 1000 + t_chain, "output_usage": {"output_tokens": g["output_tokens"]},
                                        "task_scores": score_synthetic(g["text"], qkeys[q["question_id"]]), "status": "success" if g["text"].strip() else "invalid_output"})
                            rec["timing_spans"].update({"receiver_prefill_ms": g["prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"]})
                        except Exception as e:
                            rec["error_code"] = f"{type(e).__name__}: {str(e)[:200]}"
                        write_record(rec, out_path)
                    L = hop_meta[h]["ledger"]
                    log(f"{cid:>28s} {it['document_id']} hop={h} retention={L['retention_rate']:.2f} unsupported={L['unsupported_numbers']}+{L['unsupported_identifiers']} bytes={L['message_bytes']} fetched={fetched_total}")
            # store hop messages for audit
            Path(ROOT / "runs/raw/chains").mkdir(parents=True, exist_ok=True)
            (ROOT / "runs/raw/chains" / f"{cfg['experiment_id']}_{it['document_id']}_{cid.replace('|','_').replace('-','_')}.json").write_text(json.dumps({"hops": hop_msgs, "meta": hop_meta}, ensure_ascii=False, indent=1, default=str))
        mx.clear_cache()
    log("chains end")
    return 0