"""Deterministic calculation graph helpers for reactive orchestration."""

from __future__ import annotations

import hashlib
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum


def _normalize(value: object) -> str:
    text = str(value).strip()
    if not text:
        raise ValueError("Node identifier must be non-empty.")
    return text


@dataclass(frozen=True, order=True, slots=True)
class NodeId:
    value: str

    def __post_init__(self) -> None:
        object.__setattr__(self, "value", _normalize(self.value))

    @classmethod
    def parse(cls, value: object) -> "NodeId":
        return cls(_normalize(value))

    @classmethod
    def from_string(cls, value: str) -> "NodeId":
        return cls.parse(value)

    def as_str(self) -> str:
        return self.value

    def __str__(self) -> str:  # pragma: no cover - trivial
        return self.value


@dataclass(frozen=True, slots=True)
class NodeValue:
    value: object | None = None
    revision: int = 0
    updated_at: datetime | None = None
    source: str | None = None
    stale: bool = False


class ShardStrategy(str, Enum):
    MODULO = "MODULO"
    CONSISTENT_HASH = "CONSISTENT_HASH"


@dataclass(frozen=True, slots=True)
class ShardAssignment:
    node_id: NodeId
    shard_key: str
    shard_index: int
    owner: str | None = None


@dataclass(frozen=True, slots=True)
class ShardConfig:
    shard_count: int = 1
    owners: tuple[str, ...] = ()
    strategy: ShardStrategy = ShardStrategy.CONSISTENT_HASH
    virtual_nodes: int = 16

    def __post_init__(self) -> None:
        if self.shard_count <= 0:
            raise ValueError("shard_count must be positive.")
        if self.virtual_nodes <= 0:
            raise ValueError("virtual_nodes must be positive.")
        object.__setattr__(self, "owners", tuple(item.strip() for item in self.owners if item.strip()))


@dataclass(slots=True)
class _GraphNode:
    node_id: NodeId
    dependencies: set[NodeId] = field(default_factory=set)
    dependents: set[NodeId] = field(default_factory=set)
    cached_value: NodeValue | None = None
    dirty: bool = False
    shard_key: str | None = None
    metadata: dict[str, object] = field(default_factory=dict)


def _hash_to_int(value: str) -> int:
    return int(hashlib.sha256(value.encode("utf-8")).hexdigest(), 16)


@dataclass(slots=True)
class CalculationGraph:
    _nodes: dict[NodeId, _GraphNode] = field(default_factory=dict)

    def add_node(
        self,
        node_id: NodeId | str,
        *,
        value: object | None = None,
        shard_key: str | None = None,
        metadata: dict[str, object] | None = None,
    ) -> NodeId:
        resolved = node_id if isinstance(node_id, NodeId) else NodeId.parse(node_id)
        node = self._nodes.get(resolved)
        if node is None:
            cached_value = None
            if value is not None:
                cached_value = NodeValue(value=value, revision=1, updated_at=datetime.now(UTC))
            self._nodes[resolved] = _GraphNode(
                node_id=resolved,
                cached_value=cached_value,
                shard_key=shard_key,
                metadata=dict(metadata or {}),
            )
            return resolved
        if shard_key is not None:
            node.shard_key = shard_key
        if metadata is not None:
            node.metadata.update(metadata)
        if value is not None and node.cached_value is None:
            node.cached_value = NodeValue(value=value, revision=1, updated_at=datetime.now(UTC))
        return resolved

    def has_node(self, node_id: NodeId | str) -> bool:
        resolved = node_id if isinstance(node_id, NodeId) else NodeId.parse(node_id)
        return resolved in self._nodes

    def add_dependency(self, node_id: NodeId | str, depends_on: NodeId | str) -> None:
        child = self.add_node(node_id)
        parent = self.add_node(depends_on)
        self._nodes[child].dependencies.add(parent)
        self._nodes[parent].dependents.add(child)

    def dependencies(self, node_id: NodeId | str) -> tuple[NodeId, ...]:
        resolved = self._require_node(node_id)
        return tuple(sorted(self._nodes[resolved].dependencies))

    def dependents(self, node_id: NodeId | str) -> tuple[NodeId, ...]:
        resolved = self._require_node(node_id)
        return tuple(sorted(self._nodes[resolved].dependents))

    def mark_dirty(
        self,
        node_id: NodeId | str,
        *,
        propagate: bool = True,
        include_self: bool = True,
    ) -> tuple[NodeId, ...]:
        resolved = self.add_node(node_id)
        marked: set[NodeId] = set()
        stack = [resolved]
        while stack:
            current = stack.pop()
            if current in marked and (current != resolved or include_self):
                continue
            if current != resolved or include_self:
                self._nodes[current].dirty = True
                marked.add(current)
            if propagate:
                stack.extend(sorted(self._nodes[current].dependents, reverse=True))
        return tuple(sorted(marked))

    def mark_dependents_dirty(self, node_id: NodeId | str) -> tuple[NodeId, ...]:
        return self.mark_dirty(node_id, include_self=False)

    def clear_dirty(self, node_id: NodeId | str) -> None:
        resolved = self._require_node(node_id)
        self._nodes[resolved].dirty = False

    def query_dirty(self) -> tuple[NodeId, ...]:
        return tuple(sorted(node_id for node_id, node in self._nodes.items() if node.dirty))

    def update_node_value(
        self,
        node_id: NodeId | str,
        value: object,
        *,
        revision: int | None = None,
        source: str | None = None,
        mark_clean: bool = True,
    ) -> NodeValue:
        resolved = self.add_node(node_id)
        current = self._nodes[resolved].cached_value
        next_revision = 1 if current is None else current.revision + 1
        if revision is not None:
            next_revision = revision
        cached = NodeValue(
            value=value,
            revision=next_revision,
            updated_at=datetime.now(UTC),
            source=source,
            stale=False,
        )
        self._nodes[resolved].cached_value = cached
        if mark_clean:
            self._nodes[resolved].dirty = False
        return cached

    def get_node_value(self, node_id: NodeId | str) -> NodeValue | None:
        resolved = self._require_node(node_id)
        return self._nodes[resolved].cached_value

    def revision(self, node_id: NodeId | str) -> int:
        cached = self.get_node_value(node_id)
        return 0 if cached is None else cached.revision

    def node_metadata(self, node_id: NodeId | str) -> dict[str, object]:
        resolved = self._require_node(node_id)
        return dict(self._nodes[resolved].metadata)

    def shard_assignment(self, node_id: NodeId | str, config: ShardConfig) -> ShardAssignment:
        resolved = self.add_node(node_id)
        node = self._nodes[resolved]
        shard_key = node.shard_key or resolved.as_str()
        shard_index = _hash_to_int(shard_key) % config.shard_count
        owner = None
        if config.owners:
            if config.strategy is ShardStrategy.MODULO:
                owner = config.owners[shard_index % len(config.owners)]
            else:
                owner = self._consistent_hash_owner(shard_key, config)
        return ShardAssignment(
            node_id=resolved,
            shard_key=shard_key,
            shard_index=shard_index,
            owner=owner,
        )

    def _consistent_hash_owner(self, key: str, config: ShardConfig) -> str:
        ring: list[tuple[int, str]] = []
        for owner in config.owners:
            for vnode in range(config.virtual_nodes):
                ring.append((_hash_to_int(f"{owner}:{vnode}"), owner))
        ring.sort(key=lambda item: item[0])
        target = _hash_to_int(key)
        for point, owner in ring:
            if target <= point:
                return owner
        return ring[0][1]

    def _require_node(self, node_id: NodeId | str) -> NodeId:
        resolved = node_id if isinstance(node_id, NodeId) else NodeId.parse(node_id)
        if resolved not in self._nodes:
            raise KeyError(f"Node {resolved.as_str()!r} is not registered.")
        return resolved


__all__ = [
    "CalculationGraph",
    "NodeId",
    "NodeValue",
    "ShardAssignment",
    "ShardConfig",
    "ShardStrategy",
]
