from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from nexus.core.events.bus import EventBus
from nexus.core.events.types import StateMutationEvent
from nexus.core.state.cold_layer import ColdLayer
from nexus.core.state.hot_layer import HotLayer
from nexus.core.state.warm_layer import WarmEntry, WarmLayer


class FluidMesh:
    """Three-tier state mesh orchestrator for per-agent memory management."""

    def __init__(
        self,
        bus: EventBus,
        hot_ttl: int = 300,
        hot_max_tokens: int = 8000,
        warm_max_entries: int = 500,
        cold_persist_dir: str = ".nexus/cold_store",
    ) -> None:
        self.bus = bus
        self.hot_ttl = hot_ttl
        self.hot_max_tokens = hot_max_tokens
        self.warm_max_entries = warm_max_entries
        self._hot: dict[str, HotLayer] = {}
        self._warm: dict[str, WarmLayer] = {}
        self._branches: dict[str, HotLayer] = {}
        self.cold = ColdLayer(persist_directory=cold_persist_dir)

    def get_hot(self, agent_id: str) -> HotLayer:
        """Get or create a hot layer for an agent."""
        if agent_id not in self._hot:
            self._hot[agent_id] = HotLayer(agent_id, ttl_seconds=self.hot_ttl, max_tokens=self.hot_max_tokens)
        return self._hot[agent_id]

    def get_warm(self, agent_id: str) -> WarmLayer:
        """Get or create a warm layer for an agent."""
        if agent_id not in self._warm:
            self._warm[agent_id] = WarmLayer(agent_id, max_entries=self.warm_max_entries)
        return self._warm[agent_id]

    async def set(self, agent_id: str, key: str, value: Any, tags: list[str] | None = None) -> None:
        """Write to hot layer and emit state mutation event."""
        hot = self.get_hot(agent_id)
        old_value = hot.get(key)
        hot.set(key, value, tags=tags)
        await self.bus.publish_and_wait(
            StateMutationEvent(
                source_id=agent_id,
                layer="hot",
                key=key,
                old_value=old_value,
                new_value=value,
            )
        )

    def get(self, agent_id: str, key: str) -> Any | None:
        """Lookup key in hot layer and fallback to warm recall via query."""
        hot = self.get_hot(agent_id)
        found = hot.get(key)
        if found is not None:
            return found
        recalled = self.warm_recall(agent_id, query=key, top_k=1)
        if not recalled:
            return None
        return recalled[0].content

    def warm_recall(self, agent_id: str, query: str, top_k: int = 5) -> list[WarmEntry]:
        """Recall entries from warm layer and inject into hot layer."""
        warm = self.get_warm(agent_id)
        hot = self.get_hot(agent_id)
        recalled = warm.recall(query, top_k=top_k)
        for idx, entry in enumerate(recalled):
            hot.set(f"warm_recall_{idx}_{entry.entry_id}", entry.content, tags=["warm_recall"])
        return recalled

    async def run_eviction_cycle(self, agent_id: str) -> None:
        """Evict stale hot entries, compress, and push to warm layer."""
        hot = self.get_hot(agent_id)
        warm = self.get_warm(agent_id)
        evicted = hot.evict_stale()
        if not evicted:
            return
        content = "\n".join(str(e.value) for e in evicted)
        keys = [e.key for e in evicted]
        warm.add(content=content, original_keys=keys, compression_ratio=1.0)
        for key in keys:
            await self.bus.publish_and_wait(
                StateMutationEvent(
                    source_id=agent_id,
                    layer="warm",
                    key=key,
                    old_value="migrated_from_hot",
                    new_value="stored_in_warm",
                )
            )

    def fork_state(self, agent_id: str, branch_id: str) -> str:
        """Fork an agent hot state for branch execution."""
        key = f"{agent_id}:{branch_id}"
        self._branches[key] = self.get_hot(agent_id).fork()
        return key

    async def merge_branch(self, branch_id: str, target_agent_id: str, strategy: str = "latest_wins") -> None:
        """Merge branch hot state into target agent hot state."""
        branch = self._branches.get(branch_id)
        if branch is None:
            return
        target = self.get_hot(target_agent_id)
        target.merge_from(branch, strategy=strategy)
        await self.bus.publish_and_wait(
            StateMutationEvent(
                source_id=target_agent_id,
                layer="hot",
                key=branch_id,
                old_value="branch",
                new_value="merged",
            )
        )

    def snapshot(self, agent_id: str) -> dict[str, Any]:
        """Snapshot hot layer state for checkpointing."""
        return self.get_hot(agent_id).to_dict()

    def restore_snapshot(self, agent_id: str, snapshot: dict[str, Any]) -> None:
        """Restore hot layer snapshot for an agent."""
        hot = self.get_hot(agent_id)
        hot.restore(snapshot)
