"""Data models used for cache file reproduction across the scripts."""

from __future__ import annotations

import logging
import struct
from datetime import datetime
from enum import StrEnum
from pathlib import Path
from py_compile import PycInvalidationMode
from typing import Annotated, Any, NamedTuple

from pydantic import (
    AliasChoices,
    BaseModel,
    Field,
    ValidatorFunctionWrapHandler,
    field_serializer,
    field_validator,
    model_validator,
)


class Package(BaseModel):
    name: str
    version: str
    proton_path: Path
    distributions: list[Distribution]

    @field_serializer("proton_path")
    def serialize_path(self, value: Path) -> str:
        return str(value)


class Distribution(BaseModel):
    kind: DistributionKind
    dist_path: Path
    cache_files: list[CacheFile]

    @field_serializer("dist_path")
    def serialize_dist_path(self, value: Path) -> str:
        return str(value)


class CacheFile(BaseModel):
    source_path: Annotated[
        Path, Field(validation_alias=AliasChoices("source", "source_path"))
    ]
    cache_path: Annotated[
        Path, Field(validation_alias=AliasChoices("cache", "cache_path"))
    ]
    python_version: Annotated[
        PythonVersion, Field(validation_alias=AliasChoices("version", "python_version"))
    ]

    @field_validator("source_path", mode="wrap")
    @classmethod
    def fix_source_path(cls, value: Any, handler: ValidatorFunctionWrapHandler):
        # TODO: this should be removed once #5 is done and the package list recreated.
        path: Path = handler(value)
        return (
            path.with_suffix("").with_suffix(".py") if path.suffix == ".pyc" else path
        )

    @field_serializer("cache_path")
    def serialize_cache_path(self, value: Path) -> str:
        return str(value)

    @field_serializer("source_path")
    def serialize_source_path(self, value: Path) -> str:
        return str(value)

    @property
    def optimization_level(self) -> int:
        stem = self.cache_path.stem
        try:
            parts = stem.split(".")
            if parts[-1].startswith("opt"):
                return int(parts[-1].split("-")[-1])
        except Exception:
            logging.exception("Failed to get optimization level for %s", self)
            return -1
        return -1

    @property
    def tag(self) -> str | None:
        stem = self.cache_path.stem
        try:
            name_parts = stem.split(".")
            for p in name_parts[::-1]:
                if p.startswith("cpython"):
                    tag_parts = p.split("-")
                    if len(tag_parts) > 2:
                        return "-".join(tag_parts[2:])
                    break
        except Exception:
            logging.exception("Failed to get tag for %s", self)
            return None
        return None


class DistributionKind(StrEnum):
    SourceDistribution = "sdist"
    Wheel = "wheel"
    Egg = "egg"


class PythonVersion(StrEnum):
    CPython26 = "cpython-26"
    CPython27 = "cpython-27"
    CPython32 = "cpython-32"
    CPython33 = "cpython-33"
    CPython34 = "cpython-34"
    CPython35 = "cpython-35"
    CPython36 = "cpython-36"
    CPython37 = "cpython-37"
    CPython38 = "cpython-38"
    CPython39 = "cpython-39"
    CPython310 = "cpython-310"
    CPython311 = "cpython-311"
    CPython312 = "cpython-312"
    CPython313 = "cpython-313"
    CPython314 = "cpython-314"

    def _extract_versions(
        self, other: Any
    ) -> tuple[MajorMinorVersion, MajorMinorVersion]:
        if not isinstance(other, PythonVersion):
            raise TypeError("'<' not supported")
        my_version = self.split("-")[-1]
        other_version = other.split("-")[-1]
        my_major, my_minor = my_version[0], my_version[1:]
        other_major, other_minor = other_version[0], other_version[1:]
        return MajorMinorVersion(int(my_major), int(my_minor)), MajorMinorVersion(
            int(other_major), int(other_minor)
        )

    @property
    def major_minor_string(self) -> str:
        my_version = self.split("-")[-1]
        my_major, my_minor = my_version[0], my_version[1:]
        return f"{my_major}.{my_minor}"

    def __gt__(self, other: Any) -> bool:
        my_version, other_version = self._extract_versions(other)
        return my_version > other_version

    def __ge__(self, other: Any) -> bool:
        my_version, other_version = self._extract_versions(other)
        return my_version >= other_version

    def __lt__(self, other: Any) -> bool:
        my_version, other_version = self._extract_versions(other)
        return my_version < other_version

    def __le__(self, other: Any) -> bool:
        my_version, other_version = self._extract_versions(other)
        return my_version <= other_version


class MajorMinorVersion(NamedTuple):
    major: int
    minor: int

    def __gt__(self, other: Any) -> bool:
        if not isinstance(other, MajorMinorVersion):
            raise TypeError
        if self.major > other.major:
            return True
        if self.minor > other.minor:
            return True
        return False

    def __ge__(self, other: Any) -> bool:
        if not isinstance(other, MajorMinorVersion):
            raise TypeError
        if self == other:
            return True
        return self > other

    def __lt__(self, other: Any) -> bool:
        if not isinstance(other, MajorMinorVersion):
            raise TypeError
        if self.major < other.major:
            return True
        if self.minor < other.minor:
            return True
        return False

    def __le__(self, other: Any) -> bool:
        if not isinstance(other, MajorMinorVersion):
            raise TypeError
        if self == other:
            return True
        return self < other


class ParseResponse(BaseModel):
    header: PyCacheHeader
    bytecode_table: dict[str, bytes]

    @field_validator("bytecode_table", mode="before")
    @classmethod
    def validate_bytecode_table(cls, value: dict[str, str]) -> dict[str, bytes]:
        d = {}
        for k, v in value.items():
            d[k] = bytes.fromhex(v)
        return d


class PyCacheHeader(BaseModel):
    magic: bytes
    invalidation_mode: PycInvalidationMode
    # only available on hash modes
    hash: bytes | None = Field(default=None)
    # only available on timestamp mode
    source_modification_time: datetime | None
    source_size: int | None

    @model_validator(mode="before")
    @classmethod
    def validate(cls, value: Any):
        source_hash = None
        mod_dt = None
        size = None

        v = bytes.fromhex(value)
        long_header = len(v) == 16
        magic_number, v = v[:4], v[4:]
        if long_header:
            flags, v = struct.unpack("<L", v[:4])[0], v[4:]
            if flags & 0x01:
                invalidation_mode = (
                    PycInvalidationMode.CHECKED_HASH
                    if flags & 0x02
                    else PycInvalidationMode.UNCHECKED_HASH
                )
            else:
                invalidation_mode = PycInvalidationMode.TIMESTAMP
        else:
            invalidation_mode = PycInvalidationMode.TIMESTAMP
        if invalidation_mode == PycInvalidationMode.TIMESTAMP:
            mod_date_bytes, v = v[:4], v[4:]
            mod_dt = datetime.fromtimestamp(struct.unpack("<L", mod_date_bytes)[0])
            size = struct.unpack("<L", v)[0]
        else:
            source_hash = v
        return {
            "magic": magic_number,
            "invalidation_mode": invalidation_mode,
            "hash": source_hash,
            "source_modification_time": mod_dt,
            "source_size": size,
        }


class ReproductionResult(BaseModel):
    package_name: str
    package_version: str
    distributions: list[DistributionResult]


class DistributionResult(BaseModel):
    distribution_kind: DistributionKind
    status: DistributionStatus
    files: list[FileReproductionResult]

    @staticmethod
    def unsupported(distribution_kind: DistributionKind) -> DistributionResult:
        return DistributionResult(
            distribution_kind=distribution_kind,
            status=DistributionStatus.UnsupportedKind,
            files=[],
        )


class FileReproductionResult(BaseModel):
    file_name: str
    status: FileStatus
    python_version: str
    python_version_choice: PythonVersionChoice
    invalidation_mode: InvalidationMode
    cache_valid: bool
    source_modification_time: datetime | None
    cache_modification_time: datetime | None
    optimization_level: int
    additional_tag: str | None

    @staticmethod
    def unsupported_python(file: CacheFile) -> FileReproductionResult:
        return FileReproductionResult(
            file_name=str(file.cache_path),
            status=FileStatus.UnsupportedPythonVersion,
            python_version="",
            python_version_choice=PythonVersionChoice.Unknown,
            invalidation_mode=InvalidationMode.Unknown,
            cache_valid=False,
            source_modification_time=None,
            cache_modification_time=None,
            optimization_level=file.optimization_level,
            additional_tag=file.tag,
        )

    @staticmethod
    def source_unavailable(file: CacheFile) -> FileReproductionResult:
        return FileReproductionResult(
            file_name=str(file.cache_path),
            status=FileStatus.NoSourceFile,
            python_version="",
            python_version_choice=PythonVersionChoice.Unknown,
            invalidation_mode=InvalidationMode.Unknown,
            cache_valid=False,
            source_modification_time=None,
            cache_modification_time=None,
            optimization_level=file.optimization_level,
            additional_tag=file.tag,
        )

    @staticmethod
    def cache_uncompilable(
        file: CacheFile,
        invalidation_mode: PycInvalidationMode,
        cache_valid: bool,
        source_mtime: datetime,
        cache_mtime: datetime,
    ) -> FileReproductionResult:
        return FileReproductionResult(
            file_name=str(file.cache_path),
            status=FileStatus.CacheUncompilable,
            python_version="",
            python_version_choice=PythonVersionChoice.Unknown,
            invalidation_mode=InvalidationMode.from_pyc_invalidation_mode(
                invalidation_mode
            ),
            cache_valid=cache_valid,
            source_modification_time=source_mtime,
            cache_modification_time=cache_mtime,
            optimization_level=file.optimization_level,
            additional_tag=file.tag,
        )

    @staticmethod
    def original_cache_unparsable(file: CacheFile) -> FileReproductionResult:
        return FileReproductionResult(
            file_name=str(file.cache_path),
            status=FileStatus.OriginalCacheUnparsable,
            python_version="",
            python_version_choice=PythonVersionChoice.Unknown,
            invalidation_mode=InvalidationMode.Unknown,
            cache_valid=False,
            source_modification_time=None,
            cache_modification_time=None,
            optimization_level=file.optimization_level,
            additional_tag=file.tag,
        )

    @staticmethod
    def rebuilt_cache_unparsable(
        file: CacheFile, invalidation_mode: PycInvalidationMode, cache_valid: bool
    ) -> FileReproductionResult:
        return FileReproductionResult(
            file_name=str(file.cache_path),
            status=FileStatus.RebuiltCacheUnparsable,
            python_version="",
            python_version_choice=PythonVersionChoice.Unknown,
            invalidation_mode=InvalidationMode.from_pyc_invalidation_mode(
                invalidation_mode
            ),
            cache_valid=cache_valid,
            source_modification_time=None,
            cache_modification_time=None,
            optimization_level=file.optimization_level,
            additional_tag=file.tag,
        )


class DistributionStatus(StrEnum):
    Success = "success"
    UnsupportedKind = "unsupported_kind"


class FileStatus(StrEnum):
    Reproducible = "reproducible"
    UnsupportedPythonVersion = "unsupported_python_version"
    NoSourceFile = "no_source_file"
    OriginalCacheUnparsable = "original_cache_unparsable"
    CacheUncompilable = "cache_uncompilable"
    RebuiltCacheUnparsable = "rebuilt_cache_unparsable"
    Irreproducible = "irreproducible"


class InvalidationMode(StrEnum):
    Unknown = "unknown"
    Timestamp = "timestamp"
    CheckedHash = "checked_hash"
    UncheckedHash = "unchecked_hash"

    @staticmethod
    def from_pyc_invalidation_mode(pyc_invalidation_mode: PycInvalidationMode):
        match pyc_invalidation_mode:
            case PycInvalidationMode.TIMESTAMP:
                return InvalidationMode.Timestamp
            case PycInvalidationMode.CHECKED_HASH:
                return InvalidationMode.CheckedHash
            case PycInvalidationMode.UNCHECKED_HASH:
                return InvalidationMode.UncheckedHash
            case _:
                return InvalidationMode.Unknown


class PythonVersionChoice(StrEnum):
    LatestMagic = "latest_magic"
    OtherMagic = "other_magic"
    LatestFileName = "latest_file_name"
    OtherFileName = "other_file_name"
    Unknown = ""