from __future__ import annotations

import pytest

from dovla_cil.data.schema import (
    CIL_VERSION,
    ActionChunk,
    CILGroup,
    CILRecord,
    FailureInfo,
    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")