from pathlib import Path
from typing import Tuple, List, Optional
from scipy.signal import correlate
import soundfile as sf
import numpy as np
import subprocess
import json
import tempfile
import os


def run(cmd):
    """
    Run a subprocess command.

    Args:
        cmd: Command and arguments as a list.

    Raises:
        subprocess.CalledProcessError: If the command exits with a non-zero status.
    """
    try:
        result = subprocess.run(
            cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )
        return result
    except subprocess.CalledProcessError as e:
        print(e.stdout)
        raise

def probe_video(path: Path):
    """Get video FPS, duration, codec, and stream start time."""
    import json
    cmd = [
        "ffprobe",
        "-v", "error",
        "-print_format", "json",
        "-show_streams",
        "-show_format",
        path
    ]
    data = json.loads(run(cmd).stdout)

    duration = float(data["format"]["duration"])
    for s in data["streams"]:
        if s["codec_type"] == "video":
            num, den = s["r_frame_rate"].split("/")
            fps = float(num) / float(den)
            codec = s["codec_name"]
            start_time = float(s.get("start_time", 0))
            return fps, duration, codec, start_time

    raise RuntimeError("No video stream found")

def _get_chapters(file: Path) -> List[Tuple[float, float]]:
    """Return list of (start_time, duration) for each chapter."""
    cmd = [
        "ffprobe",
        "-i", str(file),
        "-print_format", "json",
        "-show_chapters",
        "-loglevel", "error"
    ]

    result = run(cmd)
    data = json.loads(result.stdout)

    chapters = []
    for ch in data.get("chapters", []):
        start = float(ch["start_time"])
        end = float(ch["end_time"])
        duration = end - start
        chapters.append((start, duration))

    return chapters


def get_intro(file: Path, time: float) -> Tuple[float, float]:
    """
    Find the chapter whose duration is closest to `time`,
    considering only chapters with duration between time/2 and time*2.
    Returns (start_time, duration).
    """
    chapters = _get_chapters(file)

    if not chapters:
        raise ValueError("No chapters found")

    # Filter by allowed duration range
    min_len = time / 2
    max_len = time * 2
    filtered = [
        (start, dur) for start, dur in chapters
        if min_len <= dur <= max_len
    ]

    if not filtered:
        raise ValueError("No chapters match duration constraints")

    # Sort by closeness to target time
    filtered.sort(key=lambda x: abs(x[1] - time))

    # Take the two closest
    closest_two = filtered[:2]

    # Choose the better one (closest duration)
    best = min(closest_two, key=lambda x: abs(x[1] - time))

    return best


def extract_audio_segment(
    file: Path,
    intro: Tuple[float, float],
    audio_stream: Optional[int] = None
) -> Path:
    """
    Extract audio using (start, duration) from get_intro.

    Parameters:
        file: input video file
        intro: (start, duration)
        audio_stream: index of audio stream (0 = first audio track, etc.)

    Returns:
        Path to temporary audio file
    """
    start, duration = intro

    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp_path = Path(tmp.name)
    tmp.close()

    cmd = [
        "ffmpeg",
        "-y",
        "-i", str(file),
        "-ss", str(start),
        "-t", str(duration),
    ]

    # Select specific audio stream if provided
    if audio_stream is not None:
        cmd += ["-map", f"0:a:{audio_stream}"]
    else:
        cmd += ["-map", "0:a:0"]  # default: first audio track

    cmd += [
        "-vn",
        "-acodec", "pcm_s16le",
        "-ar", "44100",
        "-ac", "2",
        str(tmp_path)
    ]

    run(cmd)
    return tmp_path


def get_audio_streams(file: Path) -> List[int]:
    """
    Return a list of available audio stream indices (0-based for ffmpeg mapping).
    
    Example return:
        [0, 1]  -> corresponds to 0:a:0 and 0:a:1
    """
    cmd = [
        "ffprobe",
        "-i", str(file),
        "-select_streams", "a",
        "-show_entries", "stream=index",
        "-of", "json",
        "-loglevel", "error"
    ]

    result = run(cmd)
    data = json.loads(result.stdout)

    # ffprobe gives absolute stream indices, but we want relative audio indices
    streams = data.get("streams", [])

    # Convert to 0-based audio stream numbering
    return list(range(len(streams)))


def play_audio_segment(file: Path, intro: Tuple[float, float]) -> None:
    """
    Extract and play audio from (start, duration).
    """
    audio_file = extract_audio_segment(file, intro)

    try:
        run(
            ["ffplay", "-nodisp", "-autoexit", str(audio_file)]
        )
    finally:
        os.remove(audio_file)

def extract_audio(file: Path, duration: float):
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)

    run([
        "ffmpeg",
        "-y",
        "-i", str(file),
        "-map", "0:a:0",        # first audio stream
        "-ac", "1",             # mono
        "-ar", "44100",         # fixed sample rate
        "-t", str(duration),    # limit duration
        tmp.name
    ])

    data, sr = sf.read(tmp.name)
    return data, sr

def get_audio_offset(file1: Path, file2: Path, important_time: float):
    sig1, sr1 = extract_audio(file1, important_time)
    sig2, sr2 = extract_audio(file2, important_time)

    if sr1 != sr2:
        raise ValueError("Sample rates do not match")

    # normalize signals
    sig1 = sig1 - np.mean(sig1)
    sig2 = sig2 - np.mean(sig2)

    corr = correlate(sig1, sig2, mode="full")
    lag = np.argmax(corr) - (len(sig2) - 1)

    offset_seconds = lag / sr1
    return offset_seconds

def main(video: Path) -> bool:

    intro_duration = 90
    print("finding intro")
    intro = get_intro(video, 90.0)
    print("geting audio streams")
    audio_streams = get_audio_streams(video)
    print("extract first intro")
    extract_path1 = extract_audio_segment(video, intro, audio_streams[0])
    print("extract second intro")
    extract_path2 = extract_audio_segment(video, intro, audio_streams[-1])

    # # or directly play:
    # time = (0, 20)
    # play_audio_segment(extract_path1, time)
    # play_audio_segment(extract_path2, time)

    fps, _, _, _ = probe_video(video)

    print("Calculate")

    offset = get_audio_offset(extract_path1, extract_path2, intro[-1]-10)

    if np.abs(offset) < (1/fps):
        print(f"no offset: {offset}")
        return True
    else:
        print(f"NOT CORRECT!!!!!!!!!!!!!!!!!!!!!!!!: {offset}, {1/fps}")
        return False

extensions=("mp4", "mkv", "avi", "mov", "wmv", "flv", "webm", "lrv", "gif")

def check_video_ext(video: Path) -> bool:
    if video.suffix.strip(".").lower() in extensions:
        return True
    return False

if __name__ == '__main__':
    import sys
    # sys.argv.remove(sys.argv[0])
    if len(sys.argv) == 2:
        print("One file")
        if os.path.isfile(sys.argv[1]):
            main(Path(sys.argv[1]))
        else: raise
    elif len(sys.argv) > 2:
        sys.argv.remove(sys.argv[0])
        error_files = []
        for str_path in sys.argv:
            video = Path(str_path)
            if (check_video_ext(video)):
                print(f"{video}")
                working = main(video)
                if not working:
                    error_files.append(video)
            else: print(f"Not a compatabile Video: {video}")
        if error_files != []:
            print(f"These files are broken:")
            for file in error_files:
                print(file)
    else:
        print(f"Use this script like this:\npython detect_chapter.py /path/to/file")