from __future__ import annotations

import pytest

from dovla_cil.tasks.library import (
    HARD_CAUSALSTRESS_CATEGORIES,
    ToyTaskLibrary,
    built_in_causalstress_tasks,
    built_in_toy_tasks,
)
from dovla_cil.tasks.predicates import evaluate_predicate, evaluate_task_success
from dovla_cil.tasks.schema import RelationSpec, TaskSpec
from dovla_cil.tasks.validators import validate_task


def test_valid_builtin_tasks_pass() -> None:
    tasks = built_in_toy_tasks()
    assert len(tasks) == 10
    for task in tasks:
        validate_task(task)
        assert TaskSpec.from_dict(task.to_dict()) == task


def test_toy_task_library_cycles_tasks() -> None:
    library = ToyTaskLibrary()
    assert library.get(0).task_id == "toy_pick_red_mug"
    assert library.get(10).task_id == "toy_pick_red_mug"
    assert len(library.list()) == 10


def test_hard_causalstress_tasks_pass_validation() -> None:
    tasks = built_in_causalstress_tasks()
    assert len(tasks) == len(HARD_CAUSALSTRESS_CATEGORIES)
    assert {task.family for task in tasks} == set(HARD_CAUSALSTRESS_CATEGORIES)
    for task in tasks:
        validate_task(task)


def test_invalid_predicate_object_fails() -> None:
    task = built_in_toy_tasks()[0].model_copy(
        update={"success_predicates": [RelationSpec(name="grasped", args=["missing_mug"])]}
    )
    with pytest.raises(ValueError, match="unknown object"):
        validate_task(task)


def test_predicate_evaluator_inside_left_of_opened_lifted() -> None:
    state = {
        "objects": {
            "red_mug": {"inside": "blue_bowl", "position": [0.0, 0.0, 0.2]},
            "blue_bowl": {"position": [1.0, 0.0, 0.0]},
            "drawer": {"opened": True},
            "can": {"position": [0.0, 0.0, 0.3]},
        }
    }
    assert evaluate_predicate(RelationSpec(name="inside", args=["red_mug", "blue_bowl"]), state)
    assert evaluate_predicate(RelationSpec(name="left_of", args=["red_mug", "blue_bowl"]), state)
    assert evaluate_predicate(RelationSpec(name="opened", args=["drawer"]), state)
    assert evaluate_predicate(RelationSpec(name="lifted", args=["can"]), state)


def test_evaluate_task_success() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    state = {
        "objects": {
            "red_mug": {"inside": "blue_bowl", "position": [0.0, 0.0, 0.0]},
            "blue_bowl": {"position": [0.0, 0.0, 0.0]},
        }
    }
    assert evaluate_task_success(task, state)