"""In-memory models for Zarr group and consolidated metadata documents."""

from __future__ import annotations

import copy
import dataclasses
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Final, Literal, cast

from typing_extensions import TypedDict, Unpack

from zarr_metadata.model._array import (
    ATTRIBUTES_STORE_KEY_V2,
    ZarrV3ArrayMetadata,
    must_understand_subset,
)
from zarr_metadata.model._sentinel import UNSET
from zarr_metadata.model._validation import (
    GROUP_METADATA_STANDARD_KEYS_V3,
    MetadataValidationError,
    ValidationProblem,
    arrays_to_tuples,
    dump_store_json,
    load_store_json,
    parse_group_metadata_v2,
    parse_group_metadata_v3,
    validate_consolidated_metadata_v3,
    validate_json,
)

if TYPE_CHECKING:
    from zarr_metadata._common import JSONValue
    from zarr_metadata.model._array import ZarrV2AttributesStoreKey
    from zarr_metadata.v2.group import ZarrV2GroupMetadataJSON
    from zarr_metadata.v3.array import ZarrV3ExtensionField
    from zarr_metadata.v3.consolidated import ZarrV3ConsolidatedMetadataJSON
    from zarr_metadata.v3.group import ZarrV3GroupMetadataJSON

ZarrV3GroupMetadataStoreKey = Literal["zarr.json"]
GROUP_METADATA_STORE_KEY_V3: Final[ZarrV3GroupMetadataStoreKey] = "zarr.json"

ZarrV2GroupMetadataStoreKey = Literal[".zgroup"]
GROUP_METADATA_STORE_KEY_V2: Final[ZarrV2GroupMetadataStoreKey] = ".zgroup"

ZarrV2ConsolidatedMetadataStoreKey = Literal[".zmetadata"]
CONSOLIDATED_METADATA_STORE_KEY_V2: Final[ZarrV2ConsolidatedMetadataStoreKey] = ".zmetadata"

# The key under which consolidated metadata is embedded in a v3 group document.
# This is a reference-implementation convention (not a spec artifact), stored
# as an extension field on the group's `zarr.json`.
CONSOLIDATED_METADATA_KEY_V3: Final = "consolidated_metadata"


class ZarrV3GroupMetadataPartial(TypedDict, total=False):
    """
    Partial form of the constructor-settable fields of `ZarrV3GroupMetadata`.

    Every key is optional and typed with the model's own value types, so it
    describes valid keyword arguments to `ZarrV3GroupMetadata.update` and
    `create_default`. The `init=False` fields `zarr_format` and `node_type`
    are intentionally excluded, since they cannot be passed to
    `dataclasses.replace`.

    Drift between this type and the model's settable fields is prevented by
    `tests/model/test_group.py::test_group_partial_keys_match_settable_model_fields`.
    """

    attributes: dict[str, JSONValue]
    consolidated_metadata: ZarrV3ConsolidatedMetadata | UNSET
    extra_fields: dict[str, ZarrV3ExtensionField]


@dataclass(frozen=True, slots=True, kw_only=True)
class ZarrV3GroupMetadata:
    """In-memory model of a v3 group metadata document.

    A canonical, semantically lossless representation of the `zarr.json`
    content for a group. The `consolidated_metadata` reference-implementation
    convention is modeled as a typed field holding thin child models; every
    other unknown top-level key lands in `extra_fields` verbatim.
    """

    zarr_format: Literal[3] = field(default=3, init=False)
    node_type: Literal["group"] = field(default="group", init=False)
    attributes: dict[str, JSONValue]
    consolidated_metadata: ZarrV3ConsolidatedMetadata | UNSET
    extra_fields: dict[str, ZarrV3ExtensionField]

    def __post_init__(self) -> None:
        reserved = GROUP_METADATA_STANDARD_KEYS_V3 | {CONSOLIDATED_METADATA_KEY_V3}
        if set(self.extra_fields.keys()).intersection(reserved):
            raise MetadataValidationError(
                [
                    ValidationProblem(
                        ("extra_fields",),
                        "Extra fields cannot overlap with standard Zarr V3 group metadata fields",
                        "invalid_value",
                    )
                ]
            )

    @classmethod
    def create_default(cls, **overrides: Unpack[ZarrV3GroupMetadataPartial]) -> ZarrV3GroupMetadata:
        """
        Create a default (empty) v3 group metadata model, with optional overrides.

        The default is a structurally-valid group with no attributes — the group
        analog of `list()` returning `[]`. Any field can be overridden by keyword
        (the same fields accepted by `update`).
        """
        default = cls(attributes={}, consolidated_metadata=UNSET, extra_fields={})
        return default.update(**overrides)

    def update(self, **kwargs: Unpack[ZarrV3GroupMetadataPartial]) -> ZarrV3GroupMetadata:
        """
        Return a new `ZarrV3GroupMetadata` with the given fields updated.

        Only the constructor-settable fields listed in
        `ZarrV3GroupMetadataPartial` can be updated; the fixed `zarr_format` /
        `node_type` are rejected at the type level. Each given field fully
        replaces its previous value, including `extra_fields`.
        """
        return dataclasses.replace(self, **kwargs)

    def to_json(self) -> ZarrV3GroupMetadataJSON:
        # to_json output shares no mutable state with the model: every value
        # that can hold a mutable container is deep-copied.
        out: ZarrV3GroupMetadataJSON = {
            "zarr_format": self.zarr_format,
            "node_type": self.node_type,
        }
        if len(self.attributes) > 0:
            out["attributes"] = copy.deepcopy(self.attributes)
        if self.consolidated_metadata is not UNSET:
            # Consolidated metadata is a known non-core top-level JSON field.
            out[CONSOLIDATED_METADATA_KEY_V3] = cast(
                "ZarrV3ExtensionField", self.consolidated_metadata.to_json()
            )
        for key, value in self.extra_fields.items():
            out[key] = copy.deepcopy(value)
        return out

    @classmethod
    def from_json(cls, data: object) -> ZarrV3GroupMetadata:
        parsed = parse_group_metadata_v3(arrays_to_tuples(data))
        # Cast for narrowing across standard and arbitrary extra TypedDict items.
        consolidated_raw = cast("object", parsed.get(CONSOLIDATED_METADATA_KEY_V3, UNSET))
        consolidated: ZarrV3ConsolidatedMetadata | UNSET
        if consolidated_raw is UNSET or consolidated_raw is None:
            # consolidated_metadata: null was written by a historical
            # zarr-python bug; it gets no model representation. It is read as
            # absence and never written back — repaired, not preserved.
            consolidated = UNSET
        else:
            consolidated = ZarrV3ConsolidatedMetadata.from_json(consolidated_raw)
        # Sound cast: the TypedDict types all non-standard keys as its
        # `extra_items` (`ZarrV3ExtensionField`); the comprehension's inferred value
        # type is the union over ALL keys because the key filter cannot narrow it.
        extra_fields = cast(
            "dict[str, ZarrV3ExtensionField]",
            {
                k: v
                for k, v in parsed.items()
                if k not in GROUP_METADATA_STANDARD_KEYS_V3 and k != CONSOLIDATED_METADATA_KEY_V3
            },
        )
        return cls(
            attributes=dict(parsed.get("attributes", {})),
            consolidated_metadata=consolidated,
            extra_fields=extra_fields,
        )

    @property
    def must_understand_fields(self) -> dict[str, ZarrV3ExtensionField]:
        """Extra fields the reader is obligated to understand.

        Everything in `extra_fields` not explicitly waived with
        `must_understand: false` (the spec's implicit-true rule). A compliant
        reader MUST fail to open the group if this contains any field it does
        not recognize; the model layer only partitions by obligation, since
        recognition is reader-specific.
        """
        return must_understand_subset(self.extra_fields)

    @classmethod
    def from_key_value(cls, mapping: Mapping[str, bytes]) -> ZarrV3GroupMetadata:
        return cls.from_json(load_store_json(mapping, GROUP_METADATA_STORE_KEY_V3))

    def to_key_value(
        self, *, indent: int | str | None = None
    ) -> Mapping[ZarrV3GroupMetadataStoreKey, bytes]:
        return {GROUP_METADATA_STORE_KEY_V3: dump_store_json(self.to_json(), indent=indent)}


@dataclass(frozen=True, slots=True, kw_only=True)
class ZarrV3ConsolidatedMetadata:
    """In-memory model of v3 inline consolidated metadata.

    Models the reference-implementation convention where consolidated metadata
    is embedded as an extension field on a group's `zarr.json`. Each entry in
    `metadata` is a complete child document, held as a thin array or group
    model. `must_understand` is typed permissively as `bool` to mirror the
    document shape, but only `False` is valid; this is enforced at runtime.
    """

    kind: Literal["inline"] = field(default="inline", init=False)
    must_understand: bool = False
    metadata: dict[str, ZarrV3ArrayMetadata | ZarrV3GroupMetadata]

    def __post_init__(self) -> None:
        if self.must_understand is not False:
            raise MetadataValidationError(
                [
                    ValidationProblem(
                        ("must_understand",),
                        f"Invalid value for 'must_understand'. Expected False. "
                        f"Got {self.must_understand!r}.",
                        "invalid_value",
                    )
                ]
            )

    def to_json(self) -> ZarrV3ConsolidatedMetadataJSON:
        # `must_understand` is emitted as the literal False: the field is typed
        # permissively as `bool`, but `__post_init__` guarantees the value.
        return {
            "kind": self.kind,
            "must_understand": False,
            "metadata": {key: node.to_json() for key, node in self.metadata.items()},
        }

    @classmethod
    def from_json(cls, data: object) -> ZarrV3ConsolidatedMetadata:
        normalized = arrays_to_tuples(data)
        problems = validate_consolidated_metadata_v3(normalized)
        if problems:
            raise MetadataValidationError(problems)
        env = cast("Mapping[str, object]", normalized)
        entries: dict[str, ZarrV3ArrayMetadata | ZarrV3GroupMetadata] = {}
        for key, entry in cast("Mapping[str, object]", env["metadata"]).items():
            node_type = cast("Mapping[str, object]", entry).get("node_type")
            if node_type == "array":
                entries[key] = ZarrV3ArrayMetadata.from_json(entry)
            else:
                entries[key] = ZarrV3GroupMetadata.from_json(entry)
        return cls(metadata=entries)


class ZarrV2GroupMetadataPartial(TypedDict, total=False):
    """
    Partial form of the constructor-settable fields of `ZarrV2GroupMetadata`.

    Every key is optional and typed with the model's own value types, so it
    describes valid keyword arguments to `ZarrV2GroupMetadata.update` and
    `create_default`. The `init=False` field `zarr_format` is intentionally
    excluded, since it cannot be passed to `dataclasses.replace`.

    Drift between this type and the model's settable fields is prevented by
    `tests/model/test_group.py::test_group_partial_keys_match_settable_model_fields`.
    """

    attributes: dict[str, JSONValue] | UNSET


@dataclass(frozen=True, slots=True, kw_only=True)
class ZarrV2GroupMetadata:
    """In-memory model of a v2 group metadata document.

    A canonical, lossless representation of the `.zgroup` content plus the
    sibling `.zattrs` attributes, folded into a single in-memory value
    (mirroring the merged `ZarrV2GroupMetadataJSON` document form). `attributes` is
    `UNSET` when no `.zattrs` file (or merged `attributes` key) exists —
    distinct from an explicit empty `.zattrs`, which is `{}` and round-trips
    as a file.
    """

    zarr_format: Literal[2] = field(default=2, init=False)
    attributes: dict[str, JSONValue] | UNSET

    @classmethod
    def create_default(cls, **overrides: Unpack[ZarrV2GroupMetadataPartial]) -> ZarrV2GroupMetadata:
        """
        Create a default (empty) v2 group metadata model, with optional overrides.

        The default is a structurally-valid group with no attributes — the group
        analog of `list()` returning `[]`. Any field can be overridden by keyword
        (the same fields accepted by `update`).
        """
        default = cls(attributes=UNSET)
        return default.update(**overrides)

    def update(self, **kwargs: Unpack[ZarrV2GroupMetadataPartial]) -> ZarrV2GroupMetadata:
        """
        Return a new `ZarrV2GroupMetadata` with the given fields updated.

        Only the constructor-settable fields listed in
        `ZarrV2GroupMetadataPartial` can be updated; the fixed `zarr_format`
        is rejected at the type level. Each given field fully replaces its
        previous value.
        """
        return dataclasses.replace(self, **kwargs)

    def to_json(self) -> ZarrV2GroupMetadataJSON:
        """Return the merged in-memory document form.

        `attributes` is included when set (even empty). This is not the
        on-disk `.zgroup` content: a conforming `.zgroup` must exclude
        `attributes` (they live in the sibling `.zattrs` file). Use
        `to_key_value` to produce the spec-conforming split for storage.
        """
        # to_json output shares no mutable state with the model.
        out: ZarrV2GroupMetadataJSON = {"zarr_format": self.zarr_format}
        if self.attributes is not UNSET:
            out["attributes"] = copy.deepcopy(self.attributes)
        return out

    @classmethod
    def from_json(cls, data: object) -> ZarrV2GroupMetadata:
        parsed = parse_group_metadata_v2(arrays_to_tuples(data))
        return cls(attributes=(dict(parsed["attributes"]) if "attributes" in parsed else UNSET))

    @classmethod
    def from_key_value(cls, mapping: Mapping[str, bytes]) -> ZarrV2GroupMetadata:
        zgroup_raw = cast("object", load_store_json(mapping, GROUP_METADATA_STORE_KEY_V2))
        if not isinstance(zgroup_raw, Mapping):
            return cls.from_json(zgroup_raw)
        zgroup = cast("Mapping[str, object]", zgroup_raw)
        if "attributes" in zgroup:
            raise MetadataValidationError(
                [
                    ValidationProblem(
                        ("attributes",),
                        "unexpected document member",
                        "invalid_value",
                    )
                ]
            )
        if ATTRIBUTES_STORE_KEY_V2 in mapping:
            zattrs = cast("object", load_store_json(mapping, ATTRIBUTES_STORE_KEY_V2))
            return cls.from_json({**zgroup, "attributes": zattrs})
        return cls.from_json(zgroup)

    def to_key_value(
        self, *, indent: int | str | None = None
    ) -> Mapping[ZarrV2GroupMetadataStoreKey | ZarrV2AttributesStoreKey, bytes]:
        # Attributes live only in the sibling `.zattrs` file; the `.zgroup`
        # document must exclude them. The `.zattrs` key is present exactly
        # when attributes are set (even empty) — UNSET emits no file.
        zgroup = {k: v for k, v in self.to_json().items() if k != "attributes"}
        out: dict[ZarrV2GroupMetadataStoreKey | ZarrV2AttributesStoreKey, bytes] = {
            GROUP_METADATA_STORE_KEY_V2: dump_store_json(zgroup, indent=indent)
        }
        if self.attributes is not UNSET:
            out[ATTRIBUTES_STORE_KEY_V2] = dump_store_json(self.attributes, indent=indent)
        return out


@dataclass(frozen=True, slots=True, kw_only=True)
class ZarrV2ConsolidatedMetadata:
    """In-memory model of a v2 `.zmetadata` document.

    The `metadata` map holds the flat file-keyed entries (`"path/.zarray"`,
    `"path/.zattrs"`, ...) verbatim, preserving the normalized JSON tree.
    Entries are deliberately NOT merged into per-node models: which nodes had
    a `.zattrs` file at all is information the canonical representation must
    keep. Interpreting entries into node models is consumer work.
    """

    zarr_consolidated_format: Literal[1] = field(default=1, init=False)
    metadata: dict[str, JSONValue]

    def to_json(self) -> dict[str, JSONValue]:
        # to_json output shares no mutable state with the model.
        return {
            "zarr_consolidated_format": self.zarr_consolidated_format,
            "metadata": copy.deepcopy(self.metadata),
        }

    @classmethod
    def from_json(cls, data: object) -> ZarrV2ConsolidatedMetadata:
        normalized = arrays_to_tuples(data)
        if not isinstance(normalized, Mapping):
            raise MetadataValidationError(
                [ValidationProblem((), "expected a mapping", "invalid_type")]
            )
        doc = cast("Mapping[str, object]", normalized)
        problems: list[ValidationProblem] = [
            ValidationProblem((key,), "missing required key", "missing_key")
            for key in ("zarr_consolidated_format", "metadata")
            if key not in doc
        ]
        problems.extend(
            ValidationProblem((key,), "unexpected document member", "invalid_value")
            for key in doc.keys() - {"zarr_consolidated_format", "metadata"}
        )
        if "zarr_consolidated_format" in doc and (
            not isinstance(doc["zarr_consolidated_format"], int)
            or isinstance(doc["zarr_consolidated_format"], bool)
            or doc["zarr_consolidated_format"] != 1
        ):
            problems.append(
                ValidationProblem(
                    ("zarr_consolidated_format",),
                    f"expected 1, got {doc['zarr_consolidated_format']!r}",
                    "invalid_value",
                )
            )
        if "metadata" in doc:
            entries = doc["metadata"]
            if not isinstance(entries, Mapping) or not all(
                isinstance(k, str) for k in cast("Mapping[object, object]", entries)
            ):
                problems.append(
                    ValidationProblem(
                        ("metadata",), "expected a mapping with string keys", "invalid_type"
                    )
                )
            else:
                for key, value in cast("Mapping[str, object]", entries).items():
                    problems.extend(
                        ValidationProblem(
                            ("metadata", key, *problem.loc), problem.message, problem.kind
                        )
                        for problem in validate_json(value)
                    )
        if problems:
            raise MetadataValidationError(problems)
        entries_tupled = cast(
            "dict[str, JSONValue]",
            arrays_to_tuples(dict(cast("Mapping[str, object]", doc["metadata"]))),
        )
        return cls(metadata=entries_tupled)

    @classmethod
    def from_key_value(cls, mapping: Mapping[str, bytes]) -> ZarrV2ConsolidatedMetadata:
        return cls.from_json(load_store_json(mapping, CONSOLIDATED_METADATA_STORE_KEY_V2))

    def to_key_value(
        self, *, indent: int | str | None = None
    ) -> Mapping[ZarrV2ConsolidatedMetadataStoreKey, bytes]:
        return {CONSOLIDATED_METADATA_STORE_KEY_V2: dump_store_json(self.to_json(), indent=indent)}
