"""
Shot-boundary-aware clip extraction.

Uses .scenes.txt files from SHOTS_DIR. Each line is: start_frame end_frame
Cf. https://github.com/dbamman/movie-representation/tree/main/pipeline
"""

import os
import subprocess
import tempfile
from pathlib import Path

from util.paths import SHOTS_DIR


def get_scenes_path(video_file_id):
    """
    Look for {video_file_id}.scenes.txt in SHOTS_DIR.
    Returns Path or None if not found.
    """
    p = SHOTS_DIR / f'{video_file_id}.scenes.txt'
    if p.exists():
        return p
    return None


def load_shot_boundaries(video_file_id, fps):
    """
    Load shot boundaries for a video.
    Returns list of (start_sec, end_sec) tuples, or None if no scenes file.
    """
    scenes_path = get_scenes_path(video_file_id)
    if scenes_path is None:
        return None

    shots = []
    with open(scenes_path) as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 2:
                start_frame = int(parts[0])
                end_frame = int(parts[1])
                shots.append((start_frame / fps, end_frame / fps))
    return shots


def group_shots_into_chunks(shots, target_duration=60.0, min_duration=10.0):
    """
    Group consecutive shots into chunks of approximately target_duration seconds.
    Never cuts mid-shot. Short trailing shots are merged into the previous chunk.

    Returns list of (start_sec, end_sec, shot_indices) tuples.
    """
    if not shots:
        return []

    chunks = []
    current_start = None
    current_end = None
    current_shots = []

    for i, (s_start, s_end) in enumerate(shots):
        if current_start is None:
            current_start = s_start
            current_end = s_end
            current_shots = [i]
        else:
            chunk_duration = s_end - current_start
            if chunk_duration > target_duration and (current_end - current_start) >= min_duration:
                chunks.append((current_start, current_end, list(current_shots)))
                current_start = s_start
                current_end = s_end
                current_shots = [i]
            else:
                current_end = s_end
                current_shots.append(i)

    if current_start is not None:
        chunks.append((current_start, current_end, list(current_shots)))

    return chunks


def get_chunks_for_video(video_file_id, fps, duration_sec, mode='shot_grouped', target_duration=60.0):
    """
    Get (start_sec, end_sec) chunks for a video.

    mode:
        'shot_grouped' — group shots into ~target_duration chunks (default)
        'shot' — each shot is its own chunk
        'fixed' — fixed target_duration intervals (ignores shot boundaries)

    Returns list of (start_sec, end_sec) tuples.
    """
    if mode == 'fixed':
        import math
        n_chunks = int(math.ceil(duration_sec / target_duration))
        return [
            (i * target_duration, min((i + 1) * target_duration, duration_sec))
            for i in range(n_chunks)
        ]

    # shot_grouped (and future shot-based modes) require scenes files — fail loudly
    # if missing rather than silently falling back to 'fixed'. A silent fallback
    # causes shot_grouped and fixed runs to be indistinguishable in the outputs,
    # invalidating any comparison between the two chunk modes.
    shots = load_shot_boundaries(video_file_id, fps)
    if shots is None:
        scenes_path = get_scenes_path(video_file_id)
        raise FileNotFoundError(
            f"mode={mode!r} requires a .scenes.txt file for '{video_file_id}' "
            f'but none was found. Set the CHNB_SHOTS_DIR env var to the directory '
            f'containing scenes files, or use mode="fixed" if no scene detection is available. '
            f'(Looked up via get_scenes_path() -> {scenes_path})'
        )

    if mode == 'shot_grouped':
        grouped = group_shots_into_chunks(shots, target_duration=target_duration)
        # Clamp to actual video duration and drop chunks that become empty/tiny.
        # Prevents ffmpeg from producing header-only mp4s when scenes.txt has
        # phantom boundaries past end-of-video.
        MIN_CHUNK_SEC = 2.0   # anything shorter than this is useless
        clamped = []
        for s, e, _ in grouped:
            if s >= duration_sec:
                continue  # phantom chunk past end-of-video
            e = min(e, duration_sec)
            if (e - s) < MIN_CHUNK_SEC:
                continue
            clamped.append((s, e))
        return clamped
    else:
        raise ValueError(f'Unknown mode: {mode}')


_AUDIO_PROBE_CACHE: dict[str, bool] = {}


def _has_audio_stream(video_path) -> bool:
    """
    Return True iff ffprobe reports at least one audio stream in the file.
    Cached per-path — the underlying mp4 doesn't change during a run.
    """
    key = str(video_path)
    if key in _AUDIO_PROBE_CACHE:
        return _AUDIO_PROBE_CACHE[key]

    try:
        proc = subprocess.run(
            ['ffprobe', '-v', 'error', '-select_streams', 'a',
             '-show_entries', 'stream=codec_type', '-of', 'csv=p=0', key],
            capture_output=True, text=True, timeout=20,
        )
        has_audio = proc.returncode == 0 and 'audio' in (proc.stdout or '').lower()
    except (FileNotFoundError, subprocess.TimeoutExpired):
        # ffprobe missing or hung — assume audio exists; ffmpeg will handle it.
        has_audio = True
    _AUDIO_PROBE_CACHE[key] = has_audio
    return has_audio


# Videos where aac re-encoding failed despite ffprobe reporting an audio stream.
# We log the first time each one happens so the user can investigate, but
# subsequent clips from the same file don't spam the console.
_AUDIO_FALLBACK_WARNED: set[str] = set()


def extract_clip(video_path, start_sec, end_sec, output_path=None, with_audio=True):
    """
    Extract a clip from video using ffmpeg.
    Returns path to the extracted clip.

    Audio handling:
      - If the source has no audio stream (ffprobe), use `-an` silently.
      - If the source has audio, re-encode to aac. If that fails, alert loudly
        once per source file and fall back to `-an` so the clip still lands;
        subsequent clips from the same file use `-an` without re-warning.
    """
    duration = end_sec - start_sec
    if output_path is None:
        tmp = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
        tmp.close()
        output_path = tmp.name

    src_key = str(video_path)
    source_has_audio = _has_audio_stream(video_path) if with_audio else False
    # Previously-failed sources skip straight to -an (no repeat warnings).
    want_audio = with_audio and source_has_audio and src_key not in _AUDIO_FALLBACK_WARNED

    def _build_cmd(audio_mode):
        c = [
            'ffmpeg', '-y', '-ss', str(start_sec), '-t', str(duration),
            '-i', str(video_path),
            '-map', '0:v:0', '-c:v', 'libx264', '-preset', 'ultrafast', '-q:v', '28',
        ]
        if audio_mode == 'aac':
            c.extend(['-map', '0:a:0', '-c:a', 'aac', '-b:a', '128k'])
        else:  # 'none'
            c.append('-an')
        c.append(output_path)
        return c

    proc = subprocess.run(_build_cmd('aac' if want_audio else 'none'),
                          capture_output=True, text=True)
    if proc.returncode != 0:
        if want_audio:
            # Audio track exists but ffmpeg won't re-encode it — alert once
            # and fall back to -an so downstream doesn't block.
            stderr_tail = (proc.stderr or '').strip().splitlines()[-6:]
            print(f'  [audio-warn] aac re-encode failed for {video_path}; '
                  f'falling back to -an. Last ffmpeg stderr:')
            for line in stderr_tail:
                print(f'      {line}')
            _AUDIO_FALLBACK_WARNED.add(src_key)

            proc2 = subprocess.run(_build_cmd('none'), capture_output=True, text=True)
            if proc2.returncode != 0:
                stderr = (proc2.stderr or '')[-2000:]
                raise RuntimeError(
                    f'ffmpeg failed even with -an for {video_path} '
                    f'[{start_sec:.2f}-{end_sec:.2f}s]:\n{stderr}'
                )
        else:
            stderr = (proc.stderr or '')[-2000:]
            raise RuntimeError(
                f'ffmpeg failed for {video_path} '
                f'[{start_sec:.2f}-{end_sec:.2f}s]:\n{stderr}'
            )

    # Sanity check: ffmpeg sometimes produces a header-only mp4 (few hundred bytes,
    # zero-duration) when start_sec is past video end or ffmpeg has a weird seek
    # failure. Delete and raise so callers can retry or skip instead of silently
    # writing a broken file.
    try:
        size = os.path.getsize(output_path)
    except OSError:
        size = 0
    if size < 10_000:   # real 60s clips are megabytes
        try:
            os.unlink(output_path)
        except OSError:
            pass
        raise RuntimeError(
            f'extract_clip produced suspiciously small file '
            f'({size} bytes) for {video_path} [{start_sec:.2f}-{end_sec:.2f}s]; '
            f'likely a scenes.txt boundary past end-of-video.'
        )

    return output_path


def extract_all_clips(video_file_id, video_path, fps, duration_sec,
                      chunk_mode='fixed', with_audio=True):
    """
    Extract all clips for a video.
    Idempotent — skips clips that already exist.
    Returns list of (clip_path, start_sec, end_sec) tuples.
    """
    from util.paths import get_clips_dir

    clips_dir = get_clips_dir(video_file_id, chunk_mode)
    clips_dir.mkdir(parents=True, exist_ok=True)
    chunks = get_chunks_for_video(video_file_id, fps, duration_sec,
                                  mode=chunk_mode)

    results = []
    for i, (start_sec, end_sec) in enumerate(chunks):
        clip_path = clips_dir / f'{i:04d}.mp4'
        if not clip_path.exists():
            extract_clip(str(video_path), start_sec, end_sec,
                        output_path=str(clip_path), with_audio=with_audio)
        results.append((str(clip_path), start_sec, end_sec))

    # Save a manifest with timestamps
    manifest_path = clips_dir / 'manifest.json'
    if not manifest_path.exists():
        import json
        manifest = [
            {'clip': f'{i:04d}.mp4', 'start': s, 'end': e}
            for i, (s, e) in enumerate(chunks)
        ]
        with open(manifest_path, 'w') as f:
            json.dump(manifest, f, indent=2)

    return results


def load_clips(video_file_id, chunk_mode='fixed', min_bytes=10_000, verbose=False):
    """
    Load cached clips.
    Returns list of (clip_path, start_sec, end_sec).
    Set `verbose=True` to print which entries were skipped and why.
    """
    import json as _json
    from util.paths import get_clips_dir

    clips_dir = get_clips_dir(video_file_id, chunk_mode)
    manifest_path = clips_dir / 'manifest.json'
    if not manifest_path.exists():
        return []

    with open(manifest_path) as f:
        manifest = _json.load(f)

    valid, dropped = [], []
    for m in manifest:
        clip_path = clips_dir / m['clip']
        try:
            size = clip_path.stat().st_size if clip_path.exists() else 0
        except OSError:
            size = 0
        if size >= min_bytes:
            valid.append((str(clip_path), m['start'], m['end']))
        else:
            dropped.append((m['clip'], size))

    if verbose and dropped:
        print(f'[load_clips] {video_file_id}/{chunk_mode}: skipped {len(dropped)} invalid entries '
              f'(first few: {dropped[:3]})')

    return valid