"""Tests for Pedigree models."""

from __future__ import annotations

from ai4gh.specs.pedigree.models import (
    AffectedStatus,
    KINRelationship,
    LifeStatus,
    Pedigree,
    PedigreeIndividual,
    PedigreeRelationship,
    Sex,
)


class TestSex:
    def test_enum_values(self) -> None:
        assert Sex.MALE == "MALE"
        assert Sex.FEMALE == "FEMALE"
        assert Sex.OTHER == "OTHER_SEX"
        assert Sex.UNKNOWN == "UNKNOWN_SEX"


class TestAffectedStatus:
    def test_enum_values(self) -> None:
        assert AffectedStatus.AFFECTED == "AFFECTED"
        assert AffectedStatus.UNAFFECTED == "UNAFFECTED"
        assert AffectedStatus.MISSING == "MISSING"


class TestLifeStatus:
    def test_enum_values(self) -> None:
        assert LifeStatus.ALIVE == "ALIVE"
        assert LifeStatus.DECEASED == "DECEASED"


class TestKINRelationship:
    def test_enum_values(self) -> None:
        assert KINRelationship.PARENT == "KIN:001"
        assert KINRelationship.MOTHER == "KIN:002"
        assert KINRelationship.FATHER == "KIN:003"
        assert KINRelationship.SIBLING == "KIN:004"
        assert KINRelationship.TWIN == "KIN:005"
        assert KINRelationship.CHILD == "KIN:008"


class TestPedigreeIndividual:
    def test_minimal(self) -> None:
        ind = PedigreeIndividual(individual_id="IND-001")
        assert ind.individual_id == "IND-001"
        assert ind.sex is None
        assert ind.maternal_id is None

    def test_full(self) -> None:
        ind = PedigreeIndividual(
            individual_id="IND-001",
            sex=Sex.FEMALE,
            affected_status=AffectedStatus.AFFECTED,
            life_status=LifeStatus.ALIVE,
            maternal_id="IND-002",
            paternal_id="IND-003",
        )
        assert ind.sex is Sex.FEMALE
        assert ind.affected_status is AffectedStatus.AFFECTED
        assert ind.maternal_id == "IND-002"


class TestPedigreeRelationship:
    def test_creation(self) -> None:
        rel = PedigreeRelationship(
            subject_id="IND-001",
            object_id="IND-002",
            relationship=KINRelationship.MOTHER,
        )
        assert rel.subject_id == "IND-001"
        assert rel.relationship is KINRelationship.MOTHER


TRIO_PEDIGREE = Pedigree(
    id="fam-001",
    proband_id="child",
    individuals=[
        PedigreeIndividual(
            individual_id="mother",
            sex=Sex.FEMALE,
            affected_status=AffectedStatus.UNAFFECTED,
        ),
        PedigreeIndividual(
            individual_id="father",
            sex=Sex.MALE,
            affected_status=AffectedStatus.UNAFFECTED,
        ),
        PedigreeIndividual(
            individual_id="child",
            sex=Sex.MALE,
            affected_status=AffectedStatus.AFFECTED,
            maternal_id="mother",
            paternal_id="father",
        ),
        PedigreeIndividual(
            individual_id="sibling",
            sex=Sex.FEMALE,
            affected_status=AffectedStatus.UNAFFECTED,
            maternal_id="mother",
            paternal_id="father",
        ),
    ],
)


class TestPedigree:
    def test_creation(self) -> None:
        assert TRIO_PEDIGREE.id == "fam-001"
        assert len(TRIO_PEDIGREE.individuals) == 4

    def test_get_individual_found(self) -> None:
        ind = TRIO_PEDIGREE.get_individual("child")
        assert ind is not None
        assert ind.individual_id == "child"

    def test_get_individual_not_found(self) -> None:
        ind = TRIO_PEDIGREE.get_individual("nonexistent")
        assert ind is None

    def test_get_parents(self) -> None:
        parents = TRIO_PEDIGREE.get_parents("child")
        assert len(parents) == 2
        parent_ids = {p.individual_id for p in parents}
        assert parent_ids == {"mother", "father"}

    def test_get_parents_root(self) -> None:
        parents = TRIO_PEDIGREE.get_parents("mother")
        assert len(parents) == 0

    def test_get_parents_nonexistent(self) -> None:
        parents = TRIO_PEDIGREE.get_parents("nonexistent")
        assert len(parents) == 0

    def test_get_siblings(self) -> None:
        siblings = TRIO_PEDIGREE.get_siblings("child")
        assert len(siblings) == 1
        assert siblings[0].individual_id == "sibling"

    def test_get_siblings_root(self) -> None:
        siblings = TRIO_PEDIGREE.get_siblings("mother")
        assert len(siblings) == 0

    def test_get_affected(self) -> None:
        affected = TRIO_PEDIGREE.get_affected()
        assert len(affected) == 1
        assert affected[0].individual_id == "child"

    def test_from_dict(self) -> None:
        ped = Pedigree.model_validate(
            {
                "individuals": [
                    {"individual_id": "p1", "sex": "MALE"},
                    {"individual_id": "p2", "sex": "FEMALE"},
                ],
            }
        )
        assert len(ped.individuals) == 2

    def test_extra_fields(self) -> None:
        ind = PedigreeIndividual.model_validate({"individual_id": "p1", "custom_field": "value"})
        assert ind.custom_field == "value"
