import numpy as np
import pytest

from producer.engines.base import blend
from producer.engines.chunking import apply_chunked, plan_chunks, stitch


def test_plan_chunks_basic():
    assert plan_chunks(0, 5, 2) == []
    assert plan_chunks(4, 10, 2) == [(0, 4)]
    assert plan_chunks(10, 5, 0) == [(0, 5), (5, 10)]
    assert plan_chunks(10, 5, 2) == [(0, 5), (3, 8), (6, 10)]
    assert plan_chunks(7, 5, 2) == [(0, 5), (3, 7)]


def test_plan_chunks_last_span_covers_tail():
    spans = plan_chunks(11, 5, 4)
    assert spans[-1][1] == 11
    assert spans[-1][1] - spans[-1][0] >= 5


def test_plan_chunks_rejects_overlap_ge_chunk():
    with pytest.raises(ValueError):
        plan_chunks(100, 5, 5)
    with pytest.raises(ValueError):
        plan_chunks(100, 5, 6)


def test_stitch_identity_pieces_reconstruct_signal():
    rng = np.random.default_rng(0)
    x = rng.standard_normal(10_000).astype(np.float32)
    spans = plan_chunks(x.size, 2500, 300)
    y = stitch(spans, x.size, (x[a:b] for a, b in spans))
    assert y.dtype == np.float32
    assert y.shape == x.shape
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_stitch_single_span_passthrough():
    x = np.arange(100, dtype=np.float32)
    y = stitch([(0, 100)], 100, [x])
    np.testing.assert_array_equal(y, x)


def test_stitch_crossfade_of_complementary_pieces():
    # piece A is silence, piece B is a full-scale ramp; the crossfade zone
    # must be a smooth blend, not a jump.
    spans = [(0, 10), (5, 15)]
    pieces = [np.zeros(10, dtype=np.float32), np.ones(10, dtype=np.float32)]
    y = stitch(spans, 15, pieces)
    mid = y[7]  # 50% through the overlap
    assert 0.4 < mid < 0.6
    assert y[0] == 0.0
    assert y[14] == 1.0


def test_stitch_pads_short_pieces():
    spans = [(0, 10), (5, 15)]
    pieces = [np.arange(10, dtype=np.float32), np.arange(8, dtype=np.float32)]
    y = stitch(spans, 15, pieces)
    assert y.shape == (15,)
    assert np.all(np.isfinite(y))


def test_apply_chunked_identity_matches_whole_file():
    rng = np.random.default_rng(1)
    x = rng.standard_normal(50_000).astype(np.float32)
    sr = 8000
    calls = []

    def fn(chunk):
        calls.append(chunk.size)
        return chunk

    y = apply_chunked(x, sr, 7.0, 0.05, fn)
    assert calls == [50_000]
    np.testing.assert_array_equal(y, x)

    calls.clear()
    y = apply_chunked(x, sr, 0.5, 0.05, fn)
    assert len(calls) >= 10
    assert all(c <= 4000 for c in calls)
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_zero_disables_chunking():
    x = np.zeros(10, dtype=np.float32)
    calls = []

    def fn(chunk):
        calls.append(chunk.size)
        return chunk

    apply_chunked(x, 8000, 0.0, 0.5, fn)
    assert calls == [10]


def test_apply_chunked_context_reconstructs_identity():
    rng = np.random.default_rng(3)
    x = rng.standard_normal(30_000).astype(np.float32)
    y = apply_chunked(x, 8000, 1.5, 0.25, lambda c: c, context_s=0.5)
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_context_feeds_padded_chunks():
    x = np.ones(5_000, dtype=np.float32)
    sizes = []

    def fn(chunk):
        sizes.append(chunk.size)
        return chunk

    apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5)
    # plan_chunks(5000, 2000, 0) -> (0,2000), (2000,4000), (4000,5000)
    assert sizes == [2500, 3000, 1500]


def test_apply_chunked_context_trims_cold_start_artifacts():
    x = np.ones(6_000, dtype=np.float32)

    def fn(chunk):
        out = chunk.copy()
        out[0] = 0.0  # cold-start artifact at the start of every model call
        return out

    broken = apply_chunked(x, 1000, 2.0, 0.0, fn)
    assert int(np.sum(broken == 0.0)) > 1

    y = apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5)
    assert y[0] == 0.0  # only the true file start stays degraded
    assert np.all(y[1:] == 1.0)


def test_apply_chunked_context_survives_oom_retry():
    x = np.ones(8_000, dtype=np.float32)
    attempts = []

    def fn(chunk):
        attempts.append(chunk.size)
        if chunk.size > 3000:
            raise MemoryError("simulated oom")
        out = chunk.copy()
        out[0] = 0.0
        return out

    y = apply_chunked(x, 1000, 6.0, 0.0, fn, min_chunk_s=1.0, context_s=0.5)
    assert max(attempts) > 3000
    assert y[0] == 0.0
    assert np.all(y[1:] == 1.0)


def test_apply_chunked_retries_smaller_on_oom():
    x = np.random.default_rng(2).standard_normal(40_000).astype(np.float32)
    sr = 8000
    attempts = []

    def fn(chunk):
        attempts.append(chunk.size)
        if chunk.size > 3000:
            raise MemoryError("simulated oom")
        return chunk

    y = apply_chunked(x, sr, 5.0, 0.0, fn, min_chunk_s=0.25)
    assert attempts[0] == 40_000
    assert attempts[-1] == 2500
    assert max(a for a in attempts if a <= 3000) == 2500
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_oom_gives_up_at_min_chunk():
    x = np.zeros(40_000, dtype=np.float32)

    def fn(chunk):
        raise MemoryError("always")

    with pytest.raises(MemoryError):
        apply_chunked(x, 8000, 1.0, 0.0, fn, min_chunk_s=0.5)


def test_apply_chunked_progress_callback():
    x = np.random.default_rng(3).standard_normal(20_000).astype(np.float32)
    sr = 8000
    seen: list[tuple[int, int]] = []
    y = apply_chunked(x, sr, 0.5, 0.0, lambda c: c, on_progress=lambda d, t: seen.append((d, t)))
    total = len(plan_chunks(x.size, round(0.5 * sr), 0))
    assert seen[0] == (0, total)
    assert seen[-1] == (total, total)
    assert [d for d, _t in seen[1:]] == list(range(1, total + 1))
    np.testing.assert_allclose(y, x, atol=1e-6)


def test_apply_chunked_progress_resets_after_oom_retry():
    x = np.random.default_rng(4).standard_normal(40_000).astype(np.float32)
    sr = 8000
    seen: list[tuple[int, int]] = []

    def fn(chunk):
        if chunk.size > 3000:
            raise MemoryError("simulated oom")
        return chunk

    y = apply_chunked(
        x, sr, 5.0, 0.0, fn, min_chunk_s=0.25, on_progress=lambda d, t: seen.append((d, t))
    )
    assert seen[0] == (0, 1)  # first attempt is one whole-file chunk
    done, total = seen[-1]
    assert done == total and total > 2  # retried into smaller chunks
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_progress_absent_when_disabled():
    x = np.zeros(10_000, dtype=np.float32)
    y = apply_chunked(x, 8000, 0.5, 0.0, lambda c: c, on_progress=None)
    assert y.shape == x.shape


def test_apply_chunked_whole_file_mode_falls_back_on_oom():
    x = np.random.default_rng(5).standard_normal(10_000).astype(np.float32)
    sr = 1000

    def fn(chunk):
        if chunk.size > 1500:
            raise MemoryError("simulated oom")
        return chunk * 2

    y = apply_chunked(x, sr, 0.0, 0.0, fn, min_chunk_s=0.5)
    np.testing.assert_allclose(y, x * 2, atol=1e-6)


def test_blend_float32_no_float64_temporaries():
    x = np.full(1000, 0.25, dtype=np.float32)
    y = np.full(1000, 0.75, dtype=np.float32)
    out = blend(x, y, 0.5)
    assert out.dtype == np.float32
    np.testing.assert_allclose(out, 0.5, atol=1e-7)
    assert blend(x, y, 1.0) is y
    np.testing.assert_allclose(blend(x, y, 0.0), x, atol=1e-7)


def test_blend_realigns_resample_drift():
    # a 48k -> 16k -> 48k round trip can come back a sample or two long
    # (resample_poly emits ceil(n * up/down) per hop); blend must cope
    x = np.ones(5, dtype=np.float32)
    long_y = np.full(7, 0.5, dtype=np.float32)
    np.testing.assert_allclose(blend(x, long_y, 0.5), np.full(5, 0.75), atol=1e-7)

    short_y = np.full(4, 0.5, dtype=np.float32)
    out = blend(x, short_y, 0.5)
    assert out.shape == x.shape
    np.testing.assert_allclose(out, [0.75, 0.75, 0.75, 0.75, 0.5], atol=1e-7)