"""Tests for replay buffer strategies."""
import json
import math
import random
import tempfile
from pathlib import Path

import pytest

from llm_refresh.buffer import (
    DiverseBuffer,
    PrioritizedBuffer,
    ReservoirBuffer,
    create_buffer,
)
from llm_refresh.models import BufferStrategy


def make_examples(n: int, prefix: str = "text") -> list[dict]:
    return [{"text": f"{prefix} example {i}"} for i in range(n)]


# ---------------------------------------------------------------------------
# ReservoirBuffer
# ---------------------------------------------------------------------------

class TestReservoirBuffer:
    def test_add_under_capacity(self):
        buf = ReservoirBuffer(max_size=100)
        examples = make_examples(50)
        buf.add(examples)
        assert len(buf) == 50
        assert buf.stats()["total_seen"] == 50

    def test_add_over_capacity(self):
        buf = ReservoirBuffer(max_size=10)
        buf.add(make_examples(100))
        assert len(buf) == 10
        assert buf.stats()["total_seen"] == 100

    def test_sample_returns_correct_count(self):
        buf = ReservoirBuffer(max_size=100)
        buf.add(make_examples(50))
        sampled = buf.sample(20)
        assert len(sampled) == 20

    def test_sample_less_than_buffer(self):
        buf = ReservoirBuffer(max_size=100)
        buf.add(make_examples(5))
        sampled = buf.sample(20)
        assert len(sampled) == 5  # can't sample more than available

    def test_stats(self):
        buf = ReservoirBuffer(max_size=100)
        buf.add(make_examples(60))
        s = buf.stats()
        assert s["strategy"] == "reservoir"
        assert s["size"] == 60
        assert s["fill_pct"] == 60.0

    def test_save_load(self, tmp_path):
        buf = ReservoirBuffer(max_size=50)
        examples = make_examples(30)
        buf.add(examples)
        path = tmp_path / "reservoir.json"
        buf.save(path)

        buf2 = ReservoirBuffer()
        buf2.load(path)
        assert len(buf2) == 30
        assert buf2.max_size == 50
        assert buf2._total_seen == 30

    def test_uniformity(self):
        """Reservoir should sample roughly uniformly across all seen examples."""
        random.seed(42)
        buf = ReservoirBuffer(max_size=100)
        # Insert 1000 uniquely labelled examples
        all_examples = [{"text": f"item_{i}", "id": i} for i in range(1000)]
        buf.add(all_examples)

        # Count how many examples from first vs second half end up in buffer
        ids = {ex["id"] for ex in buf._reservoir}
        first_half = sum(1 for i in ids if i < 500)
        second_half = sum(1 for i in ids if i >= 500)

        # With uniform reservoir, expect roughly 50/50 split (±20 tolerance)
        assert abs(first_half - second_half) < 40, (
            f"Reservoir not uniform: first={first_half}, second={second_half}"
        )


# ---------------------------------------------------------------------------
# PrioritizedBuffer
# ---------------------------------------------------------------------------

class TestPrioritizedBuffer:
    def test_add_under_capacity(self):
        buf = PrioritizedBuffer(max_size=100)
        buf.add(make_examples(50))
        assert len(buf) == 50

    def test_add_over_capacity(self):
        buf = PrioritizedBuffer(max_size=10)
        buf.add(make_examples(50))
        assert len(buf) == 10

    def test_high_loss_retained(self):
        """High-loss examples should be preferentially kept."""
        buf = PrioritizedBuffer(max_size=5)
        # Add 5 low-loss examples
        low_loss = [{"text": f"easy_{i}", "_loss": 0.1} for i in range(5)]
        buf.add(low_loss)
        # Add 1 high-loss example
        hard = [{"text": "hard_example", "_loss": 10.0}]
        buf.add(hard)
        # The high-loss example should be in the buffer
        texts = [ex["text"] for ex in buf._examples]
        assert "hard_example" in texts

    def test_sample(self):
        buf = PrioritizedBuffer(max_size=100)
        buf.add(make_examples(50))
        sampled = buf.sample(20)
        assert len(sampled) == 20

    def test_stats(self):
        buf = PrioritizedBuffer(max_size=100)
        buf.add([{"text": "x", "_loss": 1.0}])
        s = buf.stats()
        assert s["strategy"] == "prioritized"
        assert s["avg_loss"] == 1.0

    def test_save_load(self, tmp_path):
        buf = PrioritizedBuffer(max_size=50)
        buf.add([{"text": f"ex_{i}", "_loss": float(i)} for i in range(30)])
        path = tmp_path / "prioritized.json"
        buf.save(path)

        buf2 = PrioritizedBuffer()
        buf2.load(path)
        assert len(buf2) == 30


# ---------------------------------------------------------------------------
# DiverseBuffer
# ---------------------------------------------------------------------------

class TestDiverseBuffer:
    def test_add_examples(self):
        buf = DiverseBuffer(max_size=64, n_buckets=8)
        buf.add(make_examples(32))
        assert len(buf) > 0
        assert len(buf) <= 64

    def test_different_texts_go_to_different_buckets(self):
        buf = DiverseBuffer(max_size=640, n_buckets=64)
        examples = [{"text": f"unique topic {i} " * 10} for i in range(200)]
        buf.add(examples)
        nonempty = buf.stats()["nonempty_buckets"]
        # Should spread across many buckets
        assert nonempty > 10

    def test_sample(self):
        buf = DiverseBuffer(max_size=128, n_buckets=8)
        buf.add(make_examples(80))
        sampled = buf.sample(20)
        assert len(sampled) <= 20

    def test_stats(self):
        buf = DiverseBuffer(max_size=64, n_buckets=8)
        buf.add(make_examples(10))
        s = buf.stats()
        assert s["strategy"] == "diverse"
        assert s["n_buckets"] == 8

    def test_save_load(self, tmp_path):
        buf = DiverseBuffer(max_size=64, n_buckets=8)
        buf.add(make_examples(20))
        original_size = len(buf)
        path = tmp_path / "diverse.json"
        buf.save(path)

        buf2 = DiverseBuffer()
        buf2.load(path)
        assert len(buf2) == original_size


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

class TestCreateBuffer:
    def test_reservoir(self):
        buf = create_buffer(BufferStrategy.RESERVOIR, max_size=100)
        assert isinstance(buf, ReservoirBuffer)

    def test_prioritized(self):
        buf = create_buffer(BufferStrategy.PRIORITIZED, max_size=100)
        assert isinstance(buf, PrioritizedBuffer)

    def test_diverse(self):
        buf = create_buffer(BufferStrategy.DIVERSE, max_size=100)
        assert isinstance(buf, DiverseBuffer)

    def test_string_strategy(self):
        buf = create_buffer("reservoir", max_size=50)
        assert isinstance(buf, ReservoirBuffer)

    def test_invalid_strategy(self):
        with pytest.raises(ValueError):
            create_buffer("unknown_strategy")
