"""Helper script for Robot Framework decision persistence integration tests.

Each test function exercises a serialization round-trip path for
the Decision domain model (model_dump / model_validate, JSON) and
prints a deterministic tag on success.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any

# Ensure src is importable when run from workspace root
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from ulid import ULID

from cleveragents.domain.models.core.decision import (
    ArtifactRef,
    ContextSnapshot,
    Decision,
    DecisionType,
    ResourceRef,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_ALL_DECISION_TYPES: list[DecisionType] = list(DecisionType)


def _make_decision(**overrides: Any) -> Decision:
    """Build a Decision with sensible defaults, merged with *overrides*."""
    defaults: dict[str, Any] = {
        "plan_id": str(ULID()),
        "sequence_number": 0,
        "decision_type": DecisionType.PROMPT_DEFINITION,
        "question": "What approach should we take?",
        "chosen_option": "Build a REST API",
    }
    dt = overrides.get("decision_type", defaults["decision_type"])
    if dt != DecisionType.PROMPT_DEFINITION and "parent_decision_id" not in overrides:
        defaults["parent_decision_id"] = str(ULID())
    defaults.update(overrides)
    return Decision(**defaults)


def _roundtrip_dump(d: Decision) -> Decision:
    data = d.model_dump()
    return Decision.model_validate(data)


def _roundtrip_json(d: Decision) -> Decision:
    json_str = d.model_dump_json()
    raw = json.loads(json_str)
    return Decision.model_validate(raw)


# ---------------------------------------------------------------------------
# Test functions
# ---------------------------------------------------------------------------


def _test_roundtrip_root() -> None:
    """Root decision survives model_dump round-trip."""
    original = _make_decision()
    restored = _roundtrip_dump(original)
    assert restored.decision_id == original.decision_id
    assert restored.decision_type == DecisionType.PROMPT_DEFINITION
    assert restored.is_root
    assert restored.question == original.question
    assert restored.chosen_option == original.chosen_option
    print("decision-persist-roundtrip-root-ok")


def _test_roundtrip_child() -> None:
    """Child decision with all fields survives model_dump round-trip."""
    parent_id = str(ULID())
    original = _make_decision(
        parent_decision_id=parent_id,
        decision_type=DecisionType.STRATEGY_CHOICE,
        question="Which framework?",
        chosen_option="FastAPI",
        alternatives_considered=["Flask", "Django", "Starlette"],
        confidence_score=0.88,
        rationale="Async support and auto docs",
    )
    restored = _roundtrip_dump(original)
    assert restored.decision_id == original.decision_id
    assert restored.parent_decision_id == parent_id
    assert not restored.is_root
    assert len(restored.alternatives_considered) == 3
    assert restored.confidence_score == 0.88
    assert restored.rationale == "Async support and auto docs"
    print("decision-persist-roundtrip-child-ok")


def _test_snapshot() -> None:
    """Context snapshot with resources survives round-trip."""
    snap = ContextSnapshot(
        hot_context_hash="sha256:persist_snap",
        hot_context_ref="store://snapshots/persist_snap",
        relevant_resources=[
            ResourceRef(resource_id=str(ULID()), path="src/main.py"),
            ResourceRef(resource_id=str(ULID()), path="src/utils.py"),
        ],
        actor_state_ref="checkpoint://actor/persist",
    )
    original = _make_decision(
        decision_type=DecisionType.IMPLEMENTATION_CHOICE,
        context_snapshot=snap,
    )
    restored = _roundtrip_dump(original)
    assert (
        restored.context_snapshot.hot_context_hash
        == original.context_snapshot.hot_context_hash
    )
    assert len(restored.context_snapshot.relevant_resources) == 2
    assert restored.context_snapshot.actor_state_ref == "checkpoint://actor/persist"
    print("decision-persist-snapshot-ok")


def _test_json_roundtrip() -> None:
    """Decision survives JSON serialization round-trip."""
    original = _make_decision(
        decision_type=DecisionType.RESOURCE_SELECTION,
        artifacts_produced=[
            ArtifactRef(artifact_path="src/api.py", artifact_type="file"),
            ArtifactRef(artifact_path="patches/fix.patch", artifact_type="patch"),
        ],
    )
    restored = _roundtrip_json(original)
    assert restored.decision_id == original.decision_id
    assert restored.decision_type == DecisionType.RESOURCE_SELECTION
    assert len(restored.artifacts_produced) == 2
    print("decision-persist-json-roundtrip-ok")


def _test_correction_chain() -> None:
    """Correction chain of 3 decisions persists and reconstructs."""
    plan_id = str(ULID())
    decisions: list[Decision] = []

    original = _make_decision(
        plan_id=plan_id,
        sequence_number=0,
        question="Initial approach?",
        chosen_option="Monolith",
    )
    decisions.append(original)

    prev_id = original.decision_id
    for i in range(1, 3):
        correction = _make_decision(
            plan_id=plan_id,
            sequence_number=i,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question=f"Correction {i}?",
            chosen_option=f"Approach v{i + 1}",
            is_correction=True,
            corrects_decision_id=prev_id,
            correction_reason=f"Reason {i}",
        )
        decisions[-1] = decisions[-1].with_superseded_by(correction.decision_id)
        decisions.append(correction)
        prev_id = correction.decision_id

    restored = [_roundtrip_dump(d) for d in decisions]
    for orig, rest in zip(decisions, restored, strict=True):
        assert rest.decision_id == orig.decision_id
        assert rest.is_correction == orig.is_correction
        assert rest.corrects_decision_id == orig.corrects_decision_id
        assert rest.superseded_by == orig.superseded_by

    for idx in range(1, len(restored)):
        assert restored[idx].corrects_decision_id == restored[idx - 1].decision_id
        assert restored[idx - 1].superseded_by == restored[idx].decision_id

    print("decision-persist-correction-chain-ok")


def _test_tree_reconstruction() -> None:
    """3-level tree (7 nodes) reconstructs from serialized data."""
    plan_id = str(ULID())
    nodes: list[Decision] = []
    seq = 0

    root = _make_decision(
        plan_id=plan_id,
        sequence_number=seq,
        question="Root?",
        chosen_option="Start",
    )
    nodes.append(root)
    seq += 1

    for ci in range(2):
        child = _make_decision(
            plan_id=plan_id,
            parent_decision_id=root.decision_id,
            sequence_number=seq,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question=f"Child {ci}?",
            chosen_option=f"Option C{ci}",
        )
        nodes.append(child)
        seq += 1
        for gi in range(2):
            grandchild = _make_decision(
                plan_id=plan_id,
                parent_decision_id=child.decision_id,
                sequence_number=seq,
                decision_type=DecisionType.IMPLEMENTATION_CHOICE,
                question=f"Grandchild {ci}.{gi}?",
                chosen_option=f"Impl {ci}.{gi}",
            )
            nodes.append(grandchild)
            seq += 1

    restored = [_roundtrip_dump(n) for n in nodes]
    assert len(restored) == 7
    id_set = {n.decision_id for n in restored}
    root_count = 0
    for n in restored:
        if n.parent_decision_id is None:
            root_count += 1
        else:
            assert n.parent_decision_id in id_set
    assert root_count == 1
    print("decision-persist-tree-reconstruction-ok")


def _test_all_types() -> None:
    """All decision types survive model_dump round-trip."""
    plan_id = str(ULID())
    for i, dt in enumerate(_ALL_DECISION_TYPES):
        kwargs: dict[str, Any] = {
            "plan_id": plan_id,
            "sequence_number": i,
            "decision_type": dt,
            "question": f"Q for {dt.value}?",
            "chosen_option": f"A for {dt.value}",
        }
        if dt != DecisionType.PROMPT_DEFINITION:
            kwargs["parent_decision_id"] = str(ULID())
        original = Decision(**kwargs)
        restored = _roundtrip_dump(original)
        assert restored.decision_id == original.decision_id
        assert restored.decision_type == original.decision_type
    print("decision-persist-all-types-ok")


# ---------------------------------------------------------------------------
# Dispatcher
# ---------------------------------------------------------------------------

_TESTS = {
    "roundtrip_root": _test_roundtrip_root,
    "roundtrip_child": _test_roundtrip_child,
    "snapshot": _test_snapshot,
    "json_roundtrip": _test_json_roundtrip,
    "correction_chain": _test_correction_chain,
    "tree_reconstruction": _test_tree_reconstruction,
    "all_types": _test_all_types,
}

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <test_name>")
        print(f"Available tests: {', '.join(sorted(_TESTS))}")
        sys.exit(1)

    test_name = sys.argv[1]
    if test_name not in _TESTS:
        print(f"Unknown test: {test_name}")
        print(f"Available: {', '.join(sorted(_TESTS))}")
        sys.exit(1)

    _TESTS[test_name]()
