"""Pydantic models for GA4GH Phenopackets v2.

Phenopackets is an open standard for sharing disease and phenotype
information. This module provides Pydantic v2 models for the core
building blocks of the schema.
"""

from __future__ import annotations

from enum import Enum
from typing import Any

from ai4gh.core.types import GA4GHBase


class Sex(str, Enum):
    """Biological sex."""

    UNKNOWN_SEX = "UNKNOWN_SEX"
    FEMALE = "FEMALE"
    MALE = "MALE"
    OTHER_SEX = "OTHER_SEX"


class KaryotypicSex(str, Enum):
    """Karyotypic sex."""

    UNKNOWN_KARYOTYPE = "UNKNOWN_KARYOTYPE"
    XX = "XX"
    XY = "XY"
    XO = "XO"
    XXY = "XXY"
    XXX = "XXX"
    XXYY = "XXYY"
    XXXY = "XXXY"
    XXXX = "XXXX"
    XYY = "XYY"
    OTHER_KARYOTYPE = "OTHER_KARYOTYPE"


class VitalStatus(str, Enum):
    """Vital status of the individual."""

    UNKNOWN_STATUS = "UNKNOWN_STATUS"
    ALIVE = "ALIVE"
    DECEASED = "DECEASED"


class AcmgPathogenicityClassification(str, Enum):
    """ACMG pathogenicity classification."""

    NOT_PROVIDED = "NOT_PROVIDED"
    BENIGN = "BENIGN"
    LIKELY_BENIGN = "LIKELY_BENIGN"
    UNCERTAIN_SIGNIFICANCE = "UNCERTAIN_SIGNIFICANCE"
    LIKELY_PATHOGENIC = "LIKELY_PATHOGENIC"
    PATHOGENIC = "PATHOGENIC"


class TherapeuticActionability(str, Enum):
    """Therapeutic actionability of a genomic interpretation."""

    UNKNOWN_ACTIONABILITY = "UNKNOWN_ACTIONABILITY"
    NOT_ACTIONABLE = "NOT_ACTIONABLE"
    ACTIONABLE = "ACTIONABLE"


class ProgressStatus(str, Enum):
    """Progress status of an interpretation."""

    UNKNOWN_PROGRESS = "UNKNOWN_PROGRESS"
    IN_PROGRESS = "IN_PROGRESS"
    COMPLETED = "COMPLETED"
    SOLVED = "SOLVED"
    UNSOLVED = "UNSOLVED"


class OntologyClass(GA4GHBase):
    """An ontology class (term) reference."""

    id: str
    label: str | None = None


class ExternalReference(GA4GHBase):
    """A reference to an external resource."""

    id: str
    reference: str | None = None
    description: str | None = None


class Evidence(GA4GHBase):
    """Evidence for a phenotypic feature or diagnosis."""

    evidence_code: OntologyClass
    reference: ExternalReference | None = None


class TimeElement(GA4GHBase):
    """Flexible time representation (age, age range, timestamp, etc.)."""

    gestational_age: dict[str, int] | None = None
    age: dict[str, str] | None = None
    age_range: dict[str, Any] | None = None
    ontology_class: OntologyClass | None = None
    timestamp: str | None = None
    interval: dict[str, Any] | None = None


class PhenotypicFeature(GA4GHBase):
    """A phenotypic feature observed in an individual."""

    type: OntologyClass
    excluded: bool = False
    severity: OntologyClass | None = None
    modifiers: list[OntologyClass] | None = None
    onset: TimeElement | None = None
    resolution: TimeElement | None = None
    evidence: list[Evidence] | None = None
    description: str | None = None


class Disease(GA4GHBase):
    """A disease diagnosed in an individual."""

    term: OntologyClass
    excluded: bool = False
    onset: TimeElement | None = None
    resolution: TimeElement | None = None
    disease_stage: list[OntologyClass] | None = None
    clinical_tnm_finding: list[OntologyClass] | None = None
    primary_site: OntologyClass | None = None
    laterality: OntologyClass | None = None


class Procedure(GA4GHBase):
    """A clinical procedure."""

    code: OntologyClass
    body_site: OntologyClass | None = None
    performed: TimeElement | None = None


class Biosample(GA4GHBase):
    """A biosample (biological sample) from an individual."""

    id: str
    individual_id: str | None = None
    derived_from_id: str | None = None
    description: str | None = None
    sampled_tissue: OntologyClass | None = None
    sample_type: OntologyClass | None = None
    phenotypic_features: list[PhenotypicFeature] | None = None
    taxonomy: OntologyClass | None = None
    time_of_collection: TimeElement | None = None
    histological_diagnosis: OntologyClass | None = None
    tumor_progression: OntologyClass | None = None
    tumor_grade: list[OntologyClass] | None = None
    diagnostic_markers: list[OntologyClass] | None = None
    procedure: Procedure | None = None
    material_sample: OntologyClass | None = None
    sample_processing: OntologyClass | None = None
    sample_storage: OntologyClass | None = None


class Individual(GA4GHBase):
    """An individual (patient, subject)."""

    id: str
    alternate_ids: list[str] | None = None
    date_of_birth: str | None = None
    time_at_last_encounter: TimeElement | None = None
    vital_status: VitalStatus | None = None
    sex: Sex | None = None
    karyotypic_sex: KaryotypicSex | None = None
    gender: OntologyClass | None = None
    taxonomy: OntologyClass | None = None


class VariantInterpretation(GA4GHBase):
    """Interpretation of a genomic variant."""

    acmg_pathogenicity_classification: AcmgPathogenicityClassification = (
        AcmgPathogenicityClassification.NOT_PROVIDED
    )
    therapeutic_actionability: TherapeuticActionability = (
        TherapeuticActionability.UNKNOWN_ACTIONABILITY
    )
    variation_descriptor: dict[str, Any] | None = None


class GenomicInterpretation(GA4GHBase):
    """A genomic interpretation for an individual or biosample."""

    subject_or_biosample_id: str
    interpretation_status: ProgressStatus = ProgressStatus.UNKNOWN_PROGRESS
    variant_interpretation: VariantInterpretation | None = None
    gene_descriptor: dict[str, Any] | None = None


class Diagnosis(GA4GHBase):
    """A diagnosis with associated genomic interpretations."""

    disease: OntologyClass
    genomic_interpretations: list[GenomicInterpretation] | None = None


class Interpretation(GA4GHBase):
    """An interpretation for a phenopacket."""

    id: str
    progress_status: ProgressStatus = ProgressStatus.UNKNOWN_PROGRESS
    diagnosis: Diagnosis | None = None
    summary: str | None = None


class MetaData(GA4GHBase):
    """Metadata about a phenopacket."""

    created: str | None = None
    created_by: str | None = None
    submitted_by: str | None = None
    resources: list[dict[str, Any]] | None = None
    phenopacket_schema_version: str | None = None
    external_references: list[ExternalReference] | None = None


class MedicalAction(GA4GHBase):
    """A medical action (treatment, procedure, etc.)."""

    procedure: Procedure | None = None
    treatment: dict[str, Any] | None = None
    radiation_therapy: dict[str, Any] | None = None
    therapeutic_regimen: dict[str, Any] | None = None
    adverse_events: list[OntologyClass] | None = None
    treatment_target: OntologyClass | None = None
    treatment_intent: OntologyClass | None = None
    response_to_treatment: OntologyClass | None = None
    treatment_termination_reason: OntologyClass | None = None


class Measurement(GA4GHBase):
    """A clinical measurement or observation."""

    assay: OntologyClass
    value: dict[str, Any] | None = None
    complex_value: dict[str, Any] | None = None
    time_observed: TimeElement | None = None
    procedure: Procedure | None = None
    description: str | None = None


class Phenopacket(GA4GHBase):
    """A GA4GH Phenopacket — the top-level element."""

    id: str
    subject: Individual | None = None
    phenotypic_features: list[PhenotypicFeature] | None = None
    measurements: list[Measurement] | None = None
    biosamples: list[Biosample] | None = None
    interpretations: list[Interpretation] | None = None
    diseases: list[Disease] | None = None
    medical_actions: list[MedicalAction] | None = None
    meta_data: MetaData | None = None
    files: list[dict[str, Any]] | None = None


class Family(GA4GHBase):
    """A collection of Phenopackets for a family."""

    id: str
    proband: Phenopacket
    relatives: list[Phenopacket] | None = None
    consanguinous_parents: bool | None = None
    pedigree: dict[str, Any] | None = None
    meta_data: MetaData | None = None
    files: list[dict[str, Any]] | None = None


class Cohort(GA4GHBase):
    """A collection of Phenopackets for a cohort study."""

    id: str
    description: str | None = None
    members: list[Phenopacket] | None = None
    meta_data: MetaData | None = None
    files: list[dict[str, Any]] | None = None
