# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import contextlib
import threading
from collections.abc import Callable, Iterator
from typing import Any, Generic, Literal, TypeVar, overload
from uuid import UUID

from pydantic import Field, PrivateAttr, field_serializer, field_validator
from typing_extensions import override

from krons.core.types import Unset, UnsetType, is_unset
from krons.errors import ExistsError, NotFoundError
from krons.protocols import Containable, Deserializable, Serializable, implements
from krons.utils import extract_types, load_type_from_string, synchronized
from krons.utils.concurrency import Lock as AsyncLock

from .element import Element
from .progression import Progression

__all__ = ("Pile",)

T = TypeVar("T", bound=Element)


@implements(
    Containable,
    Serializable,
    Deserializable,
)
class Pile(Element, Generic[T]):
    """Thread-safe typed collection with rich query interface.

    A Pile is an ordered, type-validated container for Element subclasses.
    Maintains insertion order via Progression, supports concurrent access.

    Type-dispatched __getitem__:
        pile[uuid|str]     -> T (single item by ID)
        pile[int]          -> T (single item by index)
        pile[slice]        -> Pile[T] (range)
        pile[list|tuple]   -> Pile[T] (indices or UUIDs)
        pile[Progression]  -> Pile[T] (ordered subset)
        pile[callable]     -> Pile[T] (filter function)

    Example:
        pile = Pile[Node](items=[n1, n2], item_type=Node)
        pile.add(n3)
        filtered = pile[lambda x: x.metadata.get("active")]
    """

    _items: dict[UUID, T] = PrivateAttr(default_factory=dict)
    _progression: Progression = PrivateAttr(default_factory=Progression)
    _lock: threading.RLock = PrivateAttr(default_factory=threading.RLock)
    _async_lock: AsyncLock = PrivateAttr(default_factory=AsyncLock)

    @property
    def progression(self) -> Progression:
        """Read-only copy of progression order."""
        return Progression(
            order=list(self._progression.order), name=self._progression.name
        )

    item_type: set[type] | None = Field(
        default=None,
        frozen=True,
        description="Allowed types for validation (None = any Element subclass)",
    )
    strict_type: bool = Field(
        default=False,
        frozen=True,
        description="Enforce exact type match (disallow subclasses)",
    )

    @field_validator("item_type", mode="before")
    @classmethod
    def _normalize_item_type(cls, v: Any) -> set[type] | None:
        """Normalize item_type to set[type] from various input formats."""
        if v is None:
            return None
        if isinstance(v, list) and v and isinstance(v[0], str):
            return {load_type_from_string(type_str) for type_str in v}
        return extract_types(v)

    @override
    def __init__(
        self,
        items: list[T] | None = None,
        item_type: type[T] | set[type] | list[type] | None = None,
        order: list[UUID] | Progression | None = None,
        strict_type: bool = False,
        **kwargs,
    ):
        """Initialize Pile with optional items.

        Args:
            items: Initial items to add
            item_type: Type(s) for validation (type, set, list, or Union)
            order: Custom order (list of UUIDs or Progression)
            strict_type: Enforce exact type match (no subclasses)
            **kwargs: Element fields (id, created_at, metadata)

        Raises:
            NotFoundError: If order contains UUID not in items
            TypeError: If item type validation fails
        """
        super().__init__(
            **{"item_type": item_type, "strict_type": strict_type, **kwargs}
        )

        if items:
            for item in items:
                self.add(item)

        if order:
            order_list = list(order.order) if isinstance(order, Progression) else order
            for uid in order_list:
                if uid not in self._items:
                    raise NotFoundError(f"UUID {uid} in order not found in items")
            self._progression = Progression(order=order_list)

    @field_serializer("item_type")
    def _serialize_item_type(self, v: set[type] | None) -> list[str] | None:
        """Serialize item_type to list of fully-qualified type names."""
        if v is None:
            return None
        return [f"{t.__module__}.{t.__name__}" for t in v]

    @override
    def to_dict(
        self,
        mode: Literal["python", "json", "db"] = "python",
        created_at_format: (
            Literal["datetime", "isoformat", "timestamp"] | UnsetType
        ) = Unset,
        meta_key: str | UnsetType = Unset,
        item_meta_key: str | UnsetType = Unset,
        item_created_at_format: (
            Literal["datetime", "isoformat", "timestamp"] | UnsetType
        ) = Unset,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Serialize pile with items in progression order.

        Args:
            mode: Serialization mode (python/json/db)
            created_at_format: Timestamp format for Pile itself
            meta_key: Rename Pile metadata field
            item_meta_key: Metadata key name for items
            item_created_at_format: Timestamp format for items
            **kwargs: Passed to model_dump()

        Returns:
            Dict with pile fields and serialized items in order
        """
        data = super().to_dict(
            mode=mode, created_at_format=created_at_format, meta_key=meta_key, **kwargs
        )

        actual_meta_key = (
            meta_key
            if not is_unset(meta_key)
            else ("node_metadata" if mode == "db" else "metadata")
        )

        if self._progression.name and actual_meta_key in data:
            data[actual_meta_key]["progression_name"] = self._progression.name

        item_mode = "python" if mode == "python" else "json"
        data["items"] = [
            i.to_dict(
                mode=item_mode,
                meta_key=item_meta_key,
                created_at_format=item_created_at_format,
            )
            for i in self
        ]

        return data

    # ==================== Core Operations ====================

    @synchronized
    def add(self, item: T) -> None:
        """Add item to pile. Raises ExistsError if duplicate, TypeError if invalid type."""
        self._validate_type(item)

        if item.id in self._items:
            raise ExistsError(f"Item {item.id} already exists in pile")

        self._items[item.id] = item
        self._progression.append(item.id)

    @synchronized
    def remove(self, item_id: UUID | str | Element) -> T:
        """Remove and return item. Raises NotFoundError if not found."""
        uid = self._coerce_id(item_id)

        try:
            item = self._items.pop(uid)
        except KeyError:
            raise NotFoundError(f"Item {uid} not found in pile") from None

        self._progression.remove(uid)
        return item

    @synchronized
    def pop(self, item_id: UUID | str | Element, default: Any = ...) -> T | Any:
        """Remove and return item, or default if not found. Raises NotFoundError if no default."""
        uid = self._coerce_id(item_id)

        try:
            item = self._items.pop(uid)
            self._progression.remove(uid)
            return item
        except KeyError:
            if default is ...:
                raise NotFoundError(f"Item {uid} not found in pile") from None
            return default

    @synchronized
    def get(self, item_id: UUID | str | Element, default: Any = ...) -> T | Any:
        """Get item by ID, or default if not found. Raises NotFoundError if no default."""
        uid = self._coerce_id(item_id)

        try:
            return self._items[uid]
        except KeyError:
            if default is ...:
                raise NotFoundError(f"Item {uid} not found in pile") from None
            return default

    @synchronized
    def update(self, item: T) -> None:
        """Update existing item in-place. Raises NotFoundError if not found, TypeError if invalid."""
        self._validate_type(item)

        if item.id not in self._items:
            raise NotFoundError(f"Item {item.id} not found in pile")

        self._items[item.id] = item

    @synchronized
    def clear(self) -> None:
        """Remove all items from pile."""
        self._items.clear()
        self._progression.clear()

    # ==================== Set-like Operations ====================

    @synchronized
    def include(self, item: T) -> bool:
        """Idempotent add: returns True if item is in pile after call, False on validation failure."""
        if item.id in self._items:
            return True
        try:
            self._validate_type(item)
            self._items[item.id] = item
            self._progression.append(item.id)
            return True
        except Exception:
            return False

    @synchronized
    def exclude(self, item: UUID | str | Element) -> bool:
        """Idempotent remove: returns True if item is not in pile after call, False on coercion failure."""
        try:
            uid = self._coerce_id(item)
        except Exception:
            return False
        if uid not in self._items:
            return True
        self._items.pop(uid, None)
        try:
            self._progression.remove(uid)
        except ValueError:
            pass
        return True

    # ==================== Rich __getitem__ (Type Dispatch) ====================

    @overload
    def __getitem__(self, key: UUID | str) -> T: ...
    @overload
    def __getitem__(self, key: Progression) -> Pile[T]: ...
    @overload
    def __getitem__(self, key: int) -> T: ...
    @overload
    def __getitem__(self, key: slice) -> Pile[T]: ...
    @overload
    def __getitem__(self, key: list[int] | tuple[int, ...]) -> Pile[T]: ...
    @overload
    def __getitem__(self, key: list[UUID] | tuple[UUID, ...]) -> Pile[T]: ...
    @overload
    def __getitem__(self, key: Callable[[T], bool]) -> Pile[T]: ...

    def __getitem__(self, key: Any) -> T | Pile[T]:
        """Type-dispatched query: UUID/str/int -> T; slice/list/tuple/Progression/callable -> Pile[T]."""
        if isinstance(key, (UUID, str)):
            return self.get(key)
        elif isinstance(key, int):
            return self._get_by_index(key)
        elif isinstance(key, Progression):
            return self._filter_by_progression(key)
        elif isinstance(key, slice):
            return self._get_by_slice(key)
        elif isinstance(key, (list, tuple)):
            return self._get_by_list(key)
        elif callable(key):
            return self._filter_by_function(key)
        else:
            raise TypeError(
                f"Invalid key type: {type(key)}. Expected UUID, str, int, slice, list, tuple, Progression, or callable"
            )

    def _filter_by_progression(self, prog: Progression) -> Pile[T]:
        """Return new Pile with items in progression order. Raises NotFoundError if UUID missing."""
        if any(uid not in self._items for uid in prog):
            raise NotFoundError("Some items from progression not found in pile")

        return Pile(
            items=[self._items[uid] for uid in prog],
            item_type=self.item_type,
            strict_type=self.strict_type,
        )

    @synchronized
    def _get_by_index(self, index: int) -> T:
        """Get item by position in progression order."""
        uid: UUID = self._progression[index]
        return self._items[uid]

    @synchronized
    def _get_by_slice(self, s: slice) -> Pile[T]:
        """Return new Pile with items from slice range."""
        uids: list[UUID] = self._progression[s]
        return Pile(
            items=[self._items[uid] for uid in uids],
            item_type=self.item_type,
            strict_type=self.strict_type,
        )

    @synchronized
    def _get_by_list(self, keys: list | tuple) -> Pile[T]:
        """Return new Pile from list of indices or UUIDs. No mixing allowed."""
        if not keys:
            raise ValueError("Cannot get items with empty list/tuple")

        first = keys[0]
        if isinstance(first, int):
            if not all(isinstance(k, int) for k in keys):
                raise TypeError("Cannot mix int and UUID in list/tuple indexing")
            items = [self._get_by_index(idx) for idx in keys]
        elif isinstance(first, (UUID, str)):
            if not all(isinstance(k, (UUID, str)) for k in keys):
                raise TypeError("Cannot mix int and UUID in list/tuple indexing")
            items = [self.get(uid) for uid in keys]
        else:
            raise TypeError(
                f"list/tuple must contain only int or UUID, got {type(first)}"
            )

        return Pile(
            items=items,
            item_type=self.item_type,
            strict_type=self.strict_type,
        )

    def _filter_by_function(self, func: Callable[[T], bool]) -> Pile[T]:
        """Return new Pile with items matching filter function."""
        return Pile(
            items=[item for item in self if func(item)],
            item_type=self.item_type,
            strict_type=self.strict_type,
        )

    def filter_by_type(self, item_type: type[T] | set[type] | list[type]) -> Pile[T]:
        """Return new Pile with items matching specified type(s).

        Args:
            item_type: Type(s) to filter by

        Returns:
            New Pile containing only items of requested type(s)

        Raises:
            TypeError: If requested type incompatible with pile's item_type
            NotFoundError: If no items match
        """
        types_to_filter = extract_types(item_type)

        if self.item_type is not None:
            if self.strict_type:
                invalid_types = types_to_filter - self.item_type
                if invalid_types:
                    raise TypeError(
                        f"Types {invalid_types} not allowed in pile (allowed: {self.item_type})"
                    )
            else:
                for t in types_to_filter:
                    is_compatible = any(
                        issubclass(t, allowed) or issubclass(allowed, t)
                        for allowed in self.item_type
                    )
                    if not is_compatible:
                        raise TypeError(
                            f"Type {t} not compatible with allowed types {self.item_type}"
                        )

        filtered_items = [
            item for item in self if any(isinstance(item, t) for t in types_to_filter)
        ]

        if not filtered_items:
            raise NotFoundError(f"No items of type(s) {types_to_filter} found in pile")

        return Pile(
            items=filtered_items,
            item_type=self.item_type,
            strict_type=self.strict_type,
        )

    # ==================== Context Managers ====================

    async def __aenter__(self) -> Pile[T]:
        """Acquire async lock for context manager."""
        await self._async_lock.acquire()
        return self

    async def __aexit__(self, _exc_type, _exc_val, _exc_tb) -> None:
        """Release async lock."""
        self._async_lock.release()

    # ==================== Query Operations ====================

    @synchronized
    def __contains__(self, item: UUID | str | Element) -> bool:
        """Check if item exists in pile by ID."""
        with contextlib.suppress(Exception):
            uid = self._coerce_id(item)
            return uid in self._items
        return False

    @synchronized
    def __len__(self) -> int:
        return len(self._items)

    def __bool__(self) -> bool:
        return len(self._items) > 0

    @synchronized
    def __iter__(self) -> Iterator[T]:  # type: ignore[override]
        """Iterate items in progression order. LSP override: yields T, not field tuples."""
        for uid in self._progression:
            yield self._items[uid]

    def keys(self) -> Iterator[UUID]:
        """Iterate UUIDs in progression order."""
        return iter(self._progression)

    def items(self) -> Iterator[tuple[UUID, T]]:
        """Iterate (UUID, item) pairs in progression order."""
        for i in self:
            yield (i.id, i)

    def __list__(self) -> list[T]:
        """Return items as list in progression order."""
        return [i for i in self]

    def is_empty(self) -> bool:
        """Check if pile contains no items."""
        return len(self._items) == 0

    # ==================== Validation ====================

    def _validate_type(self, item: T) -> None:
        """Validate item against pile's type constraints. Raises TypeError on failure."""
        if not isinstance(item, Element):
            raise TypeError(f"Item must be Element subclass, got {type(item)}")

        if self.item_type is not None:
            item_type_actual = type(item)

            if self.strict_type:
                if item_type_actual not in self.item_type:
                    raise TypeError(
                        f"Item type {item_type_actual} not in allowed types {self.item_type} "
                        "(strict_type=True, no subclasses allowed)"
                    )
            else:
                if not any(issubclass(item_type_actual, t) for t in self.item_type):
                    raise TypeError(
                        f"Item type {item_type_actual} is not a subclass of any allowed type {self.item_type}"
                    )

    # ==================== Deserialization ====================

    @classmethod
    @override
    def from_dict(
        cls,
        data: dict[str, Any],
        meta_key: str | UnsetType = Unset,
        item_meta_key: str | UnsetType = Unset,
        **kwargs: Any,
    ) -> Pile[T]:
        """Deserialize Pile from dict. Validates all item types before deserializing.

        Args:
            data: Serialized pile data
            meta_key: Restore metadata from this key (db mode compatibility)
            item_meta_key: Metadata key for item deserialization
            **kwargs: Additional arguments (e.g., item_type override)

        Returns:
            Reconstructed Pile with all items

        Raises:
            TypeError: If any item type incompatible with pile's constraints
        """
        from .element import Element

        data = data.copy()

        if not is_unset(meta_key) and meta_key in data:
            data["metadata"] = data.pop(meta_key)
        elif "node_metadata" in data and "metadata" not in data:
            data["metadata"] = data.pop("node_metadata")
        data.pop("node_metadata", None)

        item_type_data = data.get("item_type") or kwargs.get("item_type")
        strict_type = data.get("strict_type", False)

        items_data = data.get("items", [])
        if item_type_data is not None and items_data:
            if (
                isinstance(item_type_data, list)
                and item_type_data
                and isinstance(item_type_data[0], str)
            ):
                allowed_types = {
                    load_type_from_string(type_str) for type_str in item_type_data
                }
            else:
                allowed_types = extract_types(item_type_data)

            for item_dict in items_data:
                kron_class = item_dict.get("metadata", {}).get("kron_class")
                if kron_class:
                    try:
                        item_type_actual = load_type_from_string(kron_class)
                    except ValueError:
                        continue

                    if strict_type:
                        if item_type_actual not in allowed_types:
                            raise TypeError(
                                f"Item type {kron_class} not in allowed types {allowed_types} "
                                "(strict_type=True)"
                            )
                    else:
                        if not any(
                            issubclass(item_type_actual, t) for t in allowed_types
                        ):
                            raise TypeError(
                                f"Item type {kron_class} is not a subclass of any allowed type {allowed_types}"
                            )

        pile_data = data.copy()
        pile_data.pop("items", None)
        pile_data.pop("item_type", None)
        pile_data.pop("strict_type", None)
        pile = cls(item_type=item_type_data, strict_type=strict_type, **pile_data)

        metadata = data.get("metadata", {})
        progression_name = metadata.get("progression_name")
        if progression_name:
            pile._progression.name = progression_name

        for item_dict in items_data:
            item = Element.from_dict(item_dict, meta_key=item_meta_key)
            pile.add(item)  # type: ignore[arg-type]

        return pile

    def __repr__(self) -> str:
        return f"Pile(len={len(self)})"
