"""Audio playback utilities for patch visualization."""

from __future__ import annotations

from functools import lru_cache
from pathlib import Path
from typing import List, Tuple

import numpy as np
import soundfile as sf

TOKENS_PER_SECOND = 40  # WavTokenizer rate


@lru_cache(maxsize=16)
def _cached_audio(path: str, mtime: float) -> Tuple[int, np.ndarray]:
    """Read audio data from disk and memoize the decoded array."""
    audio, sr = sf.read(path)
    return sr, audio


def clear_audio_cache() -> None:
    """Drop all cached audio samples (useful for tests)."""
    _cached_audio.cache_clear()


def _read_audio(audio_path: str) -> Tuple[int, np.ndarray]:
    resolved = Path(audio_path).resolve()
    stat = resolved.stat()
    return _cached_audio(str(resolved), stat.st_mtime)


def load_audio_file(audio_path: str) -> Tuple[int, np.ndarray]:
    """Load audio file, reusing cached decode when possible."""
    return _read_audio(audio_path)


def get_patch_audio(
    audio_path: str,
    patches: List[List[int]],
    patch_index: int
) -> Tuple[int, np.ndarray]:
    """
    Extract audio for a specific patch.

    Args:
        audio_path: Path to audio file
        patches: List of patches
        patch_index: Index of patch to extract

    Returns:
        (sample_rate, audio_segment)
    """
    if patch_index < 0 or patch_index >= len(patches):
        return 24000, np.array([])

    sr, audio = _read_audio(audio_path)
    patch = patches[patch_index]

    if not patch:
        return sr, np.array([])

    # Convert token indices to audio samples
    start_token = patch[0]
    end_token = patch[-1]

    start_sample = int(start_token / TOKENS_PER_SECOND * sr)
    end_sample = int((end_token + 1) / TOKENS_PER_SECOND * sr)

    # Clamp to valid range
    start_sample = max(0, start_sample)
    end_sample = min(len(audio), end_sample)

    segment = audio[start_sample:end_sample]

    return sr, segment


def play_until_boundary(
    audio_path: str,
    patches: List[List[int]],
    current_token: int
) -> Tuple[int, np.ndarray]:
    """
    Get audio from current token position until next patch boundary.

    Args:
        audio_path: Path to audio file
        patches: List of patches
        current_token: Current token position

    Returns:
        (sample_rate, audio_segment)
    """
    # Find which patch contains current token
    current_patch_idx = None
    for i, patch in enumerate(patches):
        if current_token in patch:
            current_patch_idx = i
            break

    if current_patch_idx is None:
        # Token not in any patch, return empty
        return 24000, np.array([])

    # Get audio for this patch from current position onward
    sr, audio = _read_audio(audio_path)
    patch = patches[current_patch_idx]

    start_token = current_token
    end_token = patch[-1]

    # Convert to audio samples
    start_sample = int(start_token / TOKENS_PER_SECOND * sr)
    end_sample = int((end_token + 1) / TOKENS_PER_SECOND * sr)

    # Clamp to valid range
    start_sample = max(0, start_sample)
    end_sample = min(len(audio), end_sample)

    segment = audio[start_sample:end_sample]

    return sr, segment


def find_patch_for_token(patches: List[List[int]], token_idx: int) -> int:
    """
    Find which patch contains a given token.

    Args:
        patches: List of patches
        token_idx: Token index

    Returns:
        Patch index, or -1 if not found
    """
    for i, patch in enumerate(patches):
        if token_idx in patch:
            return i
    return -1


def get_patch_info(
    patches: List[List[int]],
    patch_idx: int
) -> dict:
    """
    Get information about a specific patch.

    Args:
        patches: List of patches
        patch_idx: Patch index

    Returns:
        Dictionary with patch info
    """
    if patch_idx < 0 or patch_idx >= len(patches):
        return {}

    patch = patches[patch_idx]
    start_token = patch[0]
    end_token = patch[-1]
    num_tokens = len(patch)

    start_time = start_token / TOKENS_PER_SECOND
    end_time = (end_token + 1) / TOKENS_PER_SECOND
    duration = num_tokens / TOKENS_PER_SECOND

    return {
        "patch_index": patch_idx,
        "start_token": start_token,
        "end_token": end_token,
        "num_tokens": num_tokens,
        "start_time_sec": start_time,
        "end_time_sec": end_time,
        "duration_sec": duration,
    }