from __future__ import annotations

import copy
from datetime import datetime, timezone
from typing import Any

import tiktoken
from pydantic import BaseModel, Field


class HotEntry(BaseModel):
    """Single hot-layer key/value entry with usage metadata."""

    key: str
    value: Any
    created_at: datetime
    last_accessed: datetime
    access_count: int = 0
    token_estimate: int
    tags: list[str] = Field(default_factory=list)


class HotLayer:
    """In-context mutable memory with staleness tracking and token accounting."""

    def __init__(self, agent_id: str, ttl_seconds: int = 300, max_tokens: int = 8000) -> None:
        self.agent_id = agent_id
        self.ttl_seconds = ttl_seconds
        self.max_tokens = max_tokens
        self._entries: dict[str, HotEntry] = {}
        self._encoding = tiktoken.get_encoding("cl100k_base")

    def set(self, key: str, value: Any, tags: list[str] | None = None) -> HotEntry:
        """Store or replace a value in hot memory and compute token estimate."""
        now = datetime.now(timezone.utc)
        tok = len(self._encoding.encode(str(value)))
        prior = self._entries.get(key)
        entry = HotEntry(
            key=key,
            value=value,
            created_at=prior.created_at if prior else now,
            last_accessed=now,
            access_count=(prior.access_count + 1) if prior else 1,
            token_estimate=tok,
            tags=list(tags or []),
        )
        self._entries[key] = entry
        return entry

    def get(self, key: str) -> Any | None:
        """Get a value and update access metadata."""
        entry = self._entries.get(key)
        if entry is None:
            return None
        entry.last_accessed = datetime.now(timezone.utc)
        entry.access_count += 1
        return entry.value

    def delete(self, key: str) -> bool:
        """Delete a key from hot memory."""
        return self._entries.pop(key, None) is not None

    def get_stale_keys(self, threshold_seconds: float | None = None) -> list[str]:
        """Return keys whose last access exceeds configured or provided threshold."""
        now = datetime.now(timezone.utc)
        threshold = threshold_seconds if threshold_seconds is not None else float(self.ttl_seconds)
        stale: list[str] = []
        for key, entry in self._entries.items():
            age = (now - entry.last_accessed).total_seconds()
            if age > threshold:
                stale.append(key)
        return stale

    def total_tokens(self) -> int:
        """Compute token footprint of all hot entries."""
        return sum(v.token_estimate for v in self._entries.values())

    def evict_stale(self) -> list[HotEntry]:
        """Evict stale entries and return them for warm migration."""
        stale_keys = self.get_stale_keys()
        evicted: list[HotEntry] = []
        for key in stale_keys:
            entry = self._entries.pop(key, None)
            if entry is not None:
                evicted.append(entry)
        return evicted

    def to_context_string(self) -> str:
        """Serialize all entries into deterministic context text."""
        lines: list[str] = [f"# HotLayer[{self.agent_id}]", ""]
        for key in sorted(self._entries):
            e = self._entries[key]
            lines.append(f"[{e.key}] tokens={e.token_estimate} tags={','.join(e.tags)}")
            lines.append(str(e.value))
            lines.append("")
        return "\n".join(lines).strip()

    def fork(self) -> HotLayer:
        """Create a deep-copy fork for branch copy-on-write semantics."""
        child = HotLayer(self.agent_id, ttl_seconds=self.ttl_seconds, max_tokens=self.max_tokens)
        child._entries = copy.deepcopy(self._entries)
        return child

    def merge_from(self, other: HotLayer, strategy: str = "latest_wins") -> None:
        """Merge another hot layer into this one with conflict strategy."""
        for key, other_entry in other._entries.items():
            mine = self._entries.get(key)
            if mine is None:
                self._entries[key] = copy.deepcopy(other_entry)
                continue
            if strategy == "latest_wins":
                if other_entry.last_accessed >= mine.last_accessed:
                    self._entries[key] = copy.deepcopy(other_entry)
            elif strategy == "other_wins":
                self._entries[key] = copy.deepcopy(other_entry)
            elif strategy == "union":
                suffix = 1
                new_key = f"{key}_other_{suffix}"
                while new_key in self._entries:
                    suffix += 1
                    new_key = f"{key}_other_{suffix}"
                cloned = copy.deepcopy(other_entry)
                cloned.key = new_key
                self._entries[new_key] = cloned
            else:
                raise ValueError(f"Unknown merge strategy: {strategy}")

    def to_dict(self) -> dict[str, dict[str, Any]]:
        """Return serializable mapping of key -> HotEntry dict."""
        return {k: v.model_dump() for k, v in self._entries.items()}

    def restore(self, data: dict[str, dict[str, Any]]) -> None:
        """Restore entries from serialized data."""
        self._entries = {k: HotEntry.model_validate(v) for k, v in data.items()}
