"""Pydantic models for GA4GH Cat-VRS (Categorical Variation Representation).

Cat-VRS extends VRS with categorical variation concepts:
constraints on variation that describe classes of variation
rather than specific instances (e.g., "any loss-of-function
variant in BRCA1").
"""

from __future__ import annotations

from enum import Enum

from ai4gh.core.types import GA4GHBase
from ai4gh.specs.vrs.models import (  # noqa: TC001 — needed at runtime by Pydantic
    CopyChange,
    SequenceLocation,
    SequenceReference,
)


class DefiningContextType(str, Enum):
    """Types of defining context for categorical variation."""

    GENOMIC = "genomic"
    TRANSCRIPT = "transcript"
    PROTEIN = "protein"


class CanonicalVariationConstraint(GA4GHBase):
    """Constraint describing a class of canonical variations.

    Represents a constraint that defines a category of variation
    rather than a specific allele — e.g., any variant in a gene region.
    """

    type: str = "CanonicalVariationConstraint"
    defining_context: SequenceLocation | SequenceReference | None = None
    complement: bool = False
    copies: int | None = None
    copy_change: CopyChange | None = None


class CategoricalVariation(GA4GHBase):
    """A categorical variation: a class of related variants.

    Represents a set of variations grouped by a common
    biological or clinical meaning.
    """

    type: str = "CategoricalVariation"
    id: str | None = None
    label: str | None = None
    description: str | None = None
    members: list[CanonicalVariationConstraint] | None = None
    defining_context_type: DefiningContextType | None = None


class ProteinSequenceConsequence(GA4GHBase):
    """Consequence of a variation at the protein level."""

    type: str = "ProteinSequenceConsequence"
    defining_context: SequenceLocation | None = None
    label: str | None = None
    description: str | None = None
