"""Audio analysis for Razzmatazz Kitmaker.

BPM, type classification, loop/one-shot, key, pitch, loudness, DC offset.
"""

import os
import tempfile

import librosa
import numpy as np


BPM_MIN = 60
BPM_MAX = 200

BANDS = {
    'sub':    (20,   60),
    'bass':   (60,   250),
    'lowmid': (250,  500),
    'mid':    (500,  2000),
    'highmid':(2000, 4000),
    'high':   (4000, 20000),
}

PITCH_CLASSES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']

KS_MAJOR = [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]
KS_MINOR = [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]

TYPE_SIGNATURE = [
    {'name': 'kick',  'sub': 0.3,  'bass': 0.3,  'lowmid': 0.1, 'high': 0.0,  'fast_attack': 1, 'noise': 0.0, 'pitchdrop': 0.3},
    {'name': 'snare', 'sub': 0.0,  'bass': 0.05, 'lowmid': 0.1, 'high': 0.05, 'fast_attack': 1, 'noise': 0.4, 'pitchdrop': 0.0},
    {'name': 'closedHat', 'sub': 0.0, 'bass': 0.0, 'lowmid': 0.0, 'high': 0.5, 'fast_attack': 1, 'noise': 0.6, 'pitchdrop': 0.0},
    {'name': 'openHat',   'sub': 0.0, 'bass': 0.0, 'lowmid': 0.0, 'high': 0.4, 'fast_attack': 0.5, 'noise': 0.6, 'pitchdrop': 0.0},
    {'name': 'clap', 'sub': 0.0,  'bass': 0.0,  'lowmid': 0.05, 'high': 0.1, 'fast_attack': 0, 'noise': 0.5, 'pitchdrop': 0.0},
    {'name': 'tom',  'sub': 0.1,  'bass': 0.2,  'lowmid': 0.2,  'high': 0.0,  'fast_attack': 1, 'noise': 0.0, 'pitchdrop': 0.5},
    {'name': 'perc', 'sub': 0.0,  'bass': 0.0,  'lowmid': 0.05, 'high': 0.2,  'fast_attack': 1, 'noise': 0.2, 'pitchdrop': 0.0},
    {'name': 'rim',  'sub': 0.0,  'bass': 0.0,  'lowmid': 0.0,  'high': 0.3,  'fast_attack': 1, 'noise': 0.25, 'pitchdrop': 0.0},
    {'name': 'shaker', 'sub': 0.0, 'bass': 0.0, 'lowmid': 0.0, 'high': 0.3, 'fast_attack': 0.2, 'noise': 0.5, 'pitchdrop': 0.0},
    {'name': 'crash',   'sub': 0.0, 'bass': 0.0, 'lowmid': 0.0, 'high': 0.3, 'fast_attack': 0.6, 'noise': 0.6, 'pitchdrop': 0.0},
    {'name': 'cowbell', 'sub': 0.0, 'bass': 0.0, 'lowmid': 0.2, 'high': 0.1, 'fast_attack': 1, 'noise': 0.1, 'pitchdrop': 0.0},
]

FN_PATTERNS = {
    'kick': ('kick', 'kik', 'bd ', 'bassdrum', 'bass drum'),
    'snare': ('snare', 'snr', ' sd '),
    'closedHat': ('hat', 'hh ', 'hihat', 'closed'),
    'openHat': ('open', 'ohh'),
    'clap': ('clap',),
    'tom': ('tom',),
    'perc': ('perc', 'tamb', 'clave', 'bell', 'shkr'),
    'rim': ('rim', 'rimshot'),
    'shaker': ('shaker', 'shake'),
    'crash': ('crash',),
    'cowbell': ('cowbell',),
    'ride': ('ride',),
}


def read_audio(data: bytes) -> tuple:
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
        f.write(data)
        inpath = f.name
    try:
        y, sr = librosa.load(inpath, sr=44100, mono=True, res_type='kaiser_fast')
        return y, sr
    finally:
        try:
            os.unlink(inpath)
        except OSError:
            pass


BPM_CLOSE = 5  # within 5 BPM = considered close for doubling/halving


def _readable_bpm(raw_bpm: float) -> int:
    """Clamp and round a raw BPM to a reasonable value."""
    if raw_bpm <= 0:
        return 0
    bpm = int(round(raw_bpm))
    # Only double if the raw BPM is unreasonably low for dance music
    if bpm < 70 and bpm * 2 <= BPM_MAX:
        return bpm * 2
    # Only halve if the raw BPM is unreasonably high
    if bpm > 180 and bpm / 2 >= BPM_MIN:
        return bpm // 2
    return max(BPM_MIN, min(BPM_MAX, bpm))


def detect_bpm(y: np.ndarray, sr: int) -> float:
    dur = len(y) / sr
    if dur < 0.5:
        return 0.0
    if dur < 2.0:
        return 0.0

    # Primary: onset-interval BPM — works well for rhythmic material
    onset_frames = librosa.onset.onset_detect(y=y, sr=sr, hop_length=512, backtrack=False)
    if len(onset_frames) >= 4:
        onset_times = librosa.frames_to_time(onset_frames, sr=sr, hop_length=512)
        intervals = np.diff(onset_times)
        # Remove outlier intervals
        med = np.median(intervals)
        mad = np.median(np.abs(intervals - med))
        if mad > 0:
            clean = intervals[np.abs(intervals - med) <= 2 * mad]
        else:
            clean = intervals
        if len(clean) >= 2:
            mean_interval = np.mean(clean)
            if mean_interval > 0:
                onset_bpm = 60.0 / mean_interval
                if BPM_MIN <= onset_bpm <= BPM_MAX:
                    return _readable_bpm(onset_bpm)

    # Fallback: librosa beat_track for full phrases
    tempo, _ = librosa.beat.beat_track(y=y, sr=sr, units='time')
    if tempo is None or np.all(tempo == 0):
        return 0.0
    bpm = float(np.asarray(tempo).flat[0])
    return _readable_bpm(bpm)


def detect_rms(y: np.ndarray):
    rms = np.sqrt(np.mean(y ** 2))
    if rms == 0:
        return None
    return round(20 * np.log10(rms), 1)


def detect_amplitude_envelope(y: np.ndarray, sr: int) -> list:
    if len(y) < 256:
        return []
    chunk = int(sr * 0.05)
    n = max(1, len(y) // chunk)
    rms = []
    for i in range(n):
        frame = y[i * chunk : (i + 1) * chunk]
        val = float(np.sqrt(np.mean(frame ** 2)))
        rms.append(val)
    mx = max(rms) if rms else 1
    return [round(v / mx, 4) for v in rms]


def detect_dc_offset(y: np.ndarray) -> float:
    return float(np.mean(y))


def detect_pitch(y: np.ndarray, sr: int) -> float:
    if len(y) < sr:
        return 0.0
    f0, voiced, _ = librosa.pyin(y, sr=sr, fmin=50, fmax=2000, fill_na=0.0)
    voiced_f0 = f0[voiced]
    if len(voiced_f0) == 0:
        return 0.0
    return round(float(np.median(voiced_f0)), 1)


def spectral_features(y: np.ndarray, sr: int) -> dict:
    S = np.abs(librosa.stft(y))
    freqs = librosa.fft_frequencies(sr=sr)

    if S.size == 0:
        return {
            'centroid': 0, 'rolloff': 0, 'bandwidth': 0,
            'band_energy': {}, 'zcr': 0, 'pitchdrop': 0,
        }

    total_energy = np.sum(S)
    if total_energy == 0:
        total_energy = 1

    centroid = float(np.mean(librosa.feature.spectral_centroid(S=S, sr=sr)))
    rolloff = float(np.mean(librosa.feature.spectral_rolloff(S=S, sr=sr)))
    bandwidth = float(np.mean(librosa.feature.spectral_bandwidth(S=S, sr=sr)))

    band_energy = {}
    for name, (lo, hi) in BANDS.items():
        mask = (freqs >= lo) & (freqs < hi)
        energy = np.sum(S[mask, :]) / total_energy
        band_energy[name] = float(energy)

    zcr = float(np.mean(librosa.feature.zero_crossing_rate(y)))

    half = len(y) // 2
    if half > 512:
        S_a = np.abs(librosa.stft(y[:half]))
        S_b = np.abs(librosa.stft(y[half:]))
        c_a = np.sum(freqs * np.mean(S_a, axis=1)) / max(np.sum(S_a), 1)
        min_len = min(S_a.shape[0], S_b.shape[0])
        c_b = np.sum(freqs[:min_len] * np.mean(S_b[:min_len], axis=1)) / max(np.sum(S_b[:min_len]), 1)
        pitchdrop = max(0, c_a - c_b) / max(c_a, 1)
    else:
        pitchdrop = 0.0

    onset_frames = librosa.onset.onset_detect(y=y, sr=sr, backtrack=True, units='samples')
    if len(onset_frames) > 0:
        onset = max(0, int(onset_frames[0]))
    else:
        onset = int(np.argmax(np.abs(y)))
    pre = max(0, onset - int(0.005 * sr))
    post = min(len(y), onset + int(0.020 * sr))
    energy_pre = np.sum(y[pre:onset] ** 2) if pre < onset else 1
    energy_post = np.sum(y[onset:post] ** 2)
    fast_attack = float(energy_post / max(energy_pre + energy_post, 1))

    return {
        'centroid': centroid,
        'rolloff': rolloff,
        'bandwidth': bandwidth,
        'band_energy': band_energy,
        'zcr': zcr,
        'fast_attack': fast_attack,
        'pitchdrop': float(pitchdrop),
    }


def classify_type(features: dict, filename: str = '') -> list:
    be = features['band_energy']
    fn_lower = filename.lower()

    fn_boost = {}
    for sig in TYPE_SIGNATURE:
        patterns = FN_PATTERNS.get(sig['name'], ())
        fn_boost[sig['name']] = 1 if any(p in fn_lower for p in patterns) else 0

    scores = []
    for sig in TYPE_SIGNATURE:
        score = 0.0
        for band in ('sub', 'bass', 'lowmid', 'mid', 'highmid', 'high'):
            actual = be.get(band, 0)
            expected = sig.get(band, 0)
            score -= abs(actual - expected) * 2
        noise = features['zcr']
        score -= abs(noise - sig['noise']) * 5
        score -= abs(features['fast_attack'] - sig['fast_attack']) * 3
        score -= abs(features['pitchdrop'] - sig['pitchdrop']) * 3

        if sig['name'] == 'clap' and features['fast_attack'] > 0.3:
            score -= features['fast_attack'] * 6

        score += fn_boost.get(sig['name'], 0) * 4
        scores.append((score, sig['name']))
    scores.sort(reverse=True)

    top_score = scores[0][0]
    threshold = 4.0
    results = [name for sc, name in scores if sc >= top_score - threshold]
    return results if results else [scores[0][1]]


def detect_key(y: np.ndarray, sr: int, loop_oneshot: str = '',
               features: dict | None = None) -> str:
    dur = len(y) / sr
    if dur < 2.0:
        return ''
    if loop_oneshot == 'oneshot' and dur < 4.0:
        return ''
    # Percussive sounds (low spectral centroid, high bandwidth = noise-like)
    # can't have meaningful key detection — skip early
    if features:
        cent = features.get('centroid', 0)
        bw = features.get('bandwidth', 0)
        if cent < 500 or (cent < 2000 and bw > cent * 0.8):
            return ''
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr, n_chroma=12)
    chroma_mean = np.mean(chroma, axis=1)
    total_energy = float(np.sum(chroma_mean))
    if total_energy < 2.0:
        return ''
    chroma_norm = chroma_mean / np.max(chroma_mean)
    # Reject near-uniform chroma profiles (percussive broadband energy)
    chroma_range = float(np.percentile(chroma_norm, 90) - np.percentile(chroma_norm, 10))
    if chroma_range < 0.3:
        return ''
    best_score = -np.inf
    best_key = ''
    for i in range(12):
        score_major = float(np.correlate(chroma_norm, np.roll(KS_MAJOR, i))[0])
        if score_major > best_score:
            best_score = score_major
            best_key = PITCH_CLASSES[i]
        score_minor = float(np.correlate(chroma_norm, np.roll(KS_MINOR, i))[0])
        if score_minor > best_score:
            best_score = score_minor
            best_key = PITCH_CLASSES[i] + 'm'
    if best_score < 3.5:
        return ''
    return best_key


def detect_loop_oneshot(y: np.ndarray, sr: int) -> str:
    dur = len(y) / sr

    # ── 1. Spectral continuity: do start and end sound the same? ──────
    # A seamless loop has similar spectral content at its boundaries.
    # Compare the first and last 512 samples via FFT magnitude correlation.
    win = min(512, len(y) // 4)
    if win >= 128:
        head = y[:win]
        tail = y[-win:]
        head_mag = np.abs(np.fft.rfft(head))
        tail_mag = np.abs(np.fft.rfft(tail))
        head_mag = head_mag / max(np.sum(head_mag), 1e-9)
        tail_mag = tail_mag / max(np.sum(tail_mag), 1e-9)
        continuity = float(np.sum(head_mag * tail_mag))  # cosine-like similarity
    else:
        continuity = 0.0

    # ── 2. Onset regularity: loops have rhythmic patterns ─────────────
    onset_frames = librosa.onset.onset_detect(y=y, sr=sr, backtrack=False)
    regular = False
    if len(onset_frames) >= 4:
        intervals = np.diff(onset_frames)
        if len(intervals) >= 2:
            cv = np.std(intervals) / max(np.mean(intervals), 1)
            regular = cv < 0.3

    # ── 3. Envelope symmetry: loops tend to not decay to silence ──────
    env = np.abs(y)
    env = np.convolve(env, np.hanning(256), mode='same')
    env = env / max(np.max(env), 1e-9)
    tail_energy = float(np.mean(env[-len(env)//4:]))
    head_energy = float(np.mean(env[:len(env)//4]))
    # In a one-shot, tail is much quieter than head. In a loop, they're similar.
    decay_ratio = tail_energy / max(head_energy, 1e-9)

    # ── Decision ──────────────────────────────────────────────────────
    # Strong continuity signal: end matches start → almost certainly a loop
    if continuity > 0.7 and dur > 0.5:
        return 'loop'

    # Regular onsets + low decay → rhythmic loop
    if regular and decay_ratio > 0.3:
        return 'loop'

    # Long sample with decent continuity and no massive decay
    if dur > 2.0 and continuity > 0.4 and decay_ratio > 0.2:
        return 'loop'

    return 'oneshot'


def _to_native(obj):
    """Recursively convert numpy types to native Python types."""
    if isinstance(obj, dict):
        return {k: _to_native(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_to_native(v) for v in obj]
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    return obj


def analyze_audio(input_bytes: bytes, ext: str, filename: str = '') -> dict:
    y, sr = read_audio(input_bytes)
    dur = len(y) / sr

    if len(y) < 256:
        return _to_native({
            'bpm': 0,
            'sample_type': guess_type_from_name(filename),
            'loop_oneshot': 'oneshot',
            'duration': dur,
            'key': '',
            'loudness_db': None,
            'dc_offset': 0.0,
            'pitch': 0.0,
            'error': 'Audio too short',
        })

    # Detect loop/oneshot FIRST — governs which features make sense
    loop_oneshot = detect_loop_oneshot(y, sr)

    # BPM: only for loops or long samples (> 4s)
    bpm_from_name = guess_bpm_from_name(filename)
    bpm_from_audio = 0.0
    bpm_confidence = 'unknown'
    if loop_oneshot == 'loop' or dur >= 4.0:
        bpm_from_audio = detect_bpm(y, sr)
        if bpm_from_name > 0:
            bpm = bpm_from_name
            bpm_confidence = 'filename'
        else:
            bpm = bpm_from_audio
            bpm_confidence = 'likely' if bpm > 0 else 'unknown'
    elif bpm_from_name > 0:
        bpm = bpm_from_name
        bpm_confidence = 'filename'
    else:
        bpm = 0

    features = spectral_features(y, sr)
    sample_type = classify_type(features, filename)

    # Key: only for loops or long samples, with loop_oneshot hint
    key = detect_key(y, sr, loop_oneshot, features)

    loudness_db = detect_rms(y)
    dc_offset = detect_dc_offset(y)
    pitch = detect_pitch(y, sr)
    amp_env = detect_amplitude_envelope(y, sr)

    return _to_native({
        'bpm': bpm,
        'sample_type': sample_type,
        'loop_oneshot': loop_oneshot,
        'duration': round(dur, 3),
        'key': key,
        'loudness_db': loudness_db,
        'dc_offset': round(dc_offset, 6),
        'pitch': pitch,
        'amp_env': amp_env,
        'confidence': {
            'bpm': bpm_confidence,
            'type': 'heuristic',
            'loop': 'heuristic',
        },
    })


def guess_bpm_from_name(name: str) -> float:
    n = name.lower()
    # "120bpm", "120_bpm", "bpm_120", "bpm120", "(120 bpm)"
    import re
    m = re.search(r'(\d{2,3})\s*[_\s]?bpm|bpm\s*[_\s]?(\d{2,3})', n)
    if m:
        val = int(m.group(1) or m.group(2))
        if BPM_MIN <= val <= BPM_MAX:
            return float(val)
    return 0.0


def guess_type_from_name(name: str) -> list:
    n = name.lower()
    if any(k in n for k in ('kick', 'kik', 'bd', 'bassdrum', 'bass drum', ' kick')):
        return ['kick']
    if any(s in n for s in ('snare', 'snr', ' sd ', ' snare', 'snare ')):
        return ['snare']
    if any(h in n for h in ('hihat', 'hi-hat', 'hi hat', 'hh ', 'closed', 'hat closed')):
        return ['closedHat']
    if any(o in n for o in ('open', 'ohh', 'crash', 'ride')):
        return ['openHat']
    if 'clap' in n:
        return ['clap']
    if any(t in n for t in ('tom', 'tom ')):
        return ['tom']
    if any(p in n for p in ('rim', 'rimshot', 'perc', 'cowbell', 'tamb', 'shaker', 'clave', 'bell', 'shkr')):
        return ['perc']
    if any(pat in n for pat in ('loop', 'slice', 'break', 'pattern', 'beat')):
        return ['slice']
    return ['tom']