from __future__ import annotations

import copy
import json
from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, Field


class DiffRecord(BaseModel):
    """Typed state diff artifact used for forward and reverse replay."""

    diff_id: str = Field(default_factory=lambda: str(uuid4()))
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    layer: str
    agent_id: str
    added: dict[str, Any]
    removed: dict[str, Any]
    modified: dict[str, tuple[Any, Any]]
    node_id: str | None = None


class DiffEngine:
    """Compute, apply, invert, and serialize state diffs."""

    def diff(
        self,
        before: dict[str, Any],
        after: dict[str, Any],
        layer: str,
        agent_id: str,
        node_id: str | None = None,
    ) -> DiffRecord:
        """Compute added/removed/modified maps between two snapshots."""
        b = copy.deepcopy(before)
        a = copy.deepcopy(after)

        added = {k: v for k, v in a.items() if k not in b}
        removed = {k: v for k, v in b.items() if k not in a}
        modified: dict[str, tuple[Any, Any]] = {}
        for key in set(b).intersection(a):
            if b[key] != a[key]:
                modified[key] = (b[key], a[key])

        return DiffRecord(
            layer=layer,
            agent_id=agent_id,
            added=added,
            removed=removed,
            modified=modified,
            node_id=node_id,
        )

    def apply(self, base: dict[str, Any], diff: DiffRecord) -> dict[str, Any]:
        """Apply a diff forward and return new state."""
        out = copy.deepcopy(base)
        for key in diff.removed:
            out.pop(key, None)
        for key, value in diff.added.items():
            out[key] = copy.deepcopy(value)
        for key, (_, new_val) in diff.modified.items():
            out[key] = copy.deepcopy(new_val)
        return out

    def invert(self, diff: DiffRecord) -> DiffRecord:
        """Create reverse diff that undoes this diff."""
        return DiffRecord(
            layer=diff.layer,
            agent_id=diff.agent_id,
            added=copy.deepcopy(diff.removed),
            removed=copy.deepcopy(diff.added),
            modified={k: (new_val, old_val) for k, (old_val, new_val) in diff.modified.items()},
            node_id=diff.node_id,
        )

    def serialize(self, diff: DiffRecord) -> str:
        """Serialize diff as JSON string."""
        return diff.model_dump_json()

    def deserialize(self, data: str) -> DiffRecord:
        """Deserialize JSON string into diff record."""
        payload = json.loads(data)
        return DiffRecord.model_validate(payload)
