"""Pydantic models for GA4GH VRS (Variation Representation Specification) v2.

VRS provides a computational framework for the precise representation
and sharing of variation data. This module implements the core value
objects and computed identifier logic.
"""

from __future__ import annotations

import hashlib
import json
from enum import Enum
from typing import Any

from ai4gh.core.types import GA4GHBase


class SequenceString(GA4GHBase):
    """A literal sequence string."""

    sequence: str
    type: str = "SequenceString"


class SequenceReference(GA4GHBase):
    """A reference to a named sequence (by refget accession)."""

    refgetAccession: str
    type: str = "SequenceReference"
    residueAlphabet: str | None = None


class SequenceLocation(GA4GHBase):
    """A contiguous region of a sequence."""

    sequenceReference: SequenceReference | None = None
    start: int | list[int] | None = None
    end: int | list[int] | None = None
    type: str = "SequenceLocation"


class LengthExpression(GA4GHBase):
    """An expression of sequence length."""

    length: int | dict[str, int]
    type: str = "LengthExpression"
    sequence: str | None = None
    repeatSubunitLength: int | None = None


class LiteralSequenceExpression(GA4GHBase):
    """A literal sequence as the state of a variation."""

    sequence: str
    type: str = "LiteralSequenceExpression"


class ReferenceLengthExpression(GA4GHBase):
    """A length expression referencing the reference sequence."""

    length: int | dict[str, int]
    type: str = "ReferenceLengthExpression"
    repeatSubunitLength: int | None = None


class CopyChange(str, Enum):
    """Copy number change types."""

    EFO_0030069 = "efo:0030069"  # complete genomic loss
    EFO_0020073 = "efo:0020073"  # low-level loss
    EFO_0030068 = "efo:0030068"  # loss
    EFO_0030067 = "efo:0030067"  # partial loss
    EFO_0030064 = "efo:0030064"  # regional base ploidy
    EFO_0030070 = "efo:0030070"  # gain
    EFO_0030071 = "efo:0030071"  # low-level gain
    EFO_0030072 = "efo:0030072"  # high-level gain


class Allele(GA4GHBase):
    """VRS Allele — a contiguous change at a SequenceLocation."""

    location: SequenceLocation
    state: LiteralSequenceExpression | ReferenceLengthExpression
    type: str = "Allele"
    id: str | None = None


class Haplotype(GA4GHBase):
    """VRS Haplotype — a set of co-occurring Alleles on a single molecule."""

    members: list[Allele | dict[str, Any]]
    type: str = "Haplotype"
    id: str | None = None


class CopyNumberCount(GA4GHBase):
    """VRS CopyNumberCount — an absolute copy number at a location."""

    location: SequenceLocation
    copies: int | dict[str, int]
    type: str = "CopyNumberCount"
    id: str | None = None


class CopyNumberChange(GA4GHBase):
    """VRS CopyNumberChange — a relative copy number change at a location."""

    location: SequenceLocation
    copy_change: CopyChange
    type: str = "CopyNumberChange"
    id: str | None = None


class Adjacency(GA4GHBase):
    """VRS Adjacency — a junction between two sequence locations."""

    adj_left: SequenceLocation
    adj_right: SequenceLocation
    type: str = "Adjacency"
    id: str | None = None
    linker_sequence: str | None = None


def ga4gh_digest(data: dict[str, Any]) -> str:
    """Compute a GA4GH digest for a VRS object.

    The GA4GH digest uses SHA-512 truncated to 24 bytes, then
    base64url-encoded (no padding).

    Args:
        data: A serializable VRS value object dict.

    Returns:
        The base64url-encoded truncated SHA-512 digest.
    """
    import base64

    canonical = json.dumps(data, sort_keys=True, separators=(",", ":"))
    sha = hashlib.sha512(canonical.encode("utf-8")).digest()[:24]
    return base64.urlsafe_b64encode(sha).rstrip(b"=").decode("ascii")


def ga4gh_identify(vrs_obj: dict[str, Any], vrs_type: str) -> str:
    """Compute a GA4GH Computed Identifier for a VRS object.

    The identifier has the form: ga4gh:{type}.{digest}

    Args:
        vrs_obj: A serializable VRS value object dict.
        vrs_type: The VRS type prefix (e.g., "VA" for Allele, "VH" for Haplotype).

    Returns:
        A GA4GH computed identifier string.
    """
    digest = ga4gh_digest(vrs_obj)
    return f"ga4gh:{vrs_type}.{digest}"
