"""Tests for model visitor functions."""

import astropy.units as u
import pytest
from astropy.time import Time
from astropydantic import AstroPydanticQuantity, AstroPydanticTime

from ctao_datamodel import ModelBase, unflatten_model_instance, walk_model
from ctao_datamodel._visitor import (
    Relation,
    classes_in_model,
    flatten_model_instance,
    related_classes_in_model,
)


class Child1(ModelBase):
    v1: int = 1
    v2: int = 6


class Child2(ModelBase):
    v3: int = 7
    v4: int | None = 10
    c: Child1 = Child1()
    some_list: list[float] = [1.0, 2.0, 3.0]
    some_opt_list: list[float] | None = [1.0, 3.0, 5.0]


class Parent(ModelBase):
    x: Child1 | Child2 = Child1()
    y: Child2 = Child2()


class AstroFeatures(ModelBase):
    t: AstroPydanticTime = Time("2025-11-28 14:21:59.248384000")
    q: AstroPydanticQuantity = 10 * u.m / u.s
    opt_t: AstroPydanticTime | None = Time("2025-11-28 14:21:59.248384000")


@pytest.fixture(scope="session")
def simple_example_model():
    return Parent


@pytest.fixture(scope="session")
def astro_example_model():
    return AstroFeatures


def test_walk_model(simple_example_model):
    elements = {
        ".".join(x.item_name for x in p): (
            len(p),
            m.__name__,
            p[-1].is_optional,
            p[-1].parent_relation,
        )
        for m, p in walk_model(simple_example_model)
    }

    expected = {
        "parent": (1, "Parent", False, Relation.none),
        "parent.x": (2, "Child2", False, Relation.contains_one_of),
        "parent.x.v1": (3, "int", False, Relation.contains),
        "parent.x.v2": (3, "int", False, Relation.contains),
        "parent.x.v3": (3, "int", False, Relation.contains),
        "parent.x.v4": (3, "int", True, Relation.contains),
        "parent.x.c": (3, "Child1", False, Relation.contains),
        "parent.x.c.v1": (4, "int", False, Relation.contains),
        "parent.x.c.v2": (4, "int", False, Relation.contains),
        "parent.x.some_list": (3, "float", False, Relation.contains_many),
        "parent.x.some_opt_list": (3, "float", True, Relation.contains_many),
        "parent.y": (2, "Child2", False, Relation.contains),
        "parent.y.v3": (3, "int", False, Relation.contains),
        "parent.y.v4": (3, "int", True, Relation.contains),
        "parent.y.c": (3, "Child1", False, Relation.contains),
        "parent.y.c.v1": (4, "int", False, Relation.contains),
        "parent.y.c.v2": (4, "int", False, Relation.contains),
        "parent.y.some_list": (3, "float", False, Relation.contains_many),
        "parent.y.some_opt_list": (3, "float", True, Relation.contains_many),
    }

    assert len(elements) == len(expected)

    for key in elements:
        assert elements[key] == expected[key]


def test_walk_model_with_astro():
    elements = {
        ".".join(x.item_name for x in p): (
            len(p),
            m.__name__,
            p[-1].is_optional,
            p[-1].parent_relation,
        )
        for m, p in walk_model(AstroFeatures)
    }

    expected = {
        "astrofeatures": (1, "AstroFeatures", False, Relation.none),
        "astrofeatures.t": (2, "Time", False, Relation.contains),
        "astrofeatures.q": (2, "AstroPydanticQuantity", False, Relation.contains),
        "astrofeatures.opt_t": (2, "Time", True, Relation.contains),
    }

    for key in elements:
        assert elements[key] == expected[key]

    assert len(elements) == len(expected)


def test_flatten_model_instance(simple_example_model):
    """Flatten the model."""
    instance = simple_example_model()
    flat = flatten_model_instance(instance)

    expected = {
        "x.v1": 1,
        "x.v2": 6,
        "y.v3": 7,
        "y.v4": 10,
        "y.c.v1": 1,
        "y.c.v2": 6,
        "y.some_list.0": 1.0,
        "y.some_list.1": 2.0,
        "y.some_list.2": 3.0,
        "y.some_opt_list.0": 1.0,
        "y.some_opt_list.1": 3.0,
        "y.some_opt_list.2": 5.0,
    }

    assert flat == expected


def test_flatten_model_instance_with_parent(simple_example_model):
    """Try with a parent_key."""
    instance = simple_example_model()
    flat_parent = flatten_model_instance(instance, parent_key="bork")

    expected_parent = {
        "bork.x.v1": 1,
        "bork.x.v2": 6,
        "bork.y.v3": 7,
        "bork.y.v4": 10,
        "bork.y.c.v1": 1,
        "bork.y.c.v2": 6,
        "bork.y.some_list.0": 1.0,
        "bork.y.some_list.1": 2.0,
        "bork.y.some_list.2": 3.0,
        "bork.y.some_opt_list.0": 1.0,
        "bork.y.some_opt_list.1": 3.0,
        "bork.y.some_opt_list.2": 5.0,
    }

    assert flat_parent == expected_parent


def test_flatten_model_instance_no_string(simple_example_model):
    """Try with no stringification."""

    instance = simple_example_model()
    flat_nostring = flatten_model_instance(instance, to_string=False)

    expected_nostring = {
        "x.v1": 1,
        "x.v2": 6,
        "y.v3": 7,
        "y.v4": 10,
        "y.c.v1": 1,
        "y.c.v2": 6,
        "y.some_list.0": 1.0,
        "y.some_list.1": 2.0,
        "y.some_list.2": 3.0,
        "y.some_opt_list.0": 1.0,
        "y.some_opt_list.1": 3.0,
        "y.some_opt_list.2": 5.0,
    }

    assert flat_nostring == expected_nostring


def test_unflatten_no_string_no_expand_simple(simple_example_model):
    """Try unflattening a model, flattened with to_string=False and expand_lists=False."""
    instance = simple_example_model()
    flat = flatten_model_instance(
        simple_example_model(), to_string=False, expand_lists=False
    )
    new_instance = unflatten_model_instance(flat, simple_example_model)

    assert new_instance == instance


def test_unflatten_no_string_astro(astro_example_model):
    """Try unflattening a model, flattened with to_string=False."""
    model = astro_example_model
    instance = model()
    flat = flatten_model_instance(model(), to_string=False)
    new_instance = unflatten_model_instance(flat, model)
    assert new_instance == instance


def test_unflatten_string_expand_simple(simple_example_model):
    """Try unflattening a model, flattened with to_string=True."""
    model = simple_example_model
    instance = model()
    flat = flatten_model_instance(model(), to_string=True)
    new_instance = unflatten_model_instance(flat, model)

    assert new_instance == instance


def test_unflatten_no_string_expand_simple(simple_example_model):
    """Try unflattening a model, flattened with to_string=False and expand_lists=False."""
    instance = simple_example_model()
    flat = flatten_model_instance(
        simple_example_model(), to_string=False, expand_lists=True
    )
    new_instance = unflatten_model_instance(flat, simple_example_model)

    assert new_instance == instance


def test_unflatten_string_astro(astro_example_model):
    """Try unflattening a model, flattened with to_string=True."""
    model = astro_example_model
    instance = model()
    flat = flatten_model_instance(model(), to_string=True)
    new_instance = unflatten_model_instance(flat, model)

    assert new_instance == instance


def test_classes_in_model(simple_example_model):
    """Check we get the expected class list."""

    classes = classes_in_model(simple_example_model)

    assert set(classes) == {Parent, Child1, Child2}


def test_related_classes_in_model():
    class SomeParent(ModelBase):
        pass

    class SomeChild(SomeParent):
        pass

    class SomeGrandChild(SomeChild):
        pass

    class SomeGreatGrandChild(SomeGrandChild):
        pass

    parents, children = related_classes_in_model(SomeChild)

    assert len(parents) == 1
    assert len(children) == 2
