"""DEBATE → ChatML extraction for ARCHON SFT v2 family_multi_agent.

Schema source (Multi-Agent-LLMs/DEBATE):
- globalMemory: list of turn dicts with {agent_id, persona, message, contribution, turn, agreement}
- personas: list of {agentId, persona, personaDescription}
- instruction: task prompt
- input: problem context
- paradigm: debate|relay|report|memory
- agreements: vote records

Output ChatML format:
{
  "messages": [
    {"role": "system", "content": "<peer-equal team description>"},
    {"role": "user", "content": "<instruction + input>"},
    {"role": "assistant", "content": "<multi-turn transcript with persona prefixes>"}
  ],
  "task_type": "family_negotiate|family_consult",
  "niveau": "N2|N3",
  "peers_involved": ["Marine Biologist", "Food Safety Specialist", ...],
  "language": "en",
  "source_ds": "Multi-Agent-LLMs/DEBATE",
  "paradigm": "debate|relay|report|memory"
}

CPU-only. No GPU. Streams configs sequentially.

Usage:
    python debate_extract_v2.py --output debate_chatml.jsonl --rows-per-config 30
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path

try:
    from datasets import load_dataset, get_dataset_config_names
except ImportError:
    print("ERR: pip install datasets", file=sys.stderr)
    sys.exit(1)

REPO = "Multi-Agent-LLMs/DEBATE"

PARADIGM_TO_NIVEAU = {
    "debate": ("family_negotiate", "N3"),   # multi-round propose/critique/vote
    "relay":  ("family_consult", "N2"),     # sequential delegation, peer aggregates
    "report": ("family_consult", "N2"),     # independent reports + aggregation
    "memory": ("family_consult", "N2"),     # shared memory consultation
}

SYSTEM_TEMPLATE = (
    "Vous êtes membre d'une équipe d'experts collaborant sur un problème. "
    "Aucun maître, voix égales. Chaque expert contribue sa spécialité, "
    "puis l'équipe décide par consensus ou vote. Persona attribuée : {personas}."
)


def detect_paradigm(config_name: str) -> str | None:
    """Extract paradigm from config name (e.g. 'critical_expert_debate_approval_voting' -> 'debate')."""
    for paradigm in ("debate", "relay", "report", "memory"):
        if f"_{paradigm}_" in config_name:
            return paradigm
    return None


def extract_sample(sample: dict, paradigm: str, source_ds: str = REPO) -> dict | None:
    """Convert one DEBATE sample to ChatML multi-agent transcript."""
    gm = sample.get("globalMemory")
    if not isinstance(gm, list) or len(gm) < 3:
        return None

    personas_meta = sample.get("personas") or []
    personas_seen = []
    for p in personas_meta:
        if isinstance(p, dict):
            name = p.get("persona") or "Expert"
            if name not in personas_seen:
                personas_seen.append(name)
    if not personas_seen:
        # Fallback: extract from turns
        for turn in gm:
            name = (turn.get("persona") or "").strip()
            if name and name not in personas_seen:
                personas_seen.append(name)

    if len(personas_seen) < 2:
        return None

    # 'input' can be a string OR list of strings (paradigm-dependent)
    raw_instr = sample.get("instruction")
    raw_input = sample.get("input")
    instruction = ""
    if isinstance(raw_instr, str):
        instruction = raw_instr.strip()
    elif isinstance(raw_instr, list):
        instruction = "\n".join(str(x) for x in raw_instr if x).strip()
    input_str = ""
    if isinstance(raw_input, str):
        input_str = raw_input.strip()
    elif isinstance(raw_input, list):
        input_str = "\n".join(str(x) for x in raw_input if x).strip()
    if input_str and input_str not in instruction:
        user_content = f"{instruction}\n\n{input_str}".strip()
    else:
        user_content = instruction or input_str
    if len(user_content) < 20:
        return None

    # Build assistant transcript
    lines = []
    for turn in gm:
        if not isinstance(turn, dict):
            continue
        persona = (turn.get("persona") or "Expert").strip()
        msg = (turn.get("message") or "").strip()
        contrib = (turn.get("contribution") or "").strip()
        if not msg:
            continue
        # Skip extremely short or system-noise turns
        if len(msg) < 30:
            continue
        prefix = f"**{persona}**"
        if contrib and contrib not in ("draft", ""):
            prefix += f" *({contrib})*"
        lines.append(f"{prefix}: {msg}")

    if len(lines) < 3:
        return None

    assistant_content = "\n\n".join(lines)
    if len(assistant_content) > 12000:
        # Trim oversized transcripts (training seq cap)
        assistant_content = assistant_content[:11800] + "\n\n[...transcript tronqué pour longueur...]"

    task_type, niveau = PARADIGM_TO_NIVEAU.get(paradigm, ("family_negotiate", "N3"))

    system_content = SYSTEM_TEMPLATE.format(personas=", ".join(personas_seen))

    return {
        "messages": [
            {"role": "system", "content": system_content},
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": assistant_content},
        ],
        "task_type": task_type,
        "niveau": niveau,
        "peers_involved": personas_seen,
        "language": "en",
        "source_ds": source_ds,
        "paradigm": paradigm,
        "n_turns": len(lines),
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--output", required=True, help="Output JSONL path")
    ap.add_argument("--rows-per-config", type=int, default=30,
                    help="Max samples to extract per config (default 30, 145 configs * 30 = 4350 target)")
    ap.add_argument("--max-configs", type=int, default=145,
                    help="Max configs to process (default all)")
    ap.add_argument("--paradigms", nargs="+",
                    default=["debate", "relay", "report", "memory"],
                    help="Paradigms to extract")
    args = ap.parse_args()

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"Listing configs...")
    cfgs = get_dataset_config_names(REPO)
    cfgs = [c for c in cfgs if c != "_preview"]
    print(f"Total configs available: {len(cfgs)}")

    # Filter by paradigm
    selected = []
    for c in cfgs:
        p = detect_paradigm(c)
        if p and p in args.paradigms:
            selected.append((c, p))
    print(f"Configs matching paradigms {args.paradigms}: {len(selected)}")
    selected = selected[:args.max_configs]
    print(f"Processing {len(selected)} configs, {args.rows_per_config} rows each (target ~{len(selected) * args.rows_per_config})")

    total_written = 0
    total_failed = 0
    paradigm_counts = {}
    t0 = time.time()

    with open(out_path, "w", encoding="utf-8") as f:
        for ci, (config_name, paradigm) in enumerate(selected, 1):
            try:
                ds = load_dataset(REPO, config_name, split="train", streaming=True)
            except Exception as e:
                print(f"[{ci}/{len(selected)}] LOAD_FAIL {config_name}: {type(e).__name__}", file=sys.stderr)
                continue

            written_this_cfg = 0
            for i, sample in enumerate(ds):
                if i >= args.rows_per_config:
                    break
                try:
                    out = extract_sample(sample, paradigm)
                    if out is None:
                        total_failed += 1
                        continue
                    f.write(json.dumps(out, ensure_ascii=False) + "\n")
                    written_this_cfg += 1
                    total_written += 1
                    paradigm_counts[paradigm] = paradigm_counts.get(paradigm, 0) + 1
                except Exception as e:
                    total_failed += 1
                    if total_failed <= 2:
                        import traceback
                        print(f"  [extract_fail] {config_name}#{i}: {type(e).__name__}: {e}", file=sys.stderr)
                        traceback.print_exc(file=sys.stderr)
                        print(f"  sample keys: {list(sample.keys())}", file=sys.stderr)
                        print(f"  personas type: {type(sample.get('personas')).__name__}", file=sys.stderr)
                        print(f"  globalMemory type: {type(sample.get('globalMemory')).__name__}", file=sys.stderr)
                        gm = sample.get('globalMemory')
                        if isinstance(gm, list) and gm:
                            print(f"  gm[0] type: {type(gm[0]).__name__}", file=sys.stderr)
                            print(f"  gm[0] repr (200ch): {repr(gm[0])[:200]}", file=sys.stderr)

            elapsed = time.time() - t0
            rate = total_written / max(1, elapsed)
            if ci % 10 == 0 or ci == len(selected):
                print(f"[{ci}/{len(selected)}] {config_name[:40]:40s} written={written_this_cfg:2d} total={total_written:5d} rate={rate:.1f}/s elapsed={elapsed:.0f}s")

    print(f"\n=== DONE ===")
    print(f"Total written: {total_written}")
    print(f"Total failed/skipped: {total_failed}")
    print(f"By paradigm: {paradigm_counts}")
    print(f"Output: {out_path}")
    print(f"Elapsed: {time.time() - t0:.0f}s")
    return 0


if __name__ == "__main__":
    sys.exit(main())