from __future__ import annotations

import pytest

from dovla_cil.data.schema import (
    CIL_VERSION,
    ActionChunk,
    CILBenchBranch,
    CILBenchGroup,
    CILGroup,
    CILRecord,
    FailureInfo,
    OutcomeVector,
    RewardInfo,
    StructuredEffect,
    compute_regret_and_ranks,
    compute_state_hash,
    make_record_id,
    validate_group,
)


def make_record(
    group_id: str, action_id: str, progress: float, *, state_hash: str = "s"
) -> CILRecord:
    return CILRecord(
        version=CIL_VERSION,
        record_id=make_record_id(group_id, action_id, seed=7),
        group_id=group_id,
        state_hash=state_hash,
        task_id="task",
        scene_id=None,
        instruction="move the mug",
        instruction_family={"family": "place"},
        observation_ref=None,
        observation_inline={"symbolic": True},
        action_chunk=ActionChunk(
            action_id=action_id,
            representation="delta_xy",
            horizon=1,
            values=[[progress, 0.0]],
            skill_type="push",
        ),
        next_observation_ref=None,
        next_observation_inline={"symbolic": True, "next": True},
        structured_effect=StructuredEffect(
            object_pose_delta={"mug": [progress, 0.0, 0.0]},
            relation_before={"inside(mug,bowl)": False},
            relation_after={"inside(mug,bowl)": progress > 0.5},
            moved_objects=["mug"] if progress else [],
            symbolic_before={"objects": {"mug": {"position": [0, 0, 0]}}},
            symbolic_after={"objects": {"mug": {"position": [progress, 0, 0]}}},
        ),
        reward=RewardInfo(
            progress=progress,
            success=progress > 0.5,
            terminal_success=progress > 0.5,
            dense_components={"progress": progress},
        ),
        regret=None,
        rank_within_group=None,
        candidate_type="sampled",
        failure=None
        if progress > 0
        else FailureInfo(type="no_motion", symbolic_reason="object did not move"),
    )


def test_schema_roundtrip() -> None:
    record = make_record("g", "a", 1.0)
    record.validate()
    assert CILRecord.from_dict(record.to_dict()) == record
    group = CILGroup.from_records([record])
    assert group.group_id == "g"


def test_group_validation() -> None:
    records = [make_record("g", "a", 0.0), make_record("g", "b", 1.0)]
    validate_group(records)
    with pytest.raises(ValueError):
        validate_group([make_record("g", "a", 0.0), make_record("other", "b", 1.0)])
    with pytest.raises(ValueError):
        validate_group([make_record("g", "a", 0.0), make_record("g", "b", 1.0, state_hash="x")])


def test_regret_and_rank_computation() -> None:
    ranked = compute_regret_and_ranks(
        [make_record("g", "low", 0.0), make_record("g", "high", 1.0)]
    )
    by_action = {record.action_chunk.action_id: record for record in ranked}
    assert by_action["high"].rank_within_group == 0
    assert by_action["high"].regret == 0.0
    assert by_action["low"].rank_within_group == 1
    assert by_action["low"].regret == 2.0


def test_deterministic_record_ids_and_state_hashes() -> None:
    assert make_record_id("g", "a", 1) == make_record_id("g", "a", 1)
    assert make_record_id("g", "a", 1) != make_record_id("g", "a", 2)
    assert compute_state_hash(b"state") == compute_state_hash(b"state")
    assert compute_state_hash(b"state") != compute_state_hash(b"other")


def test_outcome_vector_from_reward_and_utility() -> None:
    reward = RewardInfo(
        progress=0.6,
        success=True,
        terminal_success=True,
        dense_components={
            "contact_quality": 0.5,
            "task_stage_quality": 0.25,
            "smoothness": 0.8,
            "recovery": 1.0,
        },
    )

    outcome = OutcomeVector.from_reward(reward)

    assert outcome.success == 1.0
    assert outcome.progress == 0.6
    assert outcome.contact_quality == 0.5
    assert outcome.safety_violation == 0.0
    assert outcome.lexicographic_utility() > reward.score


def test_cilbench_group_roundtrip_and_same_state_contrast() -> None:
    base_action = ActionChunk(action_id="base", horizon=1, values=[[0.0, 0.0]])
    repair_action = ActionChunk(action_id="repair", horizon=1, values=[[0.1, 0.0]])
    group = CILBenchGroup(
        group_id="chart-0",
        task_id="PickCube-v1",
        split_id="train",
        simulator_state_hash="hash",
        instruction="pick the cube",
        observation_ref=None,
        observation_inline={"rgb": "obs/000.png"},
        scene_metadata={"target_object": "cube"},
        anchor_policy="h16_bc",
        branches=[
            CILBenchBranch(
                branch_id="base",
                action=base_action,
                branch_family="anchor",
                outcome=OutcomeVector(success=0.0, progress=0.2),
            ),
            CILBenchBranch(
                branch_id="repair",
                action=repair_action,
                branch_family="recovery_tangent",
                outcome=OutcomeVector(success=1.0, progress=0.8, recovery=1.0),
            ),
        ],
    )

    restored = CILBenchGroup.from_dict(group.to_dict())

    assert restored == group
    assert restored.same_state_causal_contrast("repair", "base") > 0.0
    with pytest.raises(KeyError):
        restored.same_state_causal_contrast("repair", "missing")