# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Deliberation protocols. Each coroutine returns a DebateTranscript for ONE property.

F  independent(...)            round 0 only, no communication (E = same with roles=[generalist]).
G  structured_debate(...)      round 0 independent → for each round r: critiques (CritiqueSet) then revised
                               estimates (AgentEstimate) with the anonymised committee state; pre-specified
                               stopping when the IQR of log points moves < 1 % (recorded, rounds actually run
                               reported in ``transcript.rounds``).
H  structured_debate(..., blinded_skeptic=True, evidence_verification=True)
                               the skeptic never joins the rounds; it sees the packet + final estimates only
                               and its verdict is appended; deterministic verification flags are injected
                               into the next prompt of each flagged agent.
I  structured_debate(..., placebo=True) with placebo roles (same mechanics, generic / permuted prompts).
J  cost_matched_single(...)    generalist sampled k times independently + a 2-step self-refine chain.

Agents of the same round run concurrently (asyncio.gather); the provider semaphore bounds the calls.
Failed turns (error set, estimate None) are excluded from the committee state and from aggregation.
"""
from __future__ import annotations

import asyncio
import time
from typing import Any

from ..agents.packets import committee_state, render_user_prompt, role_view, verify_evidence, visible_ids
from ..agents.roles import Role, get_role
from ..agents.runner import ProviderPool, run_turn
from ..schemas import AgentEstimate, AgentTurn, CritiqueSet, DebateTranscript, EvidencePacket
from .aggregation import aggregate_intervals, aggregate_points, interval_90
from .stopping import IQR_TOL, iqr_log_points, should_stop


def _latest_estimates(turns: list[AgentTurn]) -> dict[str, AgentEstimate]:
    """role -> last successful estimate (estimate or revision kinds), in role order of first appearance."""
    out: dict[str, AgentEstimate] = {}
    for t in turns:
        if t.kind in ("estimate", "revision") and t.estimate is not None:
            out[t.role] = t.estimate
    return out


def _finalize(tr: DebateTranscript, turns: list[AgentTurn], t0: float, weights: dict[str, Any] | None, final: dict[str, AgentEstimate] | None = None) -> DebateTranscript:
    tr.turns = turns
    est = final if final is not None else _latest_estimates(turns)
    tr.final_estimates = {r: e.point_estimate for r, e in est.items()}
    tr.final_quantiles = {r: [e.q05, e.q25, e.q50, e.q75, e.q95] for r, e in est.items()}
    tr.aggregated = aggregate_points(tr.final_estimates, weights)
    tr.aggregated_interval = interval_90(aggregate_intervals(tr.final_quantiles))
    tr.total_cost_usd = round(sum(t.cost_usd for t in turns), 6)
    tr.total_tokens = sum(t.input_tokens + t.output_tokens for t in turns)
    tr.total_latency_s = round(time.perf_counter() - t0, 3)
    return tr


async def _gather_turns(coros: list[Any]) -> list[AgentTurn]:
    return list(await asyncio.gather(*coros))


def _model_for(assignment: dict[str, str] | str, role_id: str) -> str:
    if isinstance(assignment, str):
        return assignment
    try:
        return assignment[role_id]
    except KeyError as e:
        raise KeyError(f"no model assigned to role {role_id!r}") from e


# ---------------------------------------------------------------------------------------------------
# F / E
# ---------------------------------------------------------------------------------------------------
async def independent(
    packet: EvidencePacket,
    roles: list[Role],
    committee_assignment: dict[str, str],
    pool: ProviderPool,
    *,
    arm: str = "F_independent",
    committee: str = "all_anthropic",
    effort: str = "medium",
    max_tokens: int = 4096,
    weights: dict[str, Any] | None = None,
) -> DebateTranscript:
    t0 = time.perf_counter()
    tr = DebateTranscript(tx_id=packet.subject.tx_id, arm=arm, committee=committee, protocol="independent", rounds=0, packet_hash=packet.content_hash(), turns=[], assignment=dict(committee_assignment))
    turns = await _gather_turns(
        [
            run_turn(pool.get(_model_for(committee_assignment, r.id)), r, render_user_prompt(role_view(packet, r), None, kind="estimate"), AgentEstimate, _model_for(committee_assignment, r.id), effort, 0, "estimate", max_tokens=max_tokens)
            for r in roles
        ]
    )
    return _finalize(tr, turns, t0, weights)


# ---------------------------------------------------------------------------------------------------
# G / H / I
# ---------------------------------------------------------------------------------------------------
async def structured_debate(
    packet: EvidencePacket,
    roles: list[Role],
    assignment: dict[str, str],
    pool: ProviderPool,
    rounds: int = 2,
    *,
    blinded_skeptic: bool = False,
    evidence_verification: bool = False,
    placebo: bool = False,
    arm: str = "G_debate",
    committee: str = "all_anthropic",
    effort: str = "medium",
    max_tokens: int = 4096,
    weights: dict[str, Any] | None = None,
    iqr_tol: float = IQR_TOL,
    skeptic_role_id: str = "skeptic",
) -> DebateTranscript:
    t0 = time.perf_counter()
    protocol = "placebo_debate" if placebo else ("debate_skeptic_verification" if (blinded_skeptic or evidence_verification) else "structured_debate")
    tr = DebateTranscript(tx_id=packet.subject.tx_id, arm=arm, committee=committee, protocol=protocol, rounds=rounds, packet_hash=packet.content_hash(), turns=[], assignment=dict(assignment))
    skeptic: Role | None = None
    debaters = list(roles)
    if blinded_skeptic:
        skeptic = next((r for r in roles if r.id == skeptic_role_id), None)
        if skeptic is None:
            skeptic = get_role(skeptic_role_id)
            tr.notes.append(f"blinded_skeptic: role {skeptic_role_id} added (not in role list)")
        debaters = [r for r in roles if r.id != skeptic.id]
        if skeptic.id not in assignment:
            raise KeyError(f"no model assigned to the blinded skeptic {skeptic.id!r}")
    views = {r.id: role_view(packet, r) for r in debaters}
    vis = {r.id: visible_ids(packet, r) for r in debaters}
    turns: list[AgentTurn] = []
    flags: dict[str, list[str]] = {}

    def _verify(est: dict[str, AgentEstimate], round_no: int) -> None:
        if not evidence_verification:
            return
        flags.clear()
        for rid, e in est.items():
            f = verify_evidence(e, packet, allowed_ids=vis[rid])
            if f:
                flags[rid] = f
                tr.notes.append(f"verification_flags round={round_no} role={rid}: {';'.join(f)}")

    # round 0 — independent
    turns += await _gather_turns(
        [run_turn(pool.get(_model_for(assignment, r.id)), r, render_user_prompt(views[r.id], None, kind="estimate"), AgentEstimate, _model_for(assignment, r.id), effort, 0, "estimate", max_tokens=max_tokens) for r in debaters]
    )
    current = _latest_estimates(turns)
    _verify(current, 0)
    prev_iqr = iqr_log_points({r: e.point_estimate for r, e in current.items()}) if current else None
    rounds_run = 0
    for r_no in range(1, rounds + 1):
        if len(current) < 2:
            tr.notes.append(f"stopped_early_at={r_no - 1}: fewer than two valid estimates")
            break
        # critiques: every debater sees the anonymised table
        crit_turns = await _gather_turns(
            [
                run_turn(pool.get(_model_for(assignment, r.id)), r, render_user_prompt(views[r.id], committee_state(r_no, current, for_role=r.id, visible=vis[r.id]), kind="critique", verification_flags=flags.get(r.id)), CritiqueSet, _model_for(assignment, r.id), effort, r_no, "critique", max_tokens=max_tokens)
                for r in debaters
            ]
        )
        turns += crit_turns
        addressed: dict[str, list[dict[str, Any]]] = {r.id: [] for r in debaters}
        for ct in crit_turns:
            if ct.critiques is None:
                continue
            for c in ct.critiques.critiques:
                if c.target_role in addressed and c.target_role != ct.role:
                    addressed[c.target_role].append({"from": "committee_member", "agree": c.agree, "points": c.points[:6], "evidence_ids_cited": c.evidence_ids_cited, "suggested_value": c.suggested_value})
        # revisions
        rev_turns = await _gather_turns(
            [
                run_turn(pool.get(_model_for(assignment, r.id)), r, render_user_prompt(views[r.id], committee_state(r_no, current, for_role=r.id, critiques=addressed[r.id], visible=vis[r.id]), kind="revision", verification_flags=flags.get(r.id)), AgentEstimate, _model_for(assignment, r.id), effort, r_no, "revision", max_tokens=max_tokens)
                for r in debaters
            ]
        )
        turns += rev_turns
        rounds_run = r_no
        current = _latest_estimates(turns)
        _verify(current, r_no)
        cur_iqr = iqr_log_points({r: e.point_estimate for r, e in current.items()})
        if should_stop(prev_iqr, cur_iqr, iqr_tol) and r_no < rounds:
            tr.notes.append(f"stopped_early_at={r_no}: IQR(log points) {prev_iqr:.5f} -> {cur_iqr:.5f} (< {iqr_tol:.0%} change); pre-specified rounds={rounds}")
            prev_iqr = cur_iqr
            break
        prev_iqr = cur_iqr
    tr.rounds = rounds_run
    if flags:
        tr.notes.append("verification_flags_unresolved_after_final_round: " + ",".join(sorted(flags)))
    # blinded skeptic verdict (H): packet + final estimates only, never the critiques
    if skeptic is not None:
        sk_view = role_view(packet, skeptic)
        state = committee_state(rounds_run + 1, current)
        state.pop("critiques_addressed_to_you", None)
        sk_turn = await run_turn(pool.get(_model_for(assignment, skeptic.id)), skeptic, render_user_prompt(sk_view, state, kind="skeptic_verdict"), AgentEstimate, _model_for(assignment, skeptic.id), effort, rounds_run + 1, "estimate", max_tokens=max_tokens)
        turns.append(sk_turn)
        tr.notes.append("blinded_skeptic_verdict: appended as final estimate of role 'skeptic' (round %d), discussion not shown" % (rounds_run + 1))
    return _finalize(tr, turns, t0, weights)


# ---------------------------------------------------------------------------------------------------
# J
# ---------------------------------------------------------------------------------------------------
async def cost_matched_single(
    packet: EvidencePacket,
    k: int,
    assignment: dict[str, str] | str,
    pool: ProviderPool,
    *,
    arm: str = "J_cost_matched",
    committee: str = "all_anthropic",
    effort: str = "medium",
    max_tokens: int = 4096,
    weights: dict[str, Any] | None = None,
    role_id: str = "generalist",
) -> DebateTranscript:
    """k independent generalist samples (same prompt; vendor-side sampling provides the variation — the
    mock is deterministic and returns identical samples) + a 2-step self-refine chain (self-critique then
    refined estimate) seeded by the first sample. Final estimates: sample_01..sample_k and self_refine."""
    t0 = time.perf_counter()
    role = get_role(role_id)
    model = _model_for(assignment, role.id)
    tr = DebateTranscript(tx_id=packet.subject.tx_id, arm=arm, committee=committee, protocol="cost_matched_single", rounds=0, packet_hash=packet.content_hash(), turns=[], assignment={role.id: model})
    view = role_view(packet, role)
    prompt = render_user_prompt(view, None, kind="estimate")
    samples = await _gather_turns([run_turn(pool.get(model), role, prompt, AgentEstimate, model, effort, 0, "estimate", max_tokens=max_tokens, tag=f"{role.id}:sample{i + 1}") for i in range(k)])
    turns: list[AgentTurn] = []
    final: dict[str, AgentEstimate] = {}
    for i, t in enumerate(samples):
        t.role = f"sample_{i + 1:02d}"
        turns.append(t)
        if t.estimate is not None:
            final[t.role] = t.estimate
    seed_est = next((t.estimate for t in samples if t.estimate is not None), None)
    if seed_est is not None:
        state = committee_state(1, {role.id: seed_est}, for_role=role.id)
        crit = await run_turn(pool.get(model), role, render_user_prompt(view, state, kind="self_critique"), CritiqueSet, model, effort, 1, "critique", max_tokens=max_tokens, tag=f"{role.id}:self_critique")
        crit.role = "self_refine"
        turns.append(crit)
        own = []
        if crit.critiques is not None:
            own = [{"from": "self", "agree": c.agree, "points": c.points[:6], "evidence_ids_cited": c.evidence_ids_cited, "suggested_value": c.suggested_value} for c in crit.critiques.critiques]
        state2 = committee_state(2, {role.id: seed_est}, for_role=role.id, critiques=own)
        ref = await run_turn(pool.get(model), role, render_user_prompt(view, state2, kind="self_refine"), AgentEstimate, model, effort, 2, "revision", max_tokens=max_tokens, tag=f"{role.id}:self_refine")
        ref.role = "self_refine"
        turns.append(ref)
        if ref.estimate is not None:
            final["self_refine"] = ref.estimate
        tr.rounds = 2
    else:
        tr.notes.append("self_refine skipped: no valid sample to seed the chain")
    tr.notes.append(f"k={k} independent samples of role {role.id}; self_refine = 2-step chain")
    out = _finalize(tr, turns, t0, weights, final=final)
    # convenience: the two cost-matched comparators the paper reports
    if final:
        samp = {r: e.point_estimate for r, e in final.items() if r.startswith("sample_")}
        if samp:
            out.aggregated["samples_mean"] = aggregate_points(samp)["mean"]
            out.aggregated["samples_median"] = aggregate_points(samp)["median"]
        if "self_refine" in final:
            out.aggregated["self_refine"] = final["self_refine"].point_estimate
    return out


PROTOCOLS = {"independent": independent, "structured_debate": structured_debate, "cost_matched_single": cost_matched_single}