"""Beat Alignment Score (BAS) — EDGE / CVPR 2023.

Measures synchronisation between audio beats and motion beats in
dance/music videos.  Audio beats are detected via onset-strength peaks;
motion beats are detected as local minima of joint velocity computed from
dense optical flow.  The alignment score is the fraction of audio beats
that have a nearby motion beat within a tolerance window.

Backend tiers:
  1. **librosa** — proper onset-strength envelope for beat detection
  2. **heuristic** — RMS-energy peak detection as beat proxy

bas_score — higher = better alignment (0-1)
Returns None when no audio stream is present.
"""

import logging
import subprocess
import tempfile
from pathlib import Path
from typing import Optional, List

import cv2
import numpy as np

from ayase.models import Sample, QualityMetrics
from ayase.pipeline import PipelineModule

logger = logging.getLogger(__name__)


def _extract_audio_to_wav(video_path: str, wav_path: str) -> bool:
    """Use ffmpeg to extract audio track to a mono 22050 Hz WAV file."""
    try:
        result = subprocess.run(
            [
                "ffmpeg", "-y", "-i", video_path,
                "-vn", "-ac", "1", "-ar", "22050", "-f", "wav", wav_path,
            ],
            capture_output=True, timeout=60,
        )
        return result.returncode == 0 and Path(wav_path).stat().st_size > 44
    except Exception:
        return False


def _has_audio_stream(video_path: str) -> bool:
    """Check whether the video file contains an audio stream."""
    try:
        result = subprocess.run(
            [
                "ffprobe", "-v", "error", "-select_streams", "a",
                "-show_entries", "stream=codec_type", "-of", "csv=p=0",
                video_path,
            ],
            capture_output=True, text=True, timeout=10,
        )
        return "audio" in result.stdout.lower()
    except Exception:
        return False


def _detect_audio_beats_librosa(wav_path: str, sr: int = 22050) -> np.ndarray:
    """Detect audio beat times (seconds) using librosa onset strength."""
    import librosa

    y, sr = librosa.load(wav_path, sr=sr, mono=True)
    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    # Pick peaks in onset envelope
    peaks = librosa.util.peak_pick(
        onset_env, pre_max=3, post_max=3, pre_avg=3, post_avg=5, delta=0.5, wait=10
    )
    times = librosa.frames_to_time(peaks, sr=sr)
    return times


def _detect_audio_beats_heuristic(wav_path: str) -> np.ndarray:
    """Fallback beat detection using RMS energy peaks from raw WAV."""
    import wave
    import struct

    with wave.open(wav_path, "rb") as wf:
        n_frames = wf.getnframes()
        sr = wf.getframerate()
        raw = wf.readframes(n_frames)

    # Convert to float samples
    n_samples = len(raw) // 2
    samples = np.array(struct.unpack(f"<{n_samples}h", raw[:n_samples * 2]), dtype=np.float32)
    samples /= 32768.0

    # Compute RMS in ~23ms windows (hop ~512 at 22050)
    hop = max(1, sr // 43)
    window = hop * 2
    n_windows = max(1, (len(samples) - window) // hop)

    rms = np.zeros(n_windows)
    for i in range(n_windows):
        start = i * hop
        chunk = samples[start:start + window]
        rms[i] = np.sqrt(np.mean(chunk ** 2) + 1e-10)

    # Normalise
    rms_max = rms.max()
    if rms_max > 0:
        rms /= rms_max

    # Peak picking: local maxima above threshold
    threshold = np.mean(rms) + 0.5 * np.std(rms)
    beats: List[float] = []
    min_gap = int(0.2 * sr / hop)  # At least 200ms between beats

    for i in range(1, len(rms) - 1):
        if rms[i] > rms[i - 1] and rms[i] > rms[i + 1] and rms[i] > threshold:
            t = i * hop / sr
            if not beats or (t - beats[-1]) > (min_gap * hop / sr):
                beats.append(t)

    return np.array(beats)


def _detect_motion_beats(video_path: str, subsample: int = 2) -> np.ndarray:
    """Detect motion beat times as local minima of optical-flow kinetic energy."""
    cap = cv2.VideoCapture(video_path)
    try:
        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total < 4:
            return np.array([])

        # Read frames at subsample rate
        indices = list(range(0, total, max(1, subsample)))
        prev_gray = None
        kinetic_energy: List[float] = []
        frame_times: List[float] = []

        for idx in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
            ret, frame = cap.read()
            if not ret:
                continue

            # Resize for speed
            small = cv2.resize(frame, (160, 120))
            gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)

            if prev_gray is not None:
                flow = cv2.calcOpticalFlowFarneback(
                    prev_gray, gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
                )
                mag = np.sqrt(flow[..., 0] ** 2 + flow[..., 1] ** 2)
                kinetic_energy.append(float(np.mean(mag)))
                frame_times.append(idx / fps)

            prev_gray = gray
    finally:
        cap.release()

    if len(kinetic_energy) < 3:
        return np.array([])

    ke = np.array(kinetic_energy)
    times = np.array(frame_times)

    # Motion beats = local minima of kinetic energy (body pauses at beat)
    motion_beats: List[float] = []
    for i in range(1, len(ke) - 1):
        if ke[i] < ke[i - 1] and ke[i] < ke[i + 1]:
            motion_beats.append(times[i])

    return np.array(motion_beats)


def _compute_alignment(audio_beats: np.ndarray, motion_beats: np.ndarray,
                        tolerance: float = 0.1) -> float:
    """Fraction of audio beats aligned with a motion beat within tolerance (seconds)."""
    if len(audio_beats) == 0:
        return 0.0
    if len(motion_beats) == 0:
        return 0.0

    aligned = 0
    for ab in audio_beats:
        dists = np.abs(motion_beats - ab)
        if np.min(dists) <= tolerance:
            aligned += 1

    return aligned / len(audio_beats)


class BeatAlignmentModule(PipelineModule):
    name = "beat_alignment"
    description = "BAS beat alignment score — audio-motion sync (EDGE/CVPR 2023)"
    default_config = {
        "tolerance": 0.1,   # seconds — alignment window
        "subsample": 2,     # frame subsample rate for flow
    }

    def __init__(self, config=None):
        super().__init__(config)
        self._model = None
        self._backend = "heuristic"
        self._librosa_available = False

    def setup(self) -> None:
        # Tier 1: librosa for proper onset-strength beat detection
        try:
            import librosa  # noqa: F401
            self._librosa_available = True
            self._backend = "librosa"
            logger.info("BeatAlignment initialised (librosa backend)")
            return
        except ImportError:
            pass

        # Tier 2: heuristic RMS-energy peaks
        self._backend = "heuristic"
        logger.info(
            "BeatAlignment initialised (heuristic) — "
            "install librosa for proper beat detection"
        )

    def process(self, sample: Sample) -> Sample:
        if not sample.is_video:
            return sample

        # Check for audio stream
        if not _has_audio_stream(str(sample.path)):
            logger.debug("BeatAlignment: no audio in %s, skipping", sample.path.name)
            return sample

        try:
            score = self._compute_bas(sample)
            if score is not None:
                if sample.quality_metrics is None:
                    sample.quality_metrics = QualityMetrics()
                sample.quality_metrics.bas_score = score
        except Exception as e:
            logger.warning("BeatAlignment failed for %s: %s", sample.path, e)

        return sample

    def _compute_bas(self, sample: Sample) -> Optional[float]:
        """Compute Beat Alignment Score."""
        tolerance = self.config.get("tolerance", 0.1)
        subsample = self.config.get("subsample", 2)

        # Extract audio to temp WAV
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            wav_path = tmp.name

        try:
            if not _extract_audio_to_wav(str(sample.path), wav_path):
                return None

            # Detect audio beats
            if self._backend == "librosa":
                audio_beats = _detect_audio_beats_librosa(wav_path)
            else:
                audio_beats = _detect_audio_beats_heuristic(wav_path)

            if len(audio_beats) == 0:
                return None

            # Detect motion beats from video
            motion_beats = _detect_motion_beats(str(sample.path), subsample=subsample)

            # Compute alignment
            score = _compute_alignment(audio_beats, motion_beats, tolerance=tolerance)
            return float(np.clip(score, 0.0, 1.0))
        finally:
            try:
                Path(wav_path).unlink(missing_ok=True)
            except Exception:
                pass
