"""Three replay buffer strategies for continual learning."""
from __future__ import annotations

import hashlib
import json
import random
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any

from llm_refresh.models import BufferStrategy


class BaseBuffer(ABC):
    def __init__(self, max_size: int = 10_000):
        self.max_size = max_size
        self._total_seen: int = 0

    @abstractmethod
    def add(self, examples: list[dict[str, Any]]) -> None:
        """Add examples to the buffer."""

    @abstractmethod
    def sample(self, n: int) -> list[dict[str, Any]]:
        """Sample up to n examples from the buffer."""

    @abstractmethod
    def stats(self) -> dict[str, Any]:
        """Return buffer statistics."""

    def save(self, path: Path | str) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        state = {
            "type": self.__class__.__name__,
            "max_size": self.max_size,
            "total_seen": self._total_seen,
            **self._get_state(),
        }
        path.write_text(json.dumps(state, ensure_ascii=False, indent=2))

    def load(self, path: Path | str) -> None:
        path = Path(path)
        state = json.loads(path.read_text())
        self.max_size = state["max_size"]
        self._total_seen = state["total_seen"]
        self._set_state(state)

    @abstractmethod
    def _get_state(self) -> dict[str, Any]:
        """Subclass-specific state for serialization."""

    @abstractmethod
    def _set_state(self, state: dict[str, Any]) -> None:
        """Restore subclass-specific state."""

    def __len__(self) -> int:
        return len(self.sample(self.max_size))


class ReservoirBuffer(BaseBuffer):
    """Vitter's Algorithm R — uniform random sample over all seen examples."""

    def __init__(self, max_size: int = 10_000):
        super().__init__(max_size)
        self._reservoir: list[dict[str, Any]] = []

    def add(self, examples: list[dict[str, Any]]) -> None:
        for ex in examples:
            self._total_seen += 1
            if len(self._reservoir) < self.max_size:
                self._reservoir.append(ex)
            else:
                j = random.randint(0, self._total_seen - 1)
                if j < self.max_size:
                    self._reservoir[j] = ex

    def sample(self, n: int) -> list[dict[str, Any]]:
        if n >= len(self._reservoir):
            return list(self._reservoir)
        return random.sample(self._reservoir, n)

    def stats(self) -> dict[str, Any]:
        return {
            "strategy": "reservoir",
            "size": len(self._reservoir),
            "max_size": self.max_size,
            "total_seen": self._total_seen,
            "fill_pct": round(100 * len(self._reservoir) / max(self.max_size, 1), 1),
        }

    def _get_state(self) -> dict[str, Any]:
        return {"reservoir": self._reservoir}

    def _set_state(self, state: dict[str, Any]) -> None:
        self._reservoir = state.get("reservoir", [])

    def __len__(self) -> int:
        return len(self._reservoir)


class PrioritizedBuffer(BaseBuffer):
    """Keeps examples with highest loss; evicts easiest examples first."""

    def __init__(self, max_size: int = 10_000):
        super().__init__(max_size)
        self._examples: list[dict[str, Any]] = []
        self._losses: list[float] = []

    def add(self, examples: list[dict[str, Any]]) -> None:
        for ex in examples:
            self._total_seen += 1
            loss = float(ex.get("_loss", random.uniform(0.5, 2.0)))
            if len(self._examples) < self.max_size:
                self._examples.append(ex)
                self._losses.append(loss)
            else:
                # Evict lowest-loss (easiest) example via random top-k scan
                k = min(50, len(self._examples))
                candidates = random.sample(range(len(self._examples)), k)
                min_idx = min(candidates, key=lambda i: self._losses[i])
                if loss > self._losses[min_idx]:
                    self._examples[min_idx] = ex
                    self._losses[min_idx] = loss

    def sample(self, n: int) -> list[dict[str, Any]]:
        if not self._examples:
            return []
        total_loss = sum(self._losses)
        if total_loss <= 0:
            weights = None
        else:
            weights = [l / total_loss for l in self._losses]
        k = min(n, len(self._examples))
        indices = random.choices(range(len(self._examples)), weights=weights, k=k)
        return [self._examples[i] for i in indices]

    def stats(self) -> dict[str, Any]:
        avg_loss = sum(self._losses) / max(len(self._losses), 1)
        return {
            "strategy": "prioritized",
            "size": len(self._examples),
            "max_size": self.max_size,
            "total_seen": self._total_seen,
            "avg_loss": round(avg_loss, 4),
            "fill_pct": round(100 * len(self._examples) / max(self.max_size, 1), 1),
        }

    def _get_state(self) -> dict[str, Any]:
        return {"examples": self._examples, "losses": self._losses}

    def _set_state(self, state: dict[str, Any]) -> None:
        self._examples = state.get("examples", [])
        self._losses = state.get("losses", [])

    def __len__(self) -> int:
        return len(self._examples)


class DiverseBuffer(BaseBuffer):
    """Hash-based bucketing into n_buckets slots to prevent topic dominance."""

    def __init__(self, max_size: int = 10_000, n_buckets: int = 64):
        super().__init__(max_size)
        self.n_buckets = n_buckets
        self._bucket_cap = max(1, max_size // n_buckets)
        self._buckets: list[list[dict[str, Any]]] = [[] for _ in range(n_buckets)]

    def _bucket_idx(self, ex: dict[str, Any]) -> int:
        text = ex.get("text", str(ex))
        h = hashlib.md5(text.encode("utf-8", errors="replace")).hexdigest()
        return int(h[:8], 16) % self.n_buckets

    def add(self, examples: list[dict[str, Any]]) -> None:
        for ex in examples:
            self._total_seen += 1
            idx = self._bucket_idx(ex)
            bucket = self._buckets[idx]
            if len(bucket) < self._bucket_cap:
                bucket.append(ex)
            else:
                # Replace a random item in the bucket
                bucket[random.randint(0, len(bucket) - 1)] = ex

    def sample(self, n: int) -> list[dict[str, Any]]:
        pool: list[dict[str, Any]] = []
        for bucket in self._buckets:
            pool.extend(bucket)
        if n >= len(pool):
            return pool
        return random.sample(pool, n)

    def stats(self) -> dict[str, Any]:
        sizes = [len(b) for b in self._buckets]
        total = sum(sizes)
        nonempty = sum(1 for s in sizes if s > 0)
        return {
            "strategy": "diverse",
            "size": total,
            "max_size": self.max_size,
            "total_seen": self._total_seen,
            "n_buckets": self.n_buckets,
            "nonempty_buckets": nonempty,
            "fill_pct": round(100 * total / max(self.max_size, 1), 1),
        }

    def _get_state(self) -> dict[str, Any]:
        return {"buckets": self._buckets, "n_buckets": self.n_buckets}

    def _set_state(self, state: dict[str, Any]) -> None:
        self.n_buckets = state.get("n_buckets", self.n_buckets)
        self._buckets = state.get("buckets", [[] for _ in range(self.n_buckets)])
        self._bucket_cap = max(1, self.max_size // self.n_buckets)

    def __len__(self) -> int:
        return sum(len(b) for b in self._buckets)


def create_buffer(strategy: BufferStrategy | str, **kwargs: Any) -> BaseBuffer:
    s = BufferStrategy(strategy) if isinstance(strategy, str) else strategy
    if s == BufferStrategy.RESERVOIR:
        return ReservoirBuffer(**kwargs)
    if s == BufferStrategy.PRIORITIZED:
        return PrioritizedBuffer(**kwargs)
    if s == BufferStrategy.DIVERSE:
        return DiverseBuffer(**kwargs)
    raise ValueError(f"Unknown strategy: {strategy}")
