# 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 TYPE_CHECKING, Any, Optional

import attrs

if TYPE_CHECKING:
    from geneva.manifest.builder import (
        CondaManifestBuilder,
        PipManifestBuilder,
        SiteManifestBuilder,
    )
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=True)
    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)

    requirements_path: Optional[str] = attrs.field(default=None)
    # conda: stored as JSON string in Arrow (Lance doesn't support Map types)
    # Ray accepts complex structures like {"dependencies": [...], "channels": [...]}
    conda: dict[str, Any] = attrs.field(
        default={},
        metadata={"pa_type": pa.string()},
    )
    conda_environment_path: Optional[str] = attrs.field(default=None)

    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()

    # -------------------------------------------------------------------------
    # Factory Methods for Type-Safe Builders
    # -------------------------------------------------------------------------

    @staticmethod
    def create_pip(name: str) -> "PipManifestBuilder":
        """Create a pip-based manifest builder.

        Example:
            manifest = GenevaManifest.create_pip("my-manifest").pip(["numpy"]).build()
        """
        from geneva.manifest.builder import PipManifestBuilder

        return PipManifestBuilder.create(name)

    @staticmethod
    def create_conda(name: str) -> "CondaManifestBuilder":
        """Create a conda-based manifest builder.

        Example:
            manifest = GenevaManifest.create_conda("my-manifest").conda({...}).build()
        """
        from geneva.manifest.builder import CondaManifestBuilder

        return CondaManifestBuilder.create(name)

    @staticmethod
    def create_site(name: str) -> "SiteManifestBuilder":
        """Create a site-packages manifest builder.

        Uploads local site-packages without external dependencies.
        upload_site_packages defaults to True.

        Example:
            manifest = GenevaManifest.create_site("my-manifest").build()
        """
        from geneva.manifest.builder import SiteManifestBuilder

        return SiteManifestBuilder.create(name)

    def as_dict(self) -> dict:
        return attrs.asdict(
            self,
            value_serializer=lambda obj, a, v: (
                v.value
                if isinstance(v, enum.Enum)
                else json.dumps(v)
                if isinstance(v, dict)
                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:
    import cattrs

    # Shallow copy; conda is stored as JSON string in Arrow, restore as dict.
    args = dict(args)
    if "conda" in args:
        if isinstance(args["conda"], str):
            args["conda"] = json.loads(args["conda"])
        elif args["conda"] is None:
            args["conda"] = {}

    converter = cattrs.Converter()
    converter.register_structure_hook(
        datetime,
        lambda ts, _: (
            datetime.fromisoformat(ts.replace("Z", "+00:00"))
            if isinstance(ts, str)
            else ts
        ),
    )
    return converter.structure(args, GenevaManifest)
