import logging

import librosa
import numpy as np
import pyloudnorm as pyln

from songmatch.config import SongMatchConfig

logger = logging.getLogger(__name__)


def compute_lufs_envelope(
    audio: np.ndarray, sr: int, window_sec: float = 3.0, hop_sec: float = 0.1
) -> np.ndarray:
    """Compute a short-term LUFS envelope over the audio.

    Returns a 1D array of LUFS values, one per hop.
    """
    meter = pyln.Meter(sr)
    window_samples = int(window_sec * sr)
    hop_samples = int(hop_sec * sr)

    # Ensure 2D for pyloudnorm
    if audio.ndim == 1:
        audio_2d = audio[:, np.newaxis]
    else:
        audio_2d = audio

    envelope = []
    pos = 0
    while pos + window_samples <= len(audio_2d):
        block = audio_2d[pos : pos + window_samples]
        try:
            lufs = meter.integrated_loudness(block)
        except Exception:
            lufs = -np.inf
        envelope.append(lufs)
        pos += hop_samples

    result = np.array(envelope)
    logger.debug(
        "LUFS envelope: %d frames, range %.1f to %.1f LUFS",
        len(result),
        np.min(result[np.isfinite(result)]) if np.any(np.isfinite(result)) else -np.inf,
        np.max(result[np.isfinite(result)]) if np.any(np.isfinite(result)) else -np.inf,
    )
    return result


def find_loudest_segment(
    envelope: np.ndarray, hop_sec: float, config: SongMatchConfig
) -> tuple[float, float] | None:
    """Find the loudest contiguous segment in a LUFS envelope.

    Returns (start_sec, end_sec) or None if cropping should be skipped.
    """
    finite_mask = np.isfinite(envelope)
    if not np.any(finite_mask):
        logger.warning("All LUFS values are -inf (silence). Skipping cropping.")
        return None

    finite_values = envelope[finite_mask]
    peak_lufs = float(np.max(finite_values))
    min_lufs = float(np.min(finite_values))
    loudness_range = peak_lufs - min_lufs

    logger.debug("Peak LUFS: %.1f, Min LUFS: %.1f, Range: %.1f dB", peak_lufs, min_lufs, loudness_range)

    # Flat track check
    if loudness_range < config.dynamic_range_floor:
        logger.info("Flat track (LRA %.1f dB < %.1f dB floor). Skipping cropping.", loudness_range, config.dynamic_range_floor)
        return None

    # Find loudest segment with threshold relaxation
    threshold = peak_lufs - config.loudness_threshold
    min_frames = int(config.min_segment / hop_sec)

    for attempt in range(4):
        if attempt > 0:
            threshold -= 1.0
            logger.debug("Relaxing threshold to %.1f LUFS (attempt %d)", threshold, attempt + 1)

        mask = envelope >= threshold
        regions = _find_contiguous_regions(mask)

        # Filter by minimum length
        regions = [(s, e) for s, e in regions if (e - s) >= min_frames]

        if regions:
            break
    else:
        # Last resort: use longest region regardless of min_segment
        all_regions = _find_contiguous_regions(envelope >= (peak_lufs - config.loudness_threshold - 3.0))
        if not all_regions:
            logger.warning("No suitable loud segment found. Skipping cropping.")
            return None
        regions = all_regions

    # Select longest region
    longest = max(regions, key=lambda r: r[1] - r[0])
    start_sec = longest[0] * hop_sec
    end_sec = longest[1] * hop_sec

    # Trim to max_segment from center if too long
    duration = end_sec - start_sec
    if duration > config.max_segment:
        center = (start_sec + end_sec) / 2
        start_sec = center - config.max_segment / 2
        end_sec = center + config.max_segment / 2

    logger.info("Loudest segment: %.1fs - %.1fs (%.1fs)", start_sec, end_sec, end_sec - start_sec)
    return start_sec, end_sec


def _find_contiguous_regions(mask: np.ndarray) -> list[tuple[int, int]]:
    """Find contiguous True regions in a boolean array.

    Returns list of (start_idx, end_idx) tuples.
    """
    regions = []
    in_region = False
    start = 0
    for i, val in enumerate(mask):
        if val and not in_region:
            start = i
            in_region = True
        elif not val and in_region:
            regions.append((start, i))
            in_region = False
    if in_region:
        regions.append((start, len(mask)))
    return regions


def snap_to_beats(
    start_sec: float, end_sec: float, audio: np.ndarray, sr: int
) -> tuple[float, float]:
    """Snap segment boundaries to nearest beat positions via librosa."""
    # Convert to mono for beat tracking
    if audio.ndim == 2:
        mono = np.mean(audio, axis=1)
    else:
        mono = audio

    _, beat_frames = librosa.beat.beat_track(y=mono, sr=sr)
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)

    if len(beat_times) == 0:
        logger.debug("No beats detected. Keeping original boundaries.")
        return start_sec, end_sec

    # Snap start to nearest preceding beat
    preceding = beat_times[beat_times <= start_sec]
    if len(preceding) > 0:
        start_sec = float(preceding[-1])

    # Snap end to nearest following beat
    following = beat_times[beat_times >= end_sec]
    if len(following) > 0:
        end_sec = float(following[0])

    return start_sec, end_sec


def crop_reference(
    audio: np.ndarray, sr: int, config: SongMatchConfig
) -> tuple[np.ndarray, int]:
    """Crop the reference audio to its loudest segment.

    Returns (cropped_audio, sample_rate).
    """
    duration = len(audio) / sr

    # Very short audio: skip cropping
    if duration <= config.min_segment:
        logger.info("Reference audio too short (%.1fs). Using full audio.", duration)
        return audio, sr

    envelope = compute_lufs_envelope(audio, sr)
    segment = find_loudest_segment(envelope, hop_sec=0.1, config=config)

    if segment is None:
        # Flat track or no suitable segment
        if duration <= config.max_segment:
            return audio, sr
        # Return centered chunk
        center = duration / 2
        start = center - config.max_segment / 2
        end = center + config.max_segment / 2
        start_sample = int(start * sr)
        end_sample = int(end * sr)
        logger.info("Using centered %.1fs chunk of flat reference.", config.max_segment)
        return audio[start_sample:end_sample], sr

    start_sec, end_sec = segment

    # Beat-align
    start_sec, end_sec = snap_to_beats(start_sec, end_sec, audio, sr)

    # Ensure bounds respect min/max segment after beat alignment
    duration_seg = end_sec - start_sec
    if duration_seg < config.min_segment:
        # Expand symmetrically
        deficit = config.min_segment - duration_seg
        start_sec = max(0, start_sec - deficit / 2)
        end_sec = min(duration, end_sec + deficit / 2)
    elif duration_seg > config.max_segment:
        center = (start_sec + end_sec) / 2
        start_sec = max(0, center - config.max_segment / 2)
        end_sec = min(duration, center + config.max_segment / 2)

    start_sample = max(0, int(start_sec * sr))
    end_sample = min(len(audio), int(end_sec * sr))

    cropped = audio[start_sample:end_sample]
    logger.info("Cropped reference: %.1fs - %.1fs (%.1fs)", start_sec, end_sec, (end_sample - start_sample) / sr)
    return cropped, sr
