"""Audio analysis module using essentia (primary) and librosa (supplement).

Provides comprehensive audio feature extraction including BPM, key, energy,
danceability, valence, LUFS loudness, mood tags, genre suggestions, sections,
and time signature estimation.

Target: Python 3.12, macOS arm64 (dev), Linux x86_64 (prod).
"""

from __future__ import annotations

import logging
import os
from dataclasses import asdict, dataclass, field
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)

SAMPLE_RATE = 44100
SUPPORTED_EXTENSIONS = {".wav", ".mp3", ".flac", ".ogg", ".aiff", ".aif", ".m4a"}

VALID_MOOD_TAGS = [
    "energetic", "calm", "dark", "bright", "melancholic",
    "aggressive", "uplifting", "dreamy", "tense", "playful",
]

GENRE_PROFILES: dict[str, dict[str, tuple[float, float]]] = {
    "techno":        {"bpm": (120, 150), "centroid": (2000, 6000), "onset_rate": (4, 12), "energy": (0.6, 1.0)},
    "house":         {"bpm": (118, 135), "centroid": (1500, 5000), "onset_rate": (3, 8),  "energy": (0.5, 0.9)},
    "drum and bass": {"bpm": (160, 180), "centroid": (2500, 7000), "onset_rate": (6, 15), "energy": (0.7, 1.0)},
    "hip hop":       {"bpm": (70, 110),  "centroid": (1000, 4000), "onset_rate": (2, 7),  "energy": (0.3, 0.8)},
    "pop":           {"bpm": (100, 130), "centroid": (1500, 5000), "onset_rate": (2, 6),  "energy": (0.4, 0.8)},
    "rock":          {"bpm": (100, 145), "centroid": (2000, 6000), "onset_rate": (3, 8),  "energy": (0.5, 0.9)},
    "ambient":       {"bpm": (60, 100),  "centroid": (500, 2500),  "onset_rate": (0, 3),  "energy": (0.0, 0.4)},
    "classical":     {"bpm": (50, 140),  "centroid": (800, 4000),  "onset_rate": (1, 5),  "energy": (0.1, 0.7)},
    "jazz":          {"bpm": (80, 160),  "centroid": (1200, 5000), "onset_rate": (2, 8),  "energy": (0.2, 0.7)},
    "reggae":        {"bpm": (60, 90),   "centroid": (800, 3000),  "onset_rate": (2, 5),  "energy": (0.3, 0.6)},
    "metal":         {"bpm": (100, 200), "centroid": (3000, 8000), "onset_rate": (5, 15), "energy": (0.7, 1.0)},
    "r&b":           {"bpm": (60, 100),  "centroid": (1000, 4000), "onset_rate": (2, 6),  "energy": (0.3, 0.7)},
}


@dataclass
class SectionInfo:
    start_time: float
    duration: float
    key: str
    loudness: float
    energy: float


@dataclass
class AnalysisResult:
    bpm: float
    bpm_confidence: float
    key: str
    scale: str
    key_notation: str
    key_confidence: float
    energy: float
    danceability: float
    valence: float
    loudness_lufs: float
    mood_tags: list[str]
    genre_suggestions: list[str]
    duration_seconds: float
    time_signature: int
    sections: list[dict[str, Any]]

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


# ---------------------------------------------------------------------------
# Pool helper: safely read from essentia Pool
# ---------------------------------------------------------------------------

def _pool_get(pool, key: str, default=None):
    """Safely get a value from an essentia Pool object."""
    if pool is None:
        return default
    try:
        return pool[key]
    except (KeyError, RuntimeError):
        return default


# ---------------------------------------------------------------------------
# BPM extraction
# ---------------------------------------------------------------------------

def _extract_bpm(
    audio: np.ndarray, pool
) -> tuple[float, float, np.ndarray | None]:
    """Extract BPM, confidence, and beat tick positions."""
    import essentia.standard as es

    # Always use RhythmExtractor2013 directly for confidence + ticks
    try:
        rhythm = es.RhythmExtractor2013(method="multifeature")
        bpm, ticks, confidence, _estimates, _intervals = rhythm(audio)
        bpm = float(bpm)
        confidence = float(np.clip(confidence, 0.0, 1.0))
        if bpm <= 0:
            raise ValueError("BPM <= 0")
        return round(bpm, 1), round(confidence, 3), ticks
    except Exception as exc:
        logger.warning("RhythmExtractor2013 failed: %s", exc)

    # Fallback: pool value with no confidence
    pool_bpm = _pool_get(pool, "rhythm.bpm")
    if pool_bpm is not None and float(pool_bpm) > 0:
        return round(float(pool_bpm), 1), 0.0, None

    # Last resort
    return 120.0, 0.0, None


# ---------------------------------------------------------------------------
# Key extraction
# ---------------------------------------------------------------------------

def _extract_key(
    audio: np.ndarray, pool
) -> tuple[str, str, str, float]:
    """Extract key, scale, notation (e.g. 'Dm'), and confidence."""
    import essentia.standard as es

    try:
        key_ext = es.KeyExtractor()
        key, scale, strength = key_ext(audio)
        key = str(key)
        scale = str(scale)
        strength = float(np.clip(strength, 0.0, 1.0))
    except Exception as exc:
        logger.warning("KeyExtractor failed: %s, trying pool", exc)
        key = str(_pool_get(pool, "tonal.key_edma.key", "C"))
        scale = str(_pool_get(pool, "tonal.key_edma.scale", "major"))
        strength = float(_pool_get(pool, "tonal.key_edma.strength", 0.0))

    suffix = "m" if scale == "minor" else ""
    notation = f"{key}{suffix}"
    return key, scale, notation, round(strength, 3)


# ---------------------------------------------------------------------------
# Energy (0-1)
# ---------------------------------------------------------------------------

def _compute_energy(audio: np.ndarray, pool) -> float:
    """Compute energy as weighted blend of RMS energy and spectral centroid."""
    import librosa

    # RMS
    pool_loudness = _pool_get(pool, "lowlevel.average_loudness")
    if pool_loudness is not None:
        rms_val = float(pool_loudness)
    else:
        rms = librosa.feature.rms(y=audio)[0]
        rms_val = float(np.mean(rms))
    rms_norm = float(np.clip(rms_val / 0.3, 0.0, 1.0))

    # Spectral centroid
    centroid = _get_centroid(pool, audio)
    centroid_norm = float(np.clip((centroid - 500) / 7500, 0.0, 1.0))

    energy = 0.7 * rms_norm + 0.3 * centroid_norm
    return round(float(np.clip(energy, 0.0, 1.0)), 3)


# ---------------------------------------------------------------------------
# Danceability (0-1)
# ---------------------------------------------------------------------------

def _compute_danceability(
    audio: np.ndarray,
    pool,
    beat_ticks: np.ndarray | None,
) -> float:
    """Compute danceability from beat regularity, onset rate, essentia estimate."""
    # Essentia native danceability (arbitrary scale ~0-3)
    ess_dance = _pool_get(pool, "rhythm.danceability")
    if ess_dance is not None:
        dance_base = float(np.clip(float(ess_dance) / 3.0, 0.0, 1.0))
    else:
        dance_base = 0.5

    # Beat regularity via inter-beat interval coefficient of variation
    if beat_ticks is not None and len(beat_ticks) > 2:
        ibis = np.diff(beat_ticks)
        ibi_cv = float(np.std(ibis) / (np.mean(ibis) + 1e-8))
        beat_regularity = float(np.clip(1.0 - ibi_cv * 5, 0.0, 1.0))
    else:
        beat_regularity = 0.5

    # Onset rate
    onset_rate = float(_pool_get(pool, "rhythm.onset_rate") or 3.0)
    onset_factor = float(np.clip(onset_rate / 10.0, 0.0, 1.0))

    danceability = 0.4 * dance_base + 0.35 * beat_regularity + 0.25 * onset_factor
    return round(float(np.clip(danceability, 0.0, 1.0)), 3)


# ---------------------------------------------------------------------------
# Valence (0-1)
# ---------------------------------------------------------------------------

def _compute_valence(scale: str, pool) -> float:
    """Compute valence from mode, spectral brightness, and harmonic complexity."""
    mode_score = 0.65 if scale == "major" else 0.35

    centroid = float(_pool_get(pool, "lowlevel.spectral_centroid.mean") or 3000)
    brightness = float(np.clip((centroid - 1000) / 6000, 0.0, 1.0))

    dissonance = float(_pool_get(pool, "lowlevel.dissonance.mean") or 0.4)
    harmony = float(np.clip(1.0 - dissonance * 2, 0.0, 1.0))

    valence = 0.40 * mode_score + 0.30 * brightness + 0.30 * harmony
    return round(float(np.clip(valence, 0.0, 1.0)), 3)


# ---------------------------------------------------------------------------
# LUFS loudness
# ---------------------------------------------------------------------------

def _measure_loudness_lufs(file_path: str, pool) -> float:
    """Measure integrated loudness in LUFS using pyloudnorm."""
    try:
        import soundfile as sf
        import pyloudnorm as pyln

        data, rate = sf.read(file_path)
        meter = pyln.Meter(rate)
        lufs = meter.integrated_loudness(data)
        if np.isfinite(lufs):
            return round(float(lufs), 1)
    except Exception as exc:
        logger.warning("pyloudnorm LUFS measurement failed: %s", exc)

    # Fallback: essentia EBU R128
    ebu = _pool_get(pool, "lowlevel.loudness_ebu128.integrated")
    if ebu is not None and np.isfinite(float(ebu)):
        return round(float(ebu), 1)

    return float("-inf")


# ---------------------------------------------------------------------------
# Time signature estimation
# ---------------------------------------------------------------------------

def _estimate_time_signature(
    beat_ticks: np.ndarray | None, audio: np.ndarray
) -> int:
    """Estimate time signature by analyzing accent periodicity in beat onsets."""
    if beat_ticks is None or len(beat_ticks) < 8:
        return 4

    import librosa

    onset_env = librosa.onset.onset_strength(y=audio, sr=SAMPLE_RATE)
    hop = 512
    sr_onset = SAMPLE_RATE / hop

    beat_strengths = []
    for tick in beat_ticks:
        idx = int(tick * sr_onset)
        if 0 <= idx < len(onset_env):
            beat_strengths.append(onset_env[idx])

    if len(beat_strengths) < 8:
        return 4

    bs = np.array(beat_strengths, dtype=np.float64)
    bs -= np.mean(bs)
    autocorr = np.correlate(bs, bs, mode="full")
    autocorr = autocorr[len(autocorr) // 2:]
    if autocorr[0] > 0:
        autocorr /= autocorr[0]

    candidates = {2: 0.0, 3: 0.0, 4: 0.0, 6: 0.0}
    for period in candidates:
        if period < len(autocorr):
            candidates[period] = float(autocorr[period])

    return max(candidates, key=candidates.get)  # type: ignore[arg-type]


# ---------------------------------------------------------------------------
# Section detection
# ---------------------------------------------------------------------------

def _detect_sections(
    audio: np.ndarray, pool
) -> list[dict[str, Any]]:
    """Detect structural sections using SBic on MFCCs, with librosa fallback."""
    import essentia.standard as es

    duration = len(audio) / SAMPLE_RATE
    frame_size = 2048
    hop_size = 1024

    # Compute MFCC matrix for SBic
    try:
        w = es.Windowing(type="blackmanharris62")
        spectrum = es.Spectrum()
        mfcc_algo = es.MFCC(numberCoefficients=13)

        mfcc_frames: list[np.ndarray] = []
        for frame in es.FrameGenerator(audio, frameSize=frame_size, hopSize=hop_size):
            windowed = w(frame)
            spec = spectrum(windowed)
            _, coeffs = mfcc_algo(spec)
            mfcc_frames.append(coeffs)

        if len(mfcc_frames) < 20:
            raise ValueError("Audio too short for SBic segmentation")

        mfcc_matrix = np.array(mfcc_frames).T  # (13, n_frames)

        sbic = es.SBic(minLength=10, inc1=60, inc2=20, cpw=1.5)
        boundaries = sbic(mfcc_matrix)
        boundary_frames = [int(b) for b in boundaries]
    except Exception as exc:
        logger.warning("SBic segmentation failed: %s, using librosa fallback", exc)
        try:
            import librosa
            chroma = librosa.feature.chroma_cqt(y=audio, sr=SAMPLE_RATE, hop_length=hop_size)
            n_clusters = max(3, int(duration / 30))
            boundary_frames = librosa.segment.agglomerative(chroma, k=n_clusters).tolist()
        except Exception as exc2:
            logger.warning("librosa segmentation also failed: %s", exc2)
            return []

    # Convert frame indices → times
    boundary_times = sorted({0.0} | {b * hop_size / SAMPLE_RATE for b in boundary_frames} | {duration})

    # Per-section analysis
    import essentia.standard as es_std
    import librosa as lr

    sections: list[dict[str, Any]] = []
    for i in range(len(boundary_times) - 1):
        start = boundary_times[i]
        end = boundary_times[i + 1]
        seg_dur = end - start
        if seg_dur < 0.5:
            continue

        seg_audio = audio[int(start * SAMPLE_RATE): int(end * SAMPLE_RATE)]
        if len(seg_audio) < int(SAMPLE_RATE * 0.5):
            continue

        # Key
        try:
            sk, ss, _ = es_std.KeyExtractor()(seg_audio)
            sec_key = f"{sk}{'m' if ss == 'minor' else ''}"
        except Exception:
            sec_key = "N/A"

        # Energy (RMS)
        sec_rms = float(np.mean(lr.feature.rms(y=seg_audio)[0]))
        sec_energy = round(float(np.clip(sec_rms / 0.3, 0.0, 1.0)), 3)

        # Loudness (dB RMS, not LUFS — segments may be too short for pyloudnorm)
        sec_loudness = round(float(20.0 * np.log10(sec_rms + 1e-10)), 1)

        sections.append({
            "start_time": round(start, 2),
            "duration": round(seg_dur, 2),
            "key": sec_key,
            "loudness": sec_loudness,
            "energy": sec_energy,
        })

    return sections


# ---------------------------------------------------------------------------
# Feature helpers
# ---------------------------------------------------------------------------

def _get_centroid(pool, audio: np.ndarray) -> float:
    """Get spectral centroid mean from pool or librosa."""
    val = _pool_get(pool, "lowlevel.spectral_centroid.mean")
    if val is not None:
        return float(val)
    import librosa
    return float(np.mean(librosa.feature.spectral_centroid(y=audio, sr=SAMPLE_RATE)))


def _get_dissonance(pool) -> float:
    val = _pool_get(pool, "lowlevel.dissonance.mean")
    return float(val) if val is not None else 0.4


def _get_onset_rate(pool) -> float:
    val = _pool_get(pool, "rhythm.onset_rate")
    return float(val) if val is not None else 4.0


# ---------------------------------------------------------------------------
# Mood tag derivation
# ---------------------------------------------------------------------------

def _derive_mood_tags(
    bpm: float,
    energy: float,
    valence: float,
    scale: str,
    centroid: float,
    dissonance: float,
    onset_rate: float,
) -> list[str]:
    """Derive mood tags from low-level features using rule-based scoring."""
    is_minor = 1.0 if scale == "minor" else 0.0
    brightness = float(np.clip((centroid - 500) / 7500, 0.0, 1.0))
    bpm_norm = float(np.clip(bpm / 160, 0.0, 1.0))
    onset_norm = float(np.clip(onset_rate / 8, 0.0, 1.0))
    diss_norm = float(np.clip(dissonance * 2, 0.0, 1.0))

    scores: dict[str, float] = {}

    scores["energetic"] = 0.5 * energy + 0.3 * bpm_norm + 0.2 * onset_norm

    scores["calm"] = 0.5 * (1 - energy) + 0.3 * (1 - bpm_norm) + 0.2 * (1 - onset_norm)

    scores["dark"] = 0.35 * (1 - brightness) + 0.35 * is_minor + 0.30 * diss_norm

    scores["bright"] = 0.45 * brightness + 0.35 * (1 - is_minor) + 0.20 * valence

    scores["melancholic"] = 0.40 * is_minor + 0.30 * (1 - energy) + 0.30 * (1 - valence)

    scores["aggressive"] = (
        0.30 * energy
        + 0.25 * float(np.clip(bpm / 180, 0, 1))
        + 0.25 * float(np.clip(onset_rate / 10, 0, 1))
        + 0.20 * diss_norm
    )

    scores["uplifting"] = 0.35 * (1 - is_minor) + 0.35 * valence + 0.30 * energy

    scores["dreamy"] = (
        0.40 * (1 - energy)
        + 0.30 * (1 - float(np.clip(onset_rate / 6, 0, 1)))
        + 0.30 * (1 - float(np.clip(bpm / 120, 0, 1)))
    )

    scores["tense"] = (
        0.35 * float(np.clip(dissonance * 2.5, 0, 1))
        + 0.35 * is_minor
        + 0.30 * float(np.clip(energy * 1.3, 0, 1))
    )

    scores["playful"] = (
        0.30 * (1 - is_minor)
        + 0.30 * valence
        + 0.20 * float(np.clip((bpm - 100) / 60, 0, 1))
        + 0.20 * float(np.clip(energy, 0.3, 0.8))
    )

    # Select tags above threshold, sorted by score descending
    threshold = 0.55
    mood_tags = [
        tag for tag, score in sorted(scores.items(), key=lambda x: -x[1])
        if score >= threshold
    ]

    # Mutual exclusivity: remove contradictions
    contradictions = [("energetic", "calm"), ("dark", "bright")]
    for a, b in contradictions:
        if a in mood_tags and b in mood_tags:
            if scores[a] >= scores[b]:
                mood_tags.remove(b)
            else:
                mood_tags.remove(a)

    return mood_tags


# ---------------------------------------------------------------------------
# Genre suggestion
# ---------------------------------------------------------------------------

def _suggest_genres(
    bpm: float,
    centroid: float,
    onset_rate: float,
    energy: float,
    scale: str,
) -> list[str]:
    """Suggest top 3 genre tags by scoring feature fit against profiles."""
    genre_scores: dict[str, float] = {}

    for genre, profile in GENRE_PROFILES.items():
        score = 0.0

        # BPM fit
        bpm_lo, bpm_hi = profile["bpm"]
        bpm_mid = (bpm_lo + bpm_hi) / 2
        bpm_range = max((bpm_hi - bpm_lo) / 2, 1)
        bpm_fit = max(0.0, 1.0 - abs(bpm - bpm_mid) / (bpm_range * 1.5))
        score += bpm_fit * 0.35

        # Centroid fit
        c_lo, c_hi = profile["centroid"]
        c_mid = (c_lo + c_hi) / 2
        c_range = max((c_hi - c_lo) / 2, 1)
        c_fit = max(0.0, 1.0 - abs(centroid - c_mid) / (c_range * 1.5))
        score += c_fit * 0.25

        # Onset rate fit
        o_lo, o_hi = profile["onset_rate"]
        o_mid = (o_lo + o_hi) / 2
        o_range = max((o_hi - o_lo) / 2, 0.5)
        o_fit = max(0.0, 1.0 - abs(onset_rate - o_mid) / (o_range * 1.5))
        score += o_fit * 0.20

        # Energy fit
        e_lo, e_hi = profile["energy"]
        if e_lo <= energy <= e_hi:
            e_fit = 1.0
        else:
            dist = min(abs(energy - e_lo), abs(energy - e_hi))
            e_fit = max(0.0, 1.0 - dist * 3)
        score += e_fit * 0.20

        genre_scores[genre] = score

    sorted_genres = sorted(genre_scores.items(), key=lambda x: -x[1])
    return [g for g, _s in sorted_genres[:3]]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def analyze(file_path: str) -> dict[str, Any]:
    """Analyze an audio file and return a comprehensive analysis dict.

    Args:
        file_path: Path to audio file (.wav, .mp3, .flac, .ogg, .aiff, .m4a).

    Returns:
        Dict with keys: bpm, bpm_confidence, key, scale, key_notation,
        key_confidence, energy, danceability, valence, loudness_lufs,
        mood_tags, genre_suggestions, duration_seconds, time_signature,
        sections.

    Raises:
        FileNotFoundError: If file does not exist.
        ValueError: If file format is unsupported or audio is empty.
        RuntimeError: If audio loading fails (corrupt file, codec error).
    """
    import essentia.standard as es

    # -- Validate input --
    if not os.path.isfile(file_path):
        raise FileNotFoundError(f"Audio file not found: {file_path}")

    ext = os.path.splitext(file_path)[1].lower()
    if ext not in SUPPORTED_EXTENSIONS:
        raise ValueError(
            f"Unsupported audio format '{ext}'. "
            f"Supported: {sorted(SUPPORTED_EXTENSIONS)}"
        )

    # -- Load audio --
    try:
        audio = es.MonoLoader(filename=file_path, sampleRate=SAMPLE_RATE)()
    except Exception as exc:
        raise RuntimeError(f"Failed to load audio: {exc}") from exc

    if len(audio) == 0:
        raise ValueError("Audio file is empty (0 samples after loading).")

    duration_seconds = round(float(len(audio)) / SAMPLE_RATE, 2)

    # -- MusicExtractor for bulk feature extraction --
    pool = None
    try:
        extractor = es.MusicExtractor(
            lowlevelStats=["mean", "stdev"],
            rhythmStats=["mean", "stdev"],
            tonalStats=["mean", "stdev"],
            analysisSampleRate=SAMPLE_RATE,
        )
        pool, _pool_frames = extractor(file_path)
    except Exception as exc:
        logger.warning("MusicExtractor failed, using individual algorithms: %s", exc)

    # -- Individual feature extraction --
    bpm, bpm_confidence, beat_ticks = _extract_bpm(audio, pool)
    key, scale, key_notation, key_confidence = _extract_key(audio, pool)
    energy = _compute_energy(audio, pool)
    danceability = _compute_danceability(audio, pool, beat_ticks)
    valence = _compute_valence(scale, pool)
    loudness_lufs = _measure_loudness_lufs(file_path, pool)
    time_signature = _estimate_time_signature(beat_ticks, audio)
    sections = _detect_sections(audio, pool)

    # -- Derived features --
    centroid = _get_centroid(pool, audio)
    dissonance = _get_dissonance(pool)
    onset_rate = _get_onset_rate(pool)
    mood_tags = _derive_mood_tags(bpm, energy, valence, scale, centroid, dissonance, onset_rate)
    genre_suggestions = _suggest_genres(bpm, centroid, onset_rate, energy, scale)

    # -- Build result --
    result = AnalysisResult(
        bpm=bpm,
        bpm_confidence=bpm_confidence,
        key=key,
        scale=scale,
        key_notation=key_notation,
        key_confidence=key_confidence,
        energy=energy,
        danceability=danceability,
        valence=valence,
        loudness_lufs=loudness_lufs,
        mood_tags=mood_tags,
        genre_suggestions=genre_suggestions,
        duration_seconds=duration_seconds,
        time_signature=time_signature,
        sections=sections,
    )
    return result.to_dict()
