#!/usr/bin/env python3
"""
Programmatic verbatim candidate extractor for Multi-Harness Agentic Python-Kernel + In-Voice Persona SFT & FTPO/DPO.
Includes all v4 audit fixes (A1..A5 and B1..B8):
  - A1: Zero synthetic "inspect the workspace..." user messages; keeps real user & <mail kind="task"> instructions,
        and prepends preceding (assistant, tool) pair with loss_mask=0 when slicing mid-session.
  - A2: Strict check_lowercase_clean (including banning capital I/I'm/I'll) + META_THINK_RE + GENZ_SLANG_RE + 0 kaomoji in <think>.
  - A4: loss_mask=0 context turns keep reasoning_content="" unless already in-voice clean (zero fake placeholder strings).
  - B1..B8: Exempts `:` before ``` code fences, strips kaomoji backticks before grounding, broadens SECRET_RE & HONEST_SLOP_RE,
            allows collaborative `let's` in visible replies while banning it in <think>, and removes lcm_*/transcript_* from PREIMPORTS.
"""
from __future__ import annotations

import argparse
import ast
import hashlib
import json
import re
import sqlite3
import struct
import sys
from collections import Counter, defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
from mine_llm_isms_from_transcripts import decode_etf_term  # noqa: E402
from prepare_phase3_sft_data import apply_speech_act_surgery  # noqa: E402

ALBEDO_TOOLS_PATH = Path("/tmp/albedo_real_tools.json")
ALBEDO_TOOLS = json.loads(ALBEDO_TOOLS_PATH.read_text(encoding="utf-8")) if ALBEDO_TOOLS_PATH.exists() else []
VALID_ALBEDO_TOOLS = {t["function"]["name"] for t in ALBEDO_TOOLS}

PRIME_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "ipython",
            "description": "Run Python code in the persistent IPython kernel for this session.",
            "parameters": {
                "type": "object",
                "properties": {"code": {"type": "string"}},
                "required": ["code"],
            },
        },
    }
]

APPEND_SYS_PATH = Path.home() / ".agents" / "APPEND_SYSTEM.md"
APPEND_SYS_TEXT = APPEND_SYS_PATH.read_text(encoding="utf-8").strip() if APPEND_SYS_PATH.exists() else ""
PRIME_RECONSTRUCTED_SYSTEM = (
    "You are an autonomous coding agent operating inside prime-agent with a persistent `ipython` tool.\n\n"
    + APPEND_SYS_TEXT
)

OOD_REPOS = (
    "/Users/dawn/proj/hydrant",
    "/Users/dawn/proj/jpshow",
    "/Users/dawn/proj/colab-mcp",
    "/Users/dawn/ark",
    "/Users/regent/Developer/niri",
)

EXCLUDED_SESSIONS = {
    "050ac5c7c1e552db24ef477295e88acb",
}

# Broadened SECRET_RE (B7)
SECRET_RE = re.compile(
    r"(?:creds\.json|access_token|refresh_token|client_secret|PRIVATE KEY"
    r"|ghp_[A-Za-z0-9]{20,}|sk-(?:ant-|proj-)?[A-Za-z0-9_\-]{18,}|ya29\.[A-Za-z0-9_\-]{20,}"
    r"|tskey-[A-Za-z0-9_\-]{15,}|AKIA[0-9A-Z]{16}|xox[bp]-[A-Za-z0-9\-]{15,}"
    r"|hf_[A-Za-z0-9]{20,}|AIza[0-9A-Za-z_\-]{30,}|Bearer\s+[A-Za-z0-9._\-]{20,}|\.key\b)",
    re.I,
)

EMOJI_RE = re.compile(r"[\U0001F300-\U0001FAFF\u2705\u274C\u2728]")

# Broadened HONEST_SLOP_RE (B4)
HONEST_SLOP_RE = re.compile(
    r"\b(?:"
    r"honest(?:ly|y)?"
    r"|to\s+be\s+(?:honest|blunt|frank|candid|transparent|straight|upfront|direct|fair|clear)"
    r"|in\s+all\s+(?:honesty|fairness|seriousness)"
    r"|being\s+(?:real|blunt|frank|candid|transparent|straight|upfront|truthful|honest)"
    r"|frankly|candidly|transparently|admittedly"
    r"|full\s+transparency|i\s+want\s+to\s+be\s+upfront"
    r"|own\s+(?:my|that|this|the)\s+mistake"
    r"|my\s+(?:mistake|bad|apologies)"
    r"|(?:that|this)(?:\s+one)?(?:'s|\s+is)\s+on\s+me"
    r"|i\s+(?:was\s+wrong|messed\s+up|screwed\s+up|apologize)"
    r"|apologies|sorry\s+about"
    r"|(?:good|great|nice|fair|sharp)\s+(?:catch|point|question|call|eye)"
    r"|you\s+caught\s+(?:me|that)"
    r"|you(?:'re|\s+are)\s+(?:completely\s+|totally\s+|absolutely\s+|100%\s+)?(?:right|correct|spot\s+on)"
    r"|spot\s+on"
    r"|(?:short|honest|real|blunt|frank|candid|straight)\s+(?:answer|take|framing|version|truth|read|summary|assessment)"
    r"|real\s+talk"
    r")\b",
    re.I,
)

# The 7 lexical cliches (slop4.h k.lex)
LEXICAL_SLOP_RE = re.compile(
    r"\b(?:"
    r"rather\s+than"
    r"|right\s+away"
    r"|end[\s\-]+to[\s\-]+end"
    r"|short\s+version"
    r"|exactly\s+what"
    r"|here(?:'s|\s+is)\s+what"
    r"|what(?:'s|\s+is)\s+actually"
    r")\b",
    re.I,
)

# Gen-Z slang gate (B5, per APPEND_SYSTEM.md)
GENZ_SLANG_RE = re.compile(
    r"\b(?:rizz|no\s+cap|poggers|pog|ngl|fr\s+fr|lowkey|highkey|bussin|slay|skibidi|gyatt)\b",
    re.I,
)

# Third-person rubric / meta-summaries in <think> (A2 + B6: bans let's in <think>, allows let's in replies)
RUBRIC_THINK_RE = re.compile(
    r"(?:\bthe\s+(?:user|instruction|task|objective|command|prompt|request)\b"
    r"|\buser\s+(?:wants|is\s+asking|asked)\b"
    r"|\bi\s+should\b|\blet\s+me\b|\blet's\b|\bi\s+will\s+now\b|\bi'll\s+(?:now|start)\b|\bmy\s+focus\b)",
    re.I,
)

RUBRIC_REPLY_RE = re.compile(
    r"(?:\bthe\s+user\b|\buser\s+(?:wants|is\s+asking|asked)\b|\bi\s+should\b|\blet\s+me\b|\bi\s+will\s+now\b|\bi'll\s+now\b)",
    re.I,
)

KAOMOJI_IN_THINK_RE = re.compile(r"(?:\([^\n()]{1,12}(?:[ωᗜ‿ᵕ•´｀¬^]|>_<)[^\n()]{0,12}\)|:[3cP]\b|<3)")

ALLOWED_ACRONYMS = {
    "AST", "API", "BPE", "CDP", "CLI", "CPU", "CSS", "CSV", "CUA", "DB", "DPO", "DOM", "EOF",
    "ETF", "FIFO", "FTS", "FTPO", "GDN", "GGUF", "GPU", "GRPO", "HTML", "HTTP", "HTTPS", "ID", "IDS",
    "IO", "IPC", "JSON", "JSONL", "JSX", "JWT", "KLD", "KV", "LCM", "LLM", "LORA", "LRU",
    "MAC", "MCP", "MD", "MIB", "GIB", "KIB", "MLP", "MMPROJ", "MOCE", "MROPE", "MSL", "MTP",
    "NIX", "NL", "NUL", "NVME", "OOD", "OOM", "OS", "PATH", "PDF", "PID", "PLE", "PNG", "POSIX",
    "PPL", "PR", "PTY", "QAT", "QSA", "RAM", "RCO", "REPL", "RGB", "RL", "RMS", "ROPE", "RPC",
    "RSS", "SCR", "SFT", "SHA", "SIMD", "SQL", "SQLITE", "SRO", "SSD", "SSH", "SSL", "STDERR",
    "STDIN", "STDOUT", "SVG", "TCP", "TG", "TLS", "TOK", "TOML", "TOR", "TS", "TSX", "TUI",
    "UI", "URI", "URL", "UTF", "UUID", "VM", "VRAM", "WFLD", "XML", "XRPC", "YAML",
}

# PREIMPORTS: only true Python kernel names (B8: lcm_* and transcript_* removed!)
PREIMPORTS = {
    "asyncio", "base64", "collections", "datetime", "functools", "hashlib", "itertools",
    "json", "math", "os", "Path", "re", "shutil", "sys", "textwrap", "time",
    "files", "run", "parallel", "cells", "output", "work", "mail", "agents", "vent", "vents",
    "resolve_vent", "memory", "remote", "colab", "browser", "cua", "view_diff", "view_code",
    "show_image", "display", "web_search", "skills", "commands", "mcp", "jobs",
} | set(dir(__import__("builtins")))

IDIOM_PATTERNS = (
    "files.read(", "files.edit(", "files.write(", "files.find(", "files.paths(", "files.ls(",
    "parallel(", ".pipe(", "view_diff(", "view_code(", "display(", "cells.run(", "output.read(",
    "work.", "mail.submit(", "agents.", "remote.connect(", "browser.", "cua.", "memory.", "vents(",
)

ANTI_PATTERNS_RE = re.compile(
    r"(?:=\s*await\s+run\(|parallel\([^\n]*await\s+run\(|\bsubprocess\.|\bos\.system\(|\.read_text\(|\.write_text\()"
)

SYSTEM_STORE: dict[str, str] = {}


def register_sys(prompt: str) -> str:
    sid = "sys_" + hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:12]
    SYSTEM_STORE[sid] = prompt
    return sid


def strip_fences(txt: str) -> str:
    return re.sub(r"```[\s\S]*?```", "", txt)


def strip_kaomoji_backticks(txt: str) -> str:
    """Strips kaomoji that contain backticks like (`・ω・´) or (´• ω •`) before backtick span extraction (B2)."""
    return re.sub(r"\([^()\n]{0,12}`[^()\n]{0,12}\)", "()", txt)


def check_lowercase_clean(txt: str) -> str | None:
    """Strict lowercase check including banning single-letter capital I (B3)."""
    outside = strip_fences(strip_kaomoji_backticks(txt))
    outside = re.sub(r"`[^`\n]+`", "", outside)
    outside = re.sub(r"https?://\S+", "", outside)
    outside = re.sub(r"(?:/[\w.\-]+)+", "", outside)
    for word in re.findall(r"[A-Za-z]+", outside):
        if word == "I":
            return "capital 'I' outside backticks (must be lowercase 'i')"
        if word.isupper() and len(word) >= 2 and word not in ALLOWED_ACRONYMS:
            return f"uppercase word '{word}' outside backticks"
        if word[0].isupper() and not word.isupper():
            return f"capitalized word '{word}' outside backticks (must be lowercase)"
    return None


def has_structural_markdown_slop(txt: str) -> bool:
    """Checks slop4.h k.md + k.colon rules outside fenced code blocks, exempting `:` right before a ``` fence (B1)."""
    lines = txt.splitlines()
    in_fence = False
    for i, ln in enumerate(lines):
        s = ln.strip()
        if s.startswith("```"):
            in_fence = not in_fence
            continue
        if in_fence or not s:
            continue
        if re.match(r"^(?:[-*+•–—→])\s+", s):
            return True
        if re.match(r"^\d+[.)]\s+", s):
            return True
        if re.match(r"^(?:#{1,6}\s+|---+\s*$|\*\*\*+\s*$)", s):
            return True
        if re.match(r"^\*\*[^*]+\*\*\s*:", s):
            return True
        if s.endswith((":", ":**", "**:")) and len(s.split()) <= 6:
            # Find next non-empty line; exempt if it opens a ``` fence (B1)
            nxt = next((lines[j].strip() for j in range(i + 1, len(lines)) if lines[j].strip()), None)
            if nxt is not None and not nxt.startswith("```"):
                return True
    return False


def has_any_prose_slop(txt: str, is_think: bool = False) -> bool:
    if not txt:
        return False
    outside = strip_fences(txt)
    if (
        HONEST_SLOP_RE.search(outside)
        or LEXICAL_SLOP_RE.search(outside)
        or GENZ_SLANG_RE.search(outside)
        or EMOJI_RE.search(outside)
    ):
        return True
    if is_think and (RUBRIC_THINK_RE.search(outside) or KAOMOJI_IN_THINK_RE.search(outside)):
        return True
    if not is_think and RUBRIC_REPLY_RE.search(outside):
        return True
    if has_structural_markdown_slop(txt):
        return True
    if check_lowercase_clean(txt) is not None:
        return True
    return False


def is_clean_in_voice_think(th: str) -> bool:
    if not th or len(th) < 25 or len(th) > 480:
        return False
    return not has_any_prose_slop(th, is_think=True)


def check_backtick_grounding(reply_txt: str, context_txt: str) -> bool:
    outside = strip_fences(strip_kaomoji_backticks(reply_txt))
    for bt in re.findall(r"`([^`\n]+)`", outside):
        s = bt.strip()
        if not s or len(s) <= 2 or s in PREIMPORTS:
            continue
        if s not in context_txt:
            return False
    return True


def assign_split(session_id: str, cwd: str | None) -> str:
    if cwd and any(cwd.startswith(r) for r in OOD_REPOS):
        return "val_ood_repo"
    h = int(hashlib.sha256(session_id.encode("utf-8")).hexdigest()[:8], 16)
    return "val_id" if (h % 100) < 15 else "train"


def reconstruct_system_prompt(ki: str, kc: bytes | None, keep_persona: bool = True) -> str:
    texts = []
    if kc and len(kc) >= 6 and kc[5] == 108:
        n_items = struct.unpack(">I", kc[6:10])[0]
        pos = 10
        for _ in range(n_items):
            blen = struct.unpack(">I", kc[pos + 1 : pos + 5])[0]
            inner = kc[pos + 5 : pos + 5 + blen]
            pos += 5 + blen
            idx = inner.find(b"user")
            if idx != -1 and inner[idx + 4] == 109:
                m_pos = idx + 4
                tlen = struct.unpack(">I", inner[m_pos + 1 : m_pos + 5])[0]
                t_str = inner[m_pos + 5 : m_pos + 5 + tlen].decode("utf-8", "replace")
                if not keep_persona and "chronically online nerd girl" in t_str:
                    continue
                texts.append(t_str)
    return ki + ("\n\n" + "\n\n".join(texts) if texts else "")


def get_defined_and_undef(code_str: str, prior_defined: set[str]) -> tuple[set[str], list[str]] | None:
    try:
        tree = ast.parse(code_str, mode="exec")
    except SyntaxError:
        return None
    defined = set(prior_defined) | PREIMPORTS
    new_defs = set()
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            for alias in node.names:
                nm = (alias.asname or alias.name).split(".")[0]
                defined.add(nm)
                new_defs.add(nm)
        elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
            defined.add(node.id)
            new_defs.add(node.id)
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            defined.add(node.name)
            new_defs.add(node.name)
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                for a in node.args.args + node.args.kwonlyargs:
                    defined.add(a.arg)
                if node.args.vararg:
                    defined.add(node.args.vararg.arg)
                if node.args.kwarg:
                    defined.add(node.args.kwarg.arg)
        elif isinstance(node, ast.ExceptHandler) and node.name:
            defined.add(node.name)
        elif isinstance(node, ast.arg):
            defined.add(node.arg)
    undef = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load):
            if node.id not in defined:
                undef.append(node.id)
    return new_defs, undef


def is_clean_idiomatic_code(code: str) -> bool:
    if not code or len(code) < 20 or len(code) > 2600:
        return False
    if SECRET_RE.search(code) or ANTI_PATTERNS_RE.search(code):
        return False
    if "open(" in code and not any(x in code for x in ("sqlite3", "os.open", "rb", "wb")):
        return False
    return any(p in code for p in IDIOM_PATTERNS) or ("run(" in code and "await " in code)


def is_failed_tool_output(out_str: str) -> bool:
    try:
        obj = json.loads(out_str)
    except Exception:
        return False
    st = obj.get("status", "")
    if st in ("error", "interrupted", "backgrounded"):
        return True
    out_txt = obj.get("output", "")
    return any(k in out_txt for k in ("Traceback (most recent call last)", "FAILED", "SyntaxError:", "ValueError:", "WorkError:"))


def clean_loss0_think(th: str | None) -> str:
    """A4 fix: never emit teacher-cased summaries or fake constant strings in loss_mask=0 context turns."""
    if th and is_clean_in_voice_think(th):
        return th
    return ""


def extract_from_sqlite(db_path: Path, source_tag: str) -> list[dict]:
    con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
    cell_map: dict[str, str] = {}
    for cid, src in con.execute("SELECT id, source FROM cells WHERE source IS NOT NULL"):
        cell_map[cid] = src
    print(f"[{source_tag}] Loaded {len(cell_map)} cells from {db_path}")

    sessions = con.execute(
        """
        SELECT id, model, cwd, known_instructions, known_context
        FROM sessions
        WHERE known_instructions IS NOT NULL AND transcript_count >= 4
        ORDER BY created_at DESC
        """
    ).fetchall()

    candidates = []
    for sid, model, cwd, ki, kc in sessions:
        if sid in EXCLUDED_SESSIONS:
            continue
        sys_prompt = reconstruct_system_prompt(ki, kc, keep_persona=True)
        sys_prompt_no_persona = reconstruct_system_prompt(ki, kc, keep_persona=False)
        if SECRET_RE.search(sys_prompt):
            continue
        if source_tag == "astaroth" and "chronically online nerd girl" not in sys_prompt and APPEND_SYS_TEXT:
            sys_prompt = sys_prompt + "\n\n" + APPEND_SYS_TEXT
        has_persona = "chronically online nerd girl" in sys_prompt
        sys_id = register_sys(sys_prompt)
        sys_id_no_persona = register_sys(sys_prompt_no_persona)
        split = assign_split(sid, cwd)

        t_rows = con.execute(
            "SELECT seq, row_class, payload FROM transcript WHERE session = ? ORDER BY seq ASC",
            (sid,),
        ).fetchall()

        parsed_msgs = []
        for sq, rc, pay in t_rows:
            try:
                dec = decode_etf_term(pay, 1)[0]
            except Exception:
                continue
            if not (isinstance(dec, tuple) and len(dec) == 2 and isinstance(dec[1], tuple)):
                continue
            tag_tuple = dec[1]
            kind = tag_tuple[0]
            if kind == ("atom", "user") and isinstance(tag_tuple[1], str):
                utxt = tag_tuple[1]
                if SECRET_RE.search(utxt):
                    continue
                parsed_msgs.append({"seq": sq, "role": "user", "content": utxt})
            elif kind == ("atom", "tool_output") and len(tag_tuple) >= 3:
                call_id = tag_tuple[1]
                out_str = tag_tuple[2]
                if not isinstance(out_str, str) or SECRET_RE.search(out_str) or len(out_str) > 12000:
                    parsed_msgs.append({"seq": sq, "role": "tool_skip", "tool_call_id": call_id})
                else:
                    parsed_msgs.append({"seq": sq, "role": "tool", "tool_call_id": call_id, "content": out_str})
            elif kind == ("atom", "replay") and len(tag_tuple) >= 2:
                ri = tag_tuple[1]
                if not (isinstance(ri, tuple) and len(ri) >= 3 and ri[0] == ("atom", "replay_item")):
                    continue
                proto = ri[1]
                body = ri[2]
                if proto == ("atom", "chat_completions") and isinstance(body, dict):
                    c = body.get("content")
                    c_str = c if isinstance(c, str) and c.strip() else None
                    rc_val = body.get("reasoning_content")
                    rc_str = rc_val.strip() if isinstance(rc_val, str) and rc_val.strip() else None
                    tcs = body.get("tool_calls")
                    norm_tcs = []
                    valid_tcs = True
                    if isinstance(tcs, list) and tcs:
                        for tc in tcs:
                            if not isinstance(tc, dict) or "function" not in tc:
                                valid_tcs = False
                                break
                            fn = tc["function"]
                            fname = fn.get("name")
                            raw_args = fn.get("arguments")
                            try:
                                args_dict = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
                            except Exception:
                                valid_tcs = False
                                break
                            if not isinstance(args_dict, dict):
                                valid_tcs = False
                                break
                            tc_id = tc.get("id") or ""
                            if fname == "python":
                                full_cid = f"{sid}/{tc_id}"
                                if full_cid not in cell_map or args_dict.get("code") != cell_map[full_cid]:
                                    valid_tcs = False
                                    break
                            norm_tcs.append(
                                {
                                    "id": tc_id,
                                    "type": "function",
                                    "function": {"name": fname, "arguments": args_dict},
                                }
                            )
                    if not valid_tcs:
                        continue
                    if c_str and SECRET_RE.search(c_str):
                        continue
                    parsed_msgs.append(
                        {
                            "seq": sq,
                            "role": "assistant",
                            "reasoning_content": rc_str,
                            "content": c_str,
                            "tool_calls": norm_tcs if norm_tcs else None,
                        }
                    )
                elif proto == ("atom", "responses") and isinstance(body, dict):
                    btype = body.get("type")
                    if btype == "function_call":
                        fname = body.get("name")
                        tc_id = body.get("call_id") or ""
                        raw_args = body.get("arguments")
                        try:
                            args_dict = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
                        except Exception:
                            continue
                        if not isinstance(args_dict, dict):
                            continue
                        if fname == "python":
                            full_cid = f"{sid}/{tc_id}"
                            if full_cid not in cell_map or args_dict.get("code") != cell_map[full_cid]:
                                continue
                        parsed_msgs.append(
                            {
                                "seq": sq,
                                "role": "assistant",
                                "reasoning_content": None,
                                "content": None,
                                "tool_calls": [
                                    {
                                        "id": tc_id,
                                        "type": "function",
                                        "function": {"name": fname, "arguments": args_dict},
                                    }
                                ],
                            }
                        )
                    elif body.get("role") == "assistant" and isinstance(body.get("content"), list):
                        txts = [
                            p.get("text", "")
                            for p in body["content"]
                            if isinstance(p, dict) and p.get("type") == "output_text" and p.get("text")
                        ]
                        c_str = "\n".join(txts).strip()
                        if c_str and not SECRET_RE.search(c_str):
                            parsed_msgs.append(
                                {
                                    "seq": sq,
                                    "role": "assistant",
                                    "reasoning_content": None,
                                    "content": c_str,
                                    "tool_calls": None,
                                }
                            )

        # A1 fix: NEVER use a synthetic placeholder string! Keep real user messages (including <mail ...> task briefs up to 4500 chars)
        last_user_txt: str | None = None
        per_session_windows = defaultdict(int)
        used_indices: set[int] = set()

        for i, m in enumerate(parsed_msgs):
            if m["role"] == "user":
                ut = m["content"].strip()
                # Skip only empty/system-notice wakeups or huge >4500-char blobs; keep <mail ...> task instructions!
                if not ut.startswith("<system-notice>") and 6 <= len(ut) <= 4500:
                    last_user_txt = ut
                continue
            if last_user_txt is None or i in used_indices or m["role"] != "assistant":
                continue

            # Helper to prepend preceding (assistant, tool) pair as loss_mask=0 context when i is not immediately after user (A1)
            prefix_ctx = []
            prev_defs: set[str] = set()
            if (
                i >= 2
                and parsed_msgs[i - 1]["role"] == "tool"
                and parsed_msgs[i - 2]["role"] == "assistant"
                and parsed_msgs[i - 2].get("tool_calls")
                and len(parsed_msgs[i - 2]["tool_calls"]) == 1
            ):
                prev_tc = parsed_msgs[i - 2]["tool_calls"][0]
                if prev_tc["function"]["name"] == "python":
                    r0 = get_defined_and_undef(prev_tc["function"]["arguments"].get("code", ""), set())
                    if r0:
                        prev_defs = r0[0]
                prefix_ctx = [
                    {
                        "role": "assistant",
                        "loss_mask": 0,
                        "reasoning_content": clean_loss0_think(parsed_msgs[i - 2].get("reasoning_content")),
                        "needs_think_rewrite": False,
                        "content": None,
                        "tool_calls": parsed_msgs[i - 2]["tool_calls"],
                    },
                    {
                        "role": "tool",
                        "tool_call_id": parsed_msgs[i - 1]["tool_call_id"],
                        "content": parsed_msgs[i - 1]["content"],
                    },
                ]

            # Bucket 3: Top-level tool discrimination
            if m.get("tool_calls") and all(
                tc["function"]["name"] in (VALID_ALBEDO_TOOLS - {"python"}) for tc in m["tool_calls"]
            ):
                if i + 1 < len(parsed_msgs) and parsed_msgs[i + 1]["role"] == "tool":
                    if parsed_msgs[i - 1]["role"] != "user" and not prefix_ctx:
                        continue
                    win_msgs = (
                        [{"role": "system", "system_id": sys_id}, {"role": "user", "content": last_user_txt}]
                        + prefix_ctx
                        + [
                            {
                                "role": "assistant",
                                "loss_mask": 1,
                                "reasoning_content": m["reasoning_content"] or "",
                                "needs_think_rewrite": not is_clean_in_voice_think(m["reasoning_content"] or ""),
                                "content": None,
                                "tool_calls": m["tool_calls"],
                            }
                        ]
                    )
                    j = i + 1
                    while j < len(parsed_msgs) and parsed_msgs[j]["role"] == "tool":
                        win_msgs.append({"role": "tool", "tool_call_id": parsed_msgs[j]["tool_call_id"], "content": parsed_msgs[j]["content"]})
                        j += 1
                    candidates.append(
                        {
                            "id": f"b3_{source_tag}_{sid[:8]}_{m['seq']}",
                            "bucket": "top_level_tool_discrimination",
                            "split": split,
                            "session_id": sid,
                            "model": model,
                            "cwd": cwd,
                            "harness": "albedo",
                            "user_provenance": "verbatim",
                            "tools": ALBEDO_TOOLS,
                            "messages": win_msgs,
                        }
                    )
                    used_indices.add(i)
                    continue

            # Bucket 2: Error recovery (turn i is loss_mask=0 failed cell -> turn i+2 is loss_mask=1 recovery cell)
            if (
                per_session_windows["error_recovery"] < 5
                and m.get("tool_calls")
                and len(m["tool_calls"]) == 1
                and m["tool_calls"][0]["function"]["name"] == "python"
                and i + 3 < len(parsed_msgs)
                and parsed_msgs[i + 1]["role"] == "tool"
                and parsed_msgs[i + 2]["role"] == "assistant"
                and parsed_msgs[i + 2].get("tool_calls")
                and len(parsed_msgs[i + 2]["tool_calls"]) == 1
                and parsed_msgs[i + 2]["tool_calls"][0]["function"]["name"] == "python"
                and parsed_msgs[i + 3]["role"] == "tool"
            ):
                out1 = parsed_msgs[i + 1]["content"]
                out2 = parsed_msgs[i + 3]["content"]
                src1 = m["tool_calls"][0]["function"]["arguments"]["code"]
                src2 = parsed_msgs[i + 2]["tool_calls"][0]["function"]["arguments"]["code"]
                if is_failed_tool_output(out1) and not is_failed_tool_output(out2) and is_clean_idiomatic_code(src2):
                    res1 = get_defined_and_undef(src1, prev_defs)
                    defs1 = prev_defs | (res1[0] if res1 else set())
                    res2 = get_defined_and_undef(src2, defs1)
                    if res2 and not res2[1]:
                        th2 = parsed_msgs[i + 2]["content"] if (sid.startswith("01a13b1a") and parsed_msgs[i + 2].get("content") and not parsed_msgs[i + 2].get("reasoning_content")) else (parsed_msgs[i + 2]["reasoning_content"] or "")
                        win_msgs = (
                            [{"role": "system", "system_id": sys_id}, {"role": "user", "content": last_user_txt}]
                            + prefix_ctx
                            + [
                                {
                                    "role": "assistant",
                                    "loss_mask": 0,
                                    "reasoning_content": clean_loss0_think(m.get("reasoning_content")),
                                    "needs_think_rewrite": False,
                                    "content": None,
                                    "tool_calls": m["tool_calls"],
                                },
                                {"role": "tool", "tool_call_id": parsed_msgs[i + 1]["tool_call_id"], "content": out1},
                                {
                                    "role": "assistant",
                                    "loss_mask": 1,
                                    "reasoning_content": th2,
                                    "needs_think_rewrite": not is_clean_in_voice_think(th2),
                                    "content": None,
                                    "tool_calls": parsed_msgs[i + 2]["tool_calls"],
                                },
                                {"role": "tool", "tool_call_id": parsed_msgs[i + 3]["tool_call_id"], "content": out2},
                            ]
                        )
                        candidates.append(
                            {
                                "id": f"b2_{source_tag}_{sid[:8]}_{m['seq']}",
                                "bucket": "error_recovery",
                                "split": split,
                                "session_id": sid,
                                "model": model,
                                "cwd": cwd,
                                "harness": "albedo",
                                "user_provenance": "verbatim",
                                "tools": ALBEDO_TOOLS,
                                "messages": win_msgs,
                            }
                        )
                        per_session_windows["error_recovery"] += 1
                        used_indices.update({i, i + 2})
                        continue

            # Bucket 1 & Bucket 5: 1..3 consecutive idiomatic python cells (+ optional grounded reply)
            is_anchor = (int(hashlib.md5(f"{sid}_{m['seq']}".encode()).hexdigest()[:4], 16) % 100) < 18
            b_key = "no_persona_coding_anchor" if (not has_persona or is_anchor) else "albedo_tool_mastery"
            cap = 8 if sid.startswith("01a13b1a") else 4
            if (
                per_session_windows[b_key] < cap
                and model != "swift-1.5-k512-iq2xs"
                and m.get("tool_calls")
                and len(m["tool_calls"]) == 1
                and m["tool_calls"][0]["function"]["name"] == "python"
                and i + 1 < len(parsed_msgs)
                and parsed_msgs[i + 1]["role"] == "tool"
                and not is_failed_tool_output(parsed_msgs[i + 1]["content"])
            ):
                # Require either immediately following user OR having prefix_ctx so turn i is causally grounded (A1)!
                if parsed_msgs[i - 1]["role"] != "user" and not prefix_ctx:
                    continue
                src1 = m["tool_calls"][0]["function"]["arguments"]["code"]
                if is_clean_idiomatic_code(src1):
                    res1 = get_defined_and_undef(src1, prev_defs)
                    if res1 and not res1[1]:
                        defs = prev_defs | set(res1[0])
                        th1 = m["content"] if (sid.startswith("01a13b1a") and m.get("content") and not m.get("reasoning_content")) else (m["reasoning_content"] or "")
                        use_sys_id = sys_id_no_persona if b_key == "no_persona_coding_anchor" else sys_id
                        win_msgs = (
                            [{"role": "system", "system_id": use_sys_id}, {"role": "user", "content": last_user_txt}]
                            + prefix_ctx
                            + [
                                {
                                    "role": "assistant",
                                    "loss_mask": 1,
                                    "reasoning_content": th1,
                                    "needs_think_rewrite": not is_clean_in_voice_think(th1),
                                    "content": None,
                                    "tool_calls": m["tool_calls"],
                                },
                                {"role": "tool", "tool_call_id": parsed_msgs[i + 1]["tool_call_id"], "content": parsed_msgs[i + 1]["content"]},
                            ]
                        )
                        used_indices.add(i)
                        cur_idx = i + 2
                        steps = 1
                        while (
                            steps < 3
                            and cur_idx + 1 < len(parsed_msgs)
                            and parsed_msgs[cur_idx]["role"] == "assistant"
                            and parsed_msgs[cur_idx].get("tool_calls")
                            and len(parsed_msgs[cur_idx]["tool_calls"]) == 1
                            and parsed_msgs[cur_idx]["tool_calls"][0]["function"]["name"] == "python"
                            and parsed_msgs[cur_idx + 1]["role"] == "tool"
                            and not is_failed_tool_output(parsed_msgs[cur_idx + 1]["content"])
                        ):
                            src_k = parsed_msgs[cur_idx]["tool_calls"][0]["function"]["arguments"]["code"]
                            if not is_clean_idiomatic_code(src_k):
                                break
                            res_k = get_defined_and_undef(src_k, defs)
                            if not res_k or res_k[1]:
                                break
                            defs |= res_k[0]
                            th_k = parsed_msgs[cur_idx]["content"] if (sid.startswith("01a13b1a") and parsed_msgs[cur_idx].get("content") and not parsed_msgs[cur_idx].get("reasoning_content")) else (parsed_msgs[cur_idx]["reasoning_content"] or "")
                            win_msgs.append(
                                {
                                    "role": "assistant",
                                    "loss_mask": 1,
                                    "reasoning_content": th_k,
                                    "needs_think_rewrite": not is_clean_in_voice_think(th_k),
                                    "content": None,
                                    "tool_calls": parsed_msgs[cur_idx]["tool_calls"],
                                }
                            )
                            win_msgs.append(
                                {
                                    "role": "tool",
                                    "tool_call_id": parsed_msgs[cur_idx + 1]["tool_call_id"],
                                    "content": parsed_msgs[cur_idx + 1]["content"],
                                }
                            )
                            used_indices.add(cur_idx)
                            cur_idx += 2
                            steps += 1

                        if (
                            has_persona
                            and b_key != "no_persona_coding_anchor"
                            and per_session_windows["in_voice_replies"] < 4
                            and cur_idx < len(parsed_msgs)
                            and parsed_msgs[cur_idx]["role"] == "assistant"
                            and not parsed_msgs[cur_idx].get("tool_calls")
                            and parsed_msgs[cur_idx].get("content")
                        ):
                            surg_txt, _ = apply_speech_act_surgery(parsed_msgs[cur_idx]["content"])
                            ctx_blob = "\n".join(
                                x.get("content", "") or json.dumps(x.get("tool_calls", ""))
                                for x in win_msgs[1:]
                            )
                            if 30 <= len(surg_txt) <= 1500 and check_backtick_grounding(surg_txt, ctx_blob):
                                th_rep = parsed_msgs[cur_idx]["reasoning_content"] or ""
                                win_msgs.append(
                                    {
                                        "role": "assistant",
                                        "loss_mask": 1,
                                        "reasoning_content": th_rep,
                                        "needs_think_rewrite": not is_clean_in_voice_think(th_rep),
                                        "needs_reply_rewrite": has_any_prose_slop(surg_txt, is_think=False),
                                        "content": surg_txt,
                                    }
                                )
                                used_indices.add(cur_idx)
                                b_key = "in_voice_replies"

                        candidates.append(
                            {
                                "id": f"b1_{source_tag}_{b_key[:6]}_{sid[:8]}_{m['seq']}",
                                "bucket": b_key,
                                "split": split,
                                "session_id": sid,
                                "model": model,
                                "cwd": cwd,
                                "harness": "albedo",
                                "user_provenance": "verbatim",
                                "tools": ALBEDO_TOOLS,
                                "messages": win_msgs,
                            }
                        )
                        per_session_windows[b_key] += 1
                        continue

            # Bucket 4: Direct conversational / design reply OR tool_call (loss_mask=0) -> tool_out -> reply (loss_mask=1)
            if (
                has_persona
                and per_session_windows["in_voice_replies"] < 5
                and model != "swift-1.5-k512-iq2xs"
                and not m.get("tool_calls")
                and m.get("content")
                and i >= 1
            ):
                surg_txt, _ = apply_speech_act_surgery(m["content"])
                if 30 <= len(surg_txt) <= 1500 and not SECRET_RE.search(surg_txt):
                    if parsed_msgs[i - 1]["role"] == "user" and not parsed_msgs[i - 1]["content"].startswith("<system-notice>"):
                        u_content = parsed_msgs[i - 1]["content"][:4000]
                        if check_backtick_grounding(surg_txt, u_content):
                            win_msgs = [
                                {"role": "system", "system_id": sys_id},
                                {"role": "user", "content": u_content},
                                {
                                    "role": "assistant",
                                    "loss_mask": 1,
                                    "reasoning_content": m["reasoning_content"] or "",
                                    "needs_think_rewrite": not is_clean_in_voice_think(m["reasoning_content"] or ""),
                                    "needs_reply_rewrite": has_any_prose_slop(surg_txt, is_think=False),
                                    "content": surg_txt,
                                },
                            ]
                            candidates.append(
                                {
                                    "id": f"b4_{source_tag}_{sid[:8]}_{m['seq']}",
                                    "bucket": "in_voice_replies",
                                    "split": split,
                                    "session_id": sid,
                                    "model": model,
                                    "cwd": cwd,
                                    "harness": "albedo",
                                    "user_provenance": "verbatim",
                                    "tools": ALBEDO_TOOLS,
                                    "messages": win_msgs,
                                }
                            )
                            per_session_windows["in_voice_replies"] += 1
                            used_indices.add(i)
                    elif prefix_ctx:
                        ctx_blob = last_user_txt + "\n" + json.dumps(prefix_ctx[0]["tool_calls"]) + "\n" + prefix_ctx[1]["content"]
                        if check_backtick_grounding(surg_txt, ctx_blob):
                            win_msgs = (
                                [{"role": "system", "system_id": sys_id}, {"role": "user", "content": last_user_txt}]
                                + prefix_ctx
                                + [
                                    {
                                        "role": "assistant",
                                        "loss_mask": 1,
                                        "reasoning_content": m["reasoning_content"] or "",
                                        "needs_think_rewrite": not is_clean_in_voice_think(m["reasoning_content"] or ""),
                                        "needs_reply_rewrite": has_any_prose_slop(surg_txt, is_think=False),
                                        "content": surg_txt,
                                    }
                                ]
                            )
                            candidates.append(
                                {
                                    "id": f"b4_t_{source_tag}_{sid[:8]}_{m['seq']}",
                                    "bucket": "in_voice_replies",
                                    "split": split,
                                    "session_id": sid,
                                    "model": model,
                                    "cwd": cwd,
                                    "harness": "albedo",
                                    "user_provenance": "verbatim",
                                    "tools": ALBEDO_TOOLS,
                                    "messages": win_msgs,
                                }
                            )
                            per_session_windows["in_voice_replies"] += 1
                            used_indices.add(i)

    con.close()
    return candidates


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", default="/Users/dawn/.albedo/albedo.sqlite")
    ap.add_argument("--tag", default="local")
    ap.add_argument("--out", default=str(ROOT / "data" / "agent_persona_v1" / "raw_candidates.jsonl"))
    ap.add_argument("--include-prime", action="store_true", default=True)
    args = ap.parse_args()

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

    candidates = extract_from_sqlite(Path(args.db), args.tag)

    prime_dir = Path.home() / ".prime" / "agent" / "sessions"
    prime_added = 0
    if args.include_prime and prime_dir.exists():
        for pfile in sorted(prime_dir.glob("*.jsonl")):
            if prime_added >= 180:
                break
            psid = pfile.stem
            split = assign_split(psid, None)
            try:
                lines = [json.loads(l) for l in pfile.open(errors="replace") if len(l) < 120000 and l.strip()]
            except Exception:
                continue
            per_p = 0
            for idx, obj in enumerate(lines):
                if per_p >= 3:
                    break
                if obj.get("type") != "message":
                    continue
                msg = obj.get("message", {})
                if msg.get("role") != "assistant" or "gemini" not in (msg.get("model") or "").lower():
                    continue
                parts = msg.get("content") or []
                if not isinstance(parts, list):
                    continue
                has_tc = any(isinstance(p, dict) and p.get("type") == "toolCall" for p in parts)
                txts = [p.get("text", "") for p in parts if isinstance(p, dict) and p.get("type") == "text" and p.get("text")]
                if has_tc or not txts:
                    continue
                raw_reply = "\n".join(txts).strip()
                surg_reply, _ = apply_speech_act_surgery(raw_reply)
                if not (35 <= len(surg_reply) <= 1200) or SECRET_RE.search(surg_reply) or EMOJI_RE.search(surg_reply):
                    continue
                u_txt = None
                for b_idx in range(idx - 1, max(-1, idx - 8), -1):
                    b_obj = lines[b_idx]
                    if b_obj.get("type") == "message" and b_obj.get("message", {}).get("role") == "user":
                        u_parts = b_obj["message"].get("content") or []
                        if isinstance(u_parts, list):
                            ut = "\n".join(p.get("text", "") for p in u_parts if isinstance(p, dict) and p.get("type") == "text").strip()
                            if 8 <= len(ut) <= 2000 and not SECRET_RE.search(ut):
                                u_txt = ut
                                break
                if not u_txt or not check_backtick_grounding(surg_reply, u_txt):
                    continue
                candidates.append(
                    {
                        "id": f"b4_prime_{psid[:8]}_{idx}",
                        "bucket": "in_voice_replies",
                        "split": split,
                        "session_id": psid,
                        "model": msg.get("model"),
                        "cwd": None,
                        "harness": "prime",
                        "user_provenance": "verbatim",
                        "tools": PRIME_TOOLS,
                        "messages": [
                            {"role": "system", "system_id": register_sys(PRIME_RECONSTRUCTED_SYSTEM)},
                            {"role": "user", "content": u_txt},
                            {
                                "role": "assistant",
                                "loss_mask": 1,
                                "reasoning_content": "",
                                "needs_think_rewrite": True,
                                "needs_reply_rewrite": has_any_prose_slop(surg_reply, is_think=False),
                                "content": surg_reply,
                            },
                        ],
                    }
                )
                per_p += 1
                prime_added += 1

    sys_path = out_path.with_name(out_path.stem + "_systems.json")
    sys_path.write_text(json.dumps(SYSTEM_STORE, ensure_ascii=False), encoding="utf-8")
    with out_path.open("w", encoding="utf-8") as fh:
        for c in candidates:
            fh.write(json.dumps(c, ensure_ascii=False) + "\n")

    bucket_counts = Counter(c["bucket"] for c in candidates)
    split_counts = Counter(c["split"] for c in candidates)
    total_sup_turns = sum(
        1 for c in candidates for m in c["messages"] if m.get("role") == "assistant" and m.get("loss_mask") == 1
    )
    needs_th = sum(
        1
        for c in candidates
        for m in c["messages"]
        if m.get("role") == "assistant" and m.get("loss_mask") == 1 and m.get("needs_think_rewrite")
    )
    needs_rep = sum(
        1
        for c in candidates
        for m in c["messages"]
        if m.get("role") == "assistant" and m.get("loss_mask") == 1 and m.get("needs_reply_rewrite")
    )
    print(f"Wrote {len(candidates)} candidate windows ({total_sup_turns} supervised turns) to {out_path}:")
    print("  By bucket:", dict(bucket_counts))
    print("  By split :", dict(split_counts))
    print(f"  Turns needing <think> rewrite: {needs_th} / {total_sup_turns}")
    print(f"  Turns needing reply slop/bullet/honest-ism rewrite: {needs_rep} / {total_sup_turns}")
    return 0


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