# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors
import enum
import hashlib
import json
import logging
from datetime import datetime, timezone
from typing import Any, Optional

import attrs
import cattrs
import pyarrow as pa

from geneva.state.manager import BaseManager
from geneva.utils import current_user, escape_sql_string, retry_lance

MANIFEST_TABLE_NAME = "geneva_manifests"

_LOG = logging.getLogger(__name__)


@attrs.define
class GenevaManifest:
    """A Geneva Manifest represents the files and dependencies used
    in the execution environment."""

    # metadata
    name: str = attrs.field()
    version: Optional[str] = attrs.field(default=None)

    # properties needed to init the cluster
    pip: list[str] = attrs.field(default=[])
    py_modules: list[str] = attrs.field(default=[])
    head_image: Optional[str] = attrs.field(default=None)
    worker_image: Optional[str] = attrs.field(default=None)

    # transient properties, only used during initial upload
    skip_site_packages: bool = attrs.field(default=False)
    delete_local_zips: bool = attrs.field(default=False)
    local_zip_output_dir: Optional[str] = attrs.field(default=None)

    # internal generated properties
    zips: list[list[str]] = attrs.field(default=[[]])
    checksum: Optional[str] = attrs.field(default=None)
    created_at: datetime = attrs.field(
        factory=lambda: datetime.now(timezone.utc),
        metadata={"pa_type": pa.timestamp("us", tz="UTC")},
    )
    created_by: str = attrs.field(factory=current_user)

    def __attrs_post_init__(self) -> None:
        self.checksum = self.compute_checksum()

    def compute_checksum(self) -> str:
        """Generate a stable checksum of the manifest, ignoring the checksum field.
        The zip file names include the checksum of the contents so this hash is
        comprehensive.
        """
        checksum_exclude_fields = {
            "name",  # excluding name since this may be autogenerated.
            "checksum",  # this is what remains unique
            "created_at",
            "created_by",
            "delete_local_zips",
            "local_zip_output_dir",
        }
        data = attrs.asdict(
            self,
            recurse=True,
            filter=lambda a, v: a.name not in checksum_exclude_fields,
        )
        payload = json.dumps(data, sort_keys=True, separators=(",", ":"), default=str)
        return hashlib.md5(payload.encode("utf-8")).hexdigest()

    def as_dict(self) -> dict:
        return attrs.asdict(
            self,
            value_serializer=lambda obj, a, v: v.value
            if isinstance(v, enum.Enum)
            else v,
        )


class ManifestConfigManager(BaseManager):
    def get_table_name(self) -> str:
        return MANIFEST_TABLE_NAME

    def get_model(self) -> Any:
        return GenevaManifest("dummy")

    @retry_lance
    def upsert(self, manifest: GenevaManifest) -> None:
        val = manifest.as_dict()
        self.delete(manifest.name)
        self.get_table().add([val])
        # # note: merge_insert with fails with schema errors - use delete+add for now
        # todo: fix schema error on merge_insert?

    @retry_lance
    def list(self, limit: int = 1000) -> list[GenevaManifest]:
        res = self.get_table(True).search().limit(limit).to_arrow().to_pylist()
        return [_make_manifest(manifest) for manifest in res]

    @retry_lance
    def load(self, name: str) -> GenevaManifest | None:
        res = (
            self.get_table(True)
            .search()
            .where(f"name = '{escape_sql_string(name)}'")
            .limit(1)
            .to_arrow()
            .to_pylist()
        )
        if not res:
            return None
        return _make_manifest(res[0])

    @retry_lance
    def delete(self, name: str) -> None:
        self.get_table().delete(f"name = '{escape_sql_string(name)}'")


def _make_manifest(args: dict) -> GenevaManifest:
    converter = cattrs.Converter()
    converter.register_structure_hook(
        datetime,
        lambda ts, _: datetime.fromisoformat(ts.replace("Z", "+00:00"))
        if isinstance(ts, str)
        else ts,
    )
    res = converter.structure(args, GenevaManifest)

    return res
