# ------------------------------------------------------------
# AAC Coder/Decoder - AAC Coder/DecoderTests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Tests for AAC Coder/Decoder module.
# ------------------------------------------------------------
from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
import soundfile as sf

from core.aac_coder import aac_coder_1, aac_coder_2, aac_coder_3, aac_read_wav_stereo_48k
from core.aac_decoder import aac_decoder_1, aac_decoder_2, aac_decoder_3, aac_remove_padding
from core.aac_utils import snr_db
from core.aac_types import *


# Helper "fixtures" for aac_coder_1 / i_aac_coder_1
# -----------------------------------------------------------------------------

@pytest.fixture()
def tmp_stereo_wav(tmp_path: Path) -> Path:
    """
    Create a temporary 48 kHz stereo WAV with random samples.
    """
    rng = np.random.default_rng(123)
    fs = 48000

    # ~1 second of audio (kept small for test speed).
    n = fs
    x: StereoSignal = rng.normal(size=(n, 2)).astype(np.float64)

    wav_path = tmp_path / "in.wav"
    sf.write(str(wav_path), x, fs)
    return wav_path


# -----------------------------------------------------------------------------
# Helper-function tests
# -----------------------------------------------------------------------------

def test_aac_read_wav_stereo_48k_roundtrip(tmp_stereo_wav: Path) -> None:
    """
    Contract test for aac_read_wav_stereo_48k():
    - Reads stereo WAV
    - Returns float64 array with shape (N,2)
    - Returns fs = 48000
    """
    x, fs = aac_read_wav_stereo_48k(tmp_stereo_wav)

    assert int(fs) == 48000
    assert isinstance(x, np.ndarray)
    assert x.dtype == np.float64
    assert x.ndim == 2
    assert x.shape[1] == 2
    assert x.shape[0] > 0


def test_aac_remove_padding_removes_hop_from_both_ends() -> None:
    """
    Contract test for aac_remove_padding():
    - Removes 'hop' samples from start and end.
    """
    hop = 1024
    n = 10000

    y_pad: StereoSignal = np.zeros((n, 2), dtype=np.float64)
    y: StereoSignal = aac_remove_padding(y_pad, hop=hop)

    assert y.shape == (n - 2 * hop, 2)
    assert y.dtype == np.float64


def test_aac_remove_padding_errors_on_too_short_input() -> None:
    """
    aac_remove_padding must raise if y_pad is shorter than 2*hop.
    """
    hop = 1024
    y_pad: StereoSignal = np.zeros((2 * hop - 1, 2), dtype=np.float64)

    with pytest.raises(ValueError):
        _ = aac_remove_padding(y_pad, hop=hop)


# -----------------------------------------------------------------------------
# Level 1 tests
# -----------------------------------------------------------------------------
def test_aac_coder_seq_schema_and_shapes(tmp_stereo_wav: Path) -> None:
    """
    Module-level contract test:
    Ensure aac_seq_1 follows the expected schema and per-frame shapes.
    """
    aac_seq: AACSeq1 = aac_coder_1(tmp_stereo_wav)

    assert isinstance(aac_seq, list)
    assert len(aac_seq) > 0

    for fr in aac_seq:
        assert isinstance(fr, dict)

        # Required keys
        assert "frame_type" in fr
        assert "win_type" in fr
        assert "chl" in fr
        assert "chr" in fr

        frame_type = fr["frame_type"]
        win_type = fr["win_type"]

        assert frame_type in ("OLS", "LSS", "ESH", "LPS")
        assert win_type in ("SIN", "KBD")

        assert isinstance(fr["chl"], dict)
        assert isinstance(fr["chr"], dict)
        assert "frame_F" in fr["chl"]
        assert "frame_F" in fr["chr"]

        chl_f = np.asarray(fr["chl"]["frame_F"], dtype=np.float64)
        chr_f = np.asarray(fr["chr"]["frame_F"], dtype=np.float64)

        if frame_type == "ESH":
            assert chl_f.shape == (128, 8)
            assert chr_f.shape == (128, 8)
        else:
            assert chl_f.shape == (1024, 1)
            assert chr_f.shape == (1024, 1)


def test_end_to_end_aac_coder_decoder_high_snr(tmp_stereo_wav: Path, tmp_path: Path) -> None:
    """
    End-to-end test:
    Encode + decode and check SNR is very high (numerical-noise only).

    The threshold is intentionally loose to avoid fragility across platforms/BLAS.
    """
    x_ref, fs = sf.read(str(tmp_stereo_wav), always_2d=True)
    x_ref = np.asarray(x_ref, dtype=np.float64)
    assert int(fs) == 48000

    out_wav = tmp_path / "out.wav"

    aac_seq = aac_coder_1(tmp_stereo_wav)
    x_hat: StereoSignal = aac_decoder_1(aac_seq, out_wav)

    # Basic sanity: output file exists and is readable
    assert out_wav.exists()
    x_hat_file, fs_hat = sf.read(str(out_wav), always_2d=True)
    assert int(fs_hat) == 48000

    # SNR against returned array (file should match closely, but we do not require it here).
    snr = snr_db(x_ref, x_hat)
    assert snr > 80.0

# -----------------------------------------------------------------------------
# Level 2 tests (new)
# -----------------------------------------------------------------------------

def test_aac_coder_2_seq_schema_and_shapes(tmp_stereo_wav: Path) -> None:
    """
    Module-level contract test (Level 2):
    Ensure aac_seq_2 follows the expected schema and per-frame shapes, including tns_coeffs.
    """
    aac_seq: AACSeq2 = aac_coder_2(tmp_stereo_wav)

    assert isinstance(aac_seq, list)
    assert len(aac_seq) > 0

    for fr in aac_seq:
        assert "frame_type" in fr
        assert "win_type" in fr
        assert "chl" in fr
        assert "chr" in fr

        frame_type: FrameType = fr["frame_type"]
        assert frame_type in ("OLS", "LSS", "ESH", "LPS")

        for ch_key in ("chl", "chr"):
            ch = fr[ch_key]
            assert "frame_F" in ch
            assert "tns_coeffs" in ch

            frame_f = np.asarray(ch["frame_F"], dtype=np.float64)
            coeffs = np.asarray(ch["tns_coeffs"], dtype=np.float64)

            if frame_type == "ESH":
                assert frame_f.shape == (128, 8)
                assert coeffs.shape[0] == 4
                assert coeffs.shape[1] == 8
            else:
                assert frame_f.shape == (1024, 1)
                assert coeffs.shape == (4, 1)


def test_end_to_end_level_2_high_snr(tmp_stereo_wav: Path, tmp_path: Path) -> None:
    """
    End-to-end test (Level 2):
    Encode + decode and check SNR remains very high.

    Level 2 is still floating-point (TNS is reversible), so reconstruction
    should remain numerical-noise only.
    """
    x_ref, fs = sf.read(str(tmp_stereo_wav), always_2d=True)
    x_ref = np.asarray(x_ref, dtype=np.float64)
    assert int(fs) == 48000

    out_wav = tmp_path / "out_l2.wav"
    aac_seq = aac_coder_2(tmp_stereo_wav)
    x_hat: StereoSignal = aac_decoder_2(aac_seq, out_wav)

    assert out_wav.exists()
    _, fs_hat = sf.read(str(out_wav), always_2d=True)
    assert int(fs_hat) == 48000

    snr = snr_db(x_ref, x_hat)
    assert snr > 80


# -----------------------------------------------------------------------------
# Level 3 tests (Quantizer + Huffman)
# -----------------------------------------------------------------------------

@pytest.fixture(scope="module")
def wav_in_path() -> Path:
    """
    Input WAV used for end-to-end tests.

    This should point to the provided test audio under material/.
    Adjust this path if your project layout differs.
    """
    # Typical layout in this project:
    # source/material/LicorDeCalandraca.wav
    return Path(__file__).resolve().parents[2] / "material" / "LicorDeCalandraca.wav"


def _assert_level3_frame_schema(frame: AACSeq3Frame) -> None:
    """
    Validate Level-3 per-frame schema (keys + basic types only).
    """
    assert "frame_type" in frame
    assert "win_type" in frame
    assert "chl" in frame
    assert "chr" in frame

    for ch_key in ("chl", "chr"):
        ch = frame[ch_key]  # type: ignore[index]
        assert "tns_coeffs" in ch
        assert "T" in ch
        assert "G" in ch
        assert "sfc" in ch
        assert "stream" in ch
        assert "codebook" in ch

        assert isinstance(ch["sfc"], str)
        assert isinstance(ch["stream"], str)
        assert isinstance(ch["codebook"], int)

        # Arrays: only check they are numpy arrays with expected dtype categories.
        assert isinstance(ch["tns_coeffs"], np.ndarray)
        assert isinstance(ch["T"], np.ndarray)

        # Global gain: long frames may be scalar float, ESH may be ndarray
        assert np.isscalar(ch["G"]) or isinstance(ch["G"], np.ndarray)


def test_aac_coder_3_seq_schema_and_shapes(wav_in_path: Path, tmp_path: Path) -> None:
    """
    Contract test:
    - aac_coder_3 returns AACSeq3
    - Per-frame keys exist and types are consistent
    - Basic shape expectations hold for ESH vs non-ESH cases

    Note:
    This test uses a short excerpt (a few frames) to keep runtime bounded.
    """
    # Use only a few frames to avoid long runtimes in the quantizer loop.
    hop = 1024
    win = 2048
    n_frames = 4
    n_samples = win + (n_frames - 1) * hop

    x, fs = aac_read_wav_stereo_48k(wav_in_path)
    x_short = x[:n_samples, :]

    short_wav = tmp_path / "input_short.wav"
    sf.write(str(short_wav), x_short, fs)

    aac_seq_3: AACSeq3 = aac_coder_3(short_wav)

    assert isinstance(aac_seq_3, list)
    assert len(aac_seq_3) > 0

    for fr in aac_seq_3:
        _assert_level3_frame_schema(fr)

        frame_type = fr["frame_type"]
        for ch_key in ("chl", "chr"):
            ch = fr[ch_key]  # type: ignore[index]

            tns = np.asarray(ch["tns_coeffs"])
            if frame_type == "ESH":
                assert tns.ndim == 2
                assert tns.shape[1] == 8
            else:
                assert tns.ndim == 2
                assert tns.shape[1] == 1

            T = np.asarray(ch["T"])
            if frame_type == "ESH":
                assert T.ndim == 2
                assert T.shape[1] == 8
            else:
                assert T.ndim == 2
                assert T.shape[1] == 1

            G = ch["G"]
            if frame_type == "ESH":
                assert isinstance(G, np.ndarray)
                assert np.asarray(G).shape == (1, 8)
            else:
                assert np.isscalar(G)

            assert isinstance(ch["sfc"], str)
            assert isinstance(ch["stream"], str)



def test_end_to_end_level_3_high_snr(wav_in_path: Path, tmp_path: Path) -> None:
    """
    End-to-end test for Level 3 (Quantizer + Huffman):

    coder_3 -> decoder_3 should reconstruct a waveform with acceptable SNR.

    Notes
    -----
    - Level 3 includes quantization, so SNR is expected to be lower than Level 1/2.
    - We intentionally use a short excerpt (few frames) to keep runtime bounded,
      since the reference quantizer implementation is computationally expensive.
    """
    # Use only a few frames to avoid long runtimes.
    hop = 1024
    win = 2048
    n_frames = 4
    n_samples = win + (n_frames - 1) * hop

    x_ref, fs = aac_read_wav_stereo_48k(wav_in_path)
    x_short = x_ref[:n_samples, :]

    short_wav = tmp_path / "input_short_l3.wav"
    sf.write(str(short_wav), x_short, fs)

    out_wav = tmp_path / "decoded_level3.wav"

    aac_seq_3: AACSeq3 = aac_coder_3(short_wav)
    y_hat: StereoSignal = aac_decoder_3(aac_seq_3, out_wav)

    # Align lengths defensively (padding removal may differ by a few samples)
    n = min(x_short.shape[0], y_hat.shape[0])
    x2 = x_short[:n, :]
    y2 = y_hat[:n, :]

    s = snr_db(x2, y2)

    # Conservative threshold: Level 3 is lossy by design.
    assert s > 10.0