"""Tests for abreka util subcommands."""

from __future__ import annotations

import json
import os
from pathlib import Path

from click.testing import CliRunner

from abreka.cli import main


def _make_session(path: Path, events: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(events, f, indent=2)
        f.write("\n")


def _delta(payload: str) -> dict:
    return {
        "type": "message_update",
        "assistantMessageEvent": {"type": "toolcall_delta", "delta": payload},
    }


def _bloated_events() -> list[dict]:
    return [
        {"type": "message_start", "message": {"role": "user"}},
        _delta("a"),
        _delta("ab"),
        _delta("abc"),
        {
            "type": "message_update",
            "assistantMessageEvent": {"type": "toolcall_end"},
        },
        {"type": "message_end", "message": {}},
    ]


def _clean_events() -> list[dict]:
    return [
        {"type": "message_start", "message": {"role": "user"}},
        {
            "type": "message_update",
            "assistantMessageEvent": {"type": "toolcall_end"},
        },
        {"type": "message_end", "message": {}},
    ]


def _snapshot_bloat_events() -> list[dict]:
    snapshot = {
        "role": "assistant",
        "content": [{"type": "thinking", "thinking": "plan"}],
    }
    return [
        {"type": "message_start", "message": {"role": "assistant"}},
        {
            "type": "message_update",
            "message": snapshot,
            "assistantMessageEvent": {
                "type": "thinking_delta",
                "delta": "x",
                "partial": snapshot,
            },
        },
        {"type": "message_end", "message": {"role": "assistant"}},
    ]


def _seed_experiment(abreka_project: Path, exp_id: str, events: list[dict]) -> Path:
    session = (
        abreka_project / "abreka" / "experiments" / exp_id / "artifacts" / "pi_session_implement.json"
    )
    _make_session(session, events)
    return session


class TestRepairDryRun:
    def test_dry_run_default_does_not_modify_files(self, abreka_project):
        session = _seed_experiment(abreka_project, "0001", _bloated_events())
        original_bytes = session.read_bytes()

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["util", "repair-session-files"])

        assert result.exit_code == 0, result.output
        assert "FIX" in result.output
        assert "DRY RUN" in result.output
        assert session.read_bytes() == original_bytes

    def test_dry_run_reports_no_work_cleanly(self, abreka_project):
        _seed_experiment(abreka_project, "0001", _clean_events())
        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["util", "repair-session-files"])
        assert result.exit_code == 0, result.output
        assert "Nothing to repair" in result.output


class TestRepairApply:
    def test_apply_with_yes_rewrites_file(self, abreka_project):
        session = _seed_experiment(abreka_project, "0001", _bloated_events())
        original_size = session.stat().st_size

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["util", "repair-session-files", "--apply", "--yes"])

        assert result.exit_code == 0, result.output
        assert "repaired" in result.output

        reloaded = json.loads(session.read_text())
        assert all(
            not (
                e.get("type") == "message_update"
                and isinstance(e.get("assistantMessageEvent"), dict)
                and e["assistantMessageEvent"].get("type") == "toolcall_delta"
            )
            for e in reloaded
        )
        assert len(reloaded) == 3
        assert reloaded[0]["type"] == "message_start"
        assert reloaded[1]["assistantMessageEvent"]["type"] == "toolcall_end"
        assert reloaded[2]["type"] == "message_end"
        assert session.stat().st_size < original_size

    def test_apply_aborts_when_confirmation_declined(self, abreka_project):
        session = _seed_experiment(abreka_project, "0001", _bloated_events())
        original_bytes = session.read_bytes()

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(
            main, ["util", "repair-session-files", "--apply"], input="n\n"
        )

        # click.confirm(..., abort=True) exits non-zero on decline.
        assert result.exit_code != 0
        assert session.read_bytes() == original_bytes

    def test_apply_skips_clean_files(self, abreka_project):
        session = _seed_experiment(abreka_project, "0001", _clean_events())
        original_bytes = session.read_bytes()
        original_mtime = session.stat().st_mtime_ns

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["util", "repair-session-files", "--apply", "--yes"])

        assert result.exit_code == 0, result.output
        assert session.read_bytes() == original_bytes
        assert session.stat().st_mtime_ns == original_mtime

    def test_apply_compacts_message_update_snapshots(self, abreka_project):
        session = _seed_experiment(abreka_project, "0001", _snapshot_bloat_events())
        original_size = session.stat().st_size

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["util", "repair-session-files", "--apply", "--yes"])

        assert result.exit_code == 0, result.output
        reloaded = json.loads(session.read_text())
        update = [e for e in reloaded if e.get("type") == "message_update"][0]
        assert "message" not in update
        assert "partial" not in update["assistantMessageEvent"]
        assert session.stat().st_size < original_size

    def test_keep_backup_creates_bak(self, abreka_project):
        session = _seed_experiment(abreka_project, "0001", _bloated_events())
        original_bytes = session.read_bytes()

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(
            main,
            ["util", "repair-session-files", "--apply", "--yes", "--keep-backup"],
        )

        assert result.exit_code == 0, result.output
        bak = session.with_suffix(session.suffix + ".bak")
        assert bak.exists()
        assert bak.read_bytes() == original_bytes
        reloaded = json.loads(session.read_text())
        assert all(
            not (
                e.get("type") == "message_update"
                and isinstance(e.get("assistantMessageEvent"), dict)
                and e["assistantMessageEvent"].get("type") == "toolcall_delta"
            )
            for e in reloaded
        )


class TestRepairPaths:
    def test_explicit_path_argument(self, abreka_project, tmp_path):
        # Seed a file in the project plus another one elsewhere; only the
        # explicit path should be touched.
        in_project = _seed_experiment(abreka_project, "0001", _bloated_events())
        other = tmp_path / "stray" / "pi_session.json"
        _make_session(other, _bloated_events())

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(
            main,
            ["util", "repair-session-files", "--apply", "--yes", str(other)],
        )

        assert result.exit_code == 0, result.output
        # The explicit one was repaired.
        reloaded = json.loads(other.read_text())
        assert all(
            not (
                e.get("type") == "message_update"
                and isinstance(e.get("assistantMessageEvent"), dict)
                and e["assistantMessageEvent"].get("type") == "toolcall_delta"
            )
            for e in reloaded
        )
        # The project one was NOT touched.
        in_project_loaded = json.loads(in_project.read_text())
        assert any(
            e.get("type") == "message_update"
            and isinstance(e.get("assistantMessageEvent"), dict)
            and e["assistantMessageEvent"].get("type") == "toolcall_delta"
            for e in in_project_loaded
        )


class TestRepairErrors:
    def test_corrupt_file_is_reported_not_crashed(self, abreka_project):
        # Create a corrupt "session" file in the expected location.
        bad = (
            abreka_project
            / "abreka"
            / "experiments"
            / "0001"
            / "artifacts"
            / "pi_session.json"
        )
        bad.parent.mkdir(parents=True, exist_ok=True)
        bad.write_text("this is not json\n")

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["util", "repair-session-files"])

        # Dry run over an errored file is exit 0 (no writes attempted).
        assert result.exit_code == 0, result.output
        assert "ERROR" in result.output
        # The corrupt file is untouched.
        assert bad.read_text() == "this is not json\n"
