import json

import pyarrow as pa
import pyarrow.parquet as pq

from albedo_eval_service.shared.submit_protocol import CANONICAL_MARKER
from local_train.pack import (
    candidate_turns,
    family_of,
    is_gold_submit,
    pack,
    phase_for,
)
from local_train.think import wrap_completion


def test_phase_and_family_helpers():
    assert phase_for(1, 9) == "cold"
    assert phase_for(7, 9) == "pre_edit"
    assert phase_for(9, 9) == "at_edit"
    assert phase_for(10, 9) == "post_edit"
    assert family_of("repo.abc.pr_12") == "pr"
    assert family_of("repo.abc.lm_9") == "lm"
    assert family_of("repo.abc.combine_x") == "combine"
    assert family_of("repo.abc.func_pm") == "mechanical"


def test_candidate_turns_skip_early_submit_and_include_edit():
    golds = [
        "```bash\nls\n```",
        "```bash\ncat a.py\n```",
        f"```bash\necho {CANONICAL_MARKER}\n```",
        "```bash\nsed -i 's/a/b/' a.py\n```",
        f"```bash\necho {CANONICAL_MARKER}\n```",
    ]
    turns = candidate_turns(len(golds), first_edit=4, golds=golds)
    idxs = {t for t, _ in turns}
    assert 1 in idxs and 2 in idxs
    assert 3 in idxs  # first-edit turn (1-based 4)
    assert 4 in idxs  # later submit kept
    # Early submit gold is not added as its own extra (turn 2 is only the official cold cut).
    assert is_gold_submit(golds[4])
    assert not is_gold_submit(golds[1])


def test_pack_from_tiny_parquet(tmp_path):
    source = tmp_path / "mini-coder" / "data"
    source.mkdir(parents=True)
    native = (
        "You are a coding agent.\n\n## Submission\n\n"
        f"```bash\necho {CANONICAL_MARKER}\n```\n"
    )
    messages = [
        {"role": "system", "content": native},
        {"role": "user", "content": "Fix foo.py"},
        {"role": "assistant", "content": "THOUGHT: look\n```bash\nls\n```"},
        {"role": "user", "content": "<returncode>0</returncode>\n<output>\nok\n</output>"},
        {"role": "assistant", "content": "THOUGHT: read\n```bash\ncat foo.py\n```"},
        {"role": "user", "content": "<returncode>0</returncode>\n<output>\nx=1\n</output>"},
        {
            "role": "assistant",
            "content": "THOUGHT: edit\n```bash\nsed -i 's/x=1/x=2/' foo.py\n```",
        },
        {"role": "user", "content": "<returncode>0</returncode>\n<output>\n\n</output>"},
        {
            "role": "assistant",
            "content": f"THOUGHT: submit\n```bash\necho {CANONICAL_MARKER}\n```",
        },
    ]
    table = pa.table(
        {
            "instance_id": ["demo.repo.pr_1"],
            "messages": [messages],
            "first_edit": [3],
            "family": ["pr"],
            "repo": ["demo"],
            "language": ["python"],
        }
    )
    pq.write_table(table, source / "train-00000.parquet")
    out = tmp_path / "pack"
    path = pack(
        dataset_root=tmp_path,
        out_dir=out,
        max_examples=16,
        seed="unit",
        n_salts=2,
    )
    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    assert rows
    kinds = {row["kind"] for row in rows}
    assert "edit" in kinds or "explore" in kinds or "submit" in kinds
    for row in rows:
        assert row["completion"].count("```bash") == 1
        assert "<think>" in row["completion"] and "</think>" in row["completion"]
        assert row["prompt"]
        if row["kind"] == "submit":
            assert row["submit_command"]
            body = row["completion"].split("</think>", 1)[-1]
            assert row["submit_command"] in body
            wrapped = wrap_completion(row["completion"])
            assert wrapped is not None
    meta = json.loads((out / "sft-16-unit.meta.json").read_text())
    assert meta["n"] == len(rows)


def test_expand_submit_salts_covers_multiple_rewrites(tmp_path):
    source = tmp_path / "mini-coder" / "data"
    source.mkdir(parents=True)
    native = (
        "You are a coding agent.\n\n## Submission\n\n"
        f"```bash\necho {CANONICAL_MARKER}\n```\n"
    )
    messages = [
        {"role": "system", "content": native},
        {"role": "user", "content": "Fix foo.py"},
        {"role": "assistant", "content": "THOUGHT: look\n```bash\nls\n```"},
        {"role": "user", "content": "<returncode>0</returncode>\n<output>\nok\n</output>"},
        {"role": "assistant", "content": "THOUGHT: read\n```bash\ncat foo.py\n```"},
        {"role": "user", "content": "<returncode>0</returncode>\n<output>\nx=1\n</output>"},
        {
            "role": "assistant",
            "content": "THOUGHT: edit\n```bash\nsed -i 's/x=1/x=2/' foo.py\n```",
        },
        {"role": "user", "content": "<returncode>0</returncode>\n<output>\n\n</output>"},
        {
            "role": "assistant",
            "content": f"THOUGHT: submit\n```bash\necho {CANONICAL_MARKER}\n```",
        },
    ]
    table = pa.table(
        {
            "instance_id": ["demo.repo.pr_1"],
            "messages": [messages],
            "first_edit": [3],
            "family": ["pr"],
            "repo": ["demo"],
            "language": ["python"],
        }
    )
    pq.write_table(table, source / "train-00000.parquet")
    path = pack(
        dataset_root=tmp_path,
        out_dir=tmp_path / "pack",
        max_examples=8,
        seed="expand",
        n_salts=4,
        submit_frac=0.8,
        edit_frac=0.2,
        expand_submit_salts=True,
    )
    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    submits = [row for row in rows if row["kind"] == "submit"]
    assert len(submits) >= 2
    assert len({row["submit_command"] for row in submits}) >= 2