#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["httpx"]
# ///
import argparse
import base64
import functools
import logging
import os
import re
import subprocess
import tempfile
import textwrap
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path

import httpx

RETRY_STATUS_CODES = {429, 500, 502, 503, 504, 524, 529}
MAX_ATTEMPTS = 5
MIN_SPLIT_DURATION = 1.0
MIN_CHUNK_DURATION = 0.1
DIALOGUE_LINE_LENGTH = 42
MAX_DIALOGUE_GAP = 1.0
MAX_DIALOGUE_DURATION = 7.0
LINE_LENGTH = 42
MAX_LINES = 2
# lowercase tokens (en, tr) whose trailing period does not end a sentence
ABBREVIATIONS = {"mr", "mrs", "ms", "dr", "prof", "e.g", "i.e", "vs", "etc", "doç", "av", "sn", "vb", "bkz"}
# Finds candidate sentence boundaries: end punctuation, optional closers, then whitespace.
# Matches in "He left.  Then" -> ".  " and in 'He said "no." Then' -> '." '.
SENTENCE_END = re.compile(
    r"""
    [.!?…]+         # one or more end marks: "." "?!" "..." or the ellipsis character
    [\"'”’)\]]*     # closing quotes or brackets that belong to the finished sentence
    \s+             # whitespace after it; one or more, so double spaces work too
    """,
    re.VERBOSE,
)

logger = logging.getLogger(__name__)


@dataclass
class Segment:
    """A timed piece of transcript; times are in seconds, speaker is a diarization index if known."""

    start: float
    end: float
    text: str
    speaker: int | None = None


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Transcribe a video into an SRT file using an OpenRouter speech-to-text model.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("video", type=Path, help="Input video (or audio) file")
    parser.add_argument("--output", "-o", type=Path, help="Output SRT path (default: <video>.srt)")
    parser.add_argument("--model", "-m", default="elevenlabs/scribe-v2", help="OpenRouter STT model id (must support diarization unless --no-diarize)")
    parser.add_argument(
        "--api-key",
        default=os.getenv("OPENROUTER_API_KEY"),
        help="OpenRouter API key (default: $OPENROUTER_API_KEY%(default).0s)",  # %(default).0s keeps the key out of --help
    )
    parser.add_argument("--language", "-l", help="ISO-639-1 source language, e.g. en or de (auto-detect if omitted)")
    parser.add_argument("--chunk-minutes", type=float, default=60.0, help="Audio chunk length in minutes (of sped-up audio)")
    parser.add_argument(
        "--speed-up",
        type=float,
        default=1.0,
        help="Speed the audio up (1.0-2.0) before transcribing to reduce cost; timestamps are scaled back and multi-sentence cues are split",
    )
    parser.add_argument(
        "--diarize",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Label speakers (needs a supporting model, e.g. elevenlabs/scribe-v2); short speaker changes become '- ' dialogue cues",
    )
    return parser.parse_args()


def run_ffmpeg_extract(video: Path, directory: Path, chunk_seconds: int, speed_up: float) -> list[Path]:
    """Extract mono 16 kHz mp3 audio into `directory` as chunks of `chunk_seconds` (measured after speed-up).

    Returns the chunk paths in order. `speed_up` > 1 applies a single ffmpeg `atempo` filter.
    """
    filter_arguments = ["-filter:a", f"atempo={speed_up}"] if speed_up != 1.0 else []
    subprocess.run(
        [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-i",
            str(video),
            "-vn",
            "-ac",
            "1",
            "-ar",
            "16000",
            *filter_arguments,
            "-c:a",
            "libmp3lame",
            "-b:a",
            "48k",
            "-f",
            "segment",
            "-segment_time",
            str(chunk_seconds),
            "-reset_timestamps",
            "1",
            str(directory / "chunk_%04d.mp3"),
        ],
        check=True,
    )
    chunks = sorted(directory.glob("chunk_*.mp3"))
    if not chunks:
        raise RuntimeError(f"ffmpeg produced no audio from {video}")
    return chunks


def get_duration(path: Path) -> float:
    """Return the media duration in seconds via ffprobe, or 0.0 if it is unreadable (e.g. an empty trailing chunk)."""
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(path)],
        check=True,
        capture_output=True,
        text=True,
    )
    try:
        return max(float(result.stdout.strip()), 0.0)
    except ValueError:
        return 0.0


def retry(
    attempts: int = MAX_ATTEMPTS,
    exceptions: tuple[type[Exception], ...] = (httpx.TransportError, httpx.HTTPStatusError),
) -> Callable:
    """Decorator factory: call the function up to `attempts` times, sleeping 2**attempt seconds between tries.

    Only `exceptions` trigger a retry; the last one is re-raised when attempts run out.
    """

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions as error:
                    if attempt == attempts:
                        raise
                    logger.warning("%s failed (%s), retry %d/%d", func.__name__, error, attempt, attempts)
                    time.sleep(2**attempt)

        return wrapper

    return decorator


@retry()
def post_request(client: httpx.Client, payload: dict) -> httpx.Response:
    """POST to the transcription endpoint. Retryable statuses (429/5xx) raise so `@retry` repeats them;
    other statuses, including 400, are returned for the caller to handle."""
    response = client.post("/audio/transcriptions", json=payload)
    if response.status_code in RETRY_STATUS_CODES:
        response.raise_for_status()
    return response


def transcribe_chunk(
    client: httpx.Client,
    path: Path,
    duration: float,
    model: str,
    language: str | None,
    diarize: bool,
) -> tuple[list[Segment], bool]:
    """Transcribe one audio chunk. Times in the result are relative to the chunk, in sped-up audio time.

    Returns (segments, has_timestamps). If the model rejects `verbose_json` the call is repeated as plain
    JSON and a single segment covering the whole chunk is returned with has_timestamps=False.
    With `diarize`, speaker turns are built from word-level labels, falling back to segment-level labels;
    a rejected diarize request exits instead of falling back.
    """
    payload: dict = {
        "model": model,
        "input_audio": {"data": base64.b64encode(path.read_bytes()).decode(), "format": "mp3"},
        "response_format": "verbose_json",
        "timestamp_granularities": ["segment", "word"] if diarize else ["segment"],
    }
    if language:
        payload["language"] = language
    if diarize:
        payload["diarize"] = True
    response = post_request(client=client, payload=payload)
    if response.status_code == 400 and diarize:
        raise SystemExit(f"Model rejected diarization: {response.text[:300]}")
    if response.status_code == 400:
        logger.warning("Model rejected verbose_json (%s); falling back to a single cue per chunk", response.text[:200])
        payload.pop("response_format")
        payload.pop("timestamp_granularities")
        response = post_request(client=client, payload=payload)
    response.raise_for_status()
    data = response.json()
    if not data.get("segments") and not data.get("words") and not str(data.get("text", "")).strip():
        raise RuntimeError(f"Empty transcription for {path.name} ({path.stat().st_size / 1e6:.1f} MB); response: {str(data)[:300]}")
    if diarize:
        turns = turns_from_words(words=data.get("words") or [])
        if turns:
            return turns, True
    segments = [
        Segment(
            start=float(segment["start"]),
            end=float(segment["end"]),
            text=str(segment["text"]).strip(),
            speaker=segment.get("speaker") if diarize else None,
        )
        for segment in data.get("segments") or []
        if str(segment.get("text", "")).strip()
    ]
    if diarize and not any(segment.speaker is not None for segment in segments):
        logger.warning("No speaker labels in the response; the model may not support diarization")
    if segments:
        return segments, True
    text = str(data.get("text", "")).strip()
    return ([Segment(start=0.0, end=duration, text=text)] if text else []), False


def turns_from_words(words: list[dict]) -> list[Segment]:
    """Group consecutive words with the same `speaker` into one Segment per speaker turn.

    Words without a speaker and `audio_event` entries (e.g. laughter) are skipped.
    """
    turns: list[Segment] = []
    for word in words:
        text = str(word.get("word", "")).strip()
        speaker = word.get("speaker")
        if not text or speaker is None or word.get("type") == "audio_event":
            continue
        if turns and turns[-1].speaker == speaker:
            turns[-1].end = float(word["end"])
            turns[-1].text += f" {text}"
        else:
            turns.append(Segment(start=float(word["start"]), end=float(word["end"]), text=text, speaker=speaker))
    for turn in turns:
        # Words are joined with spaces, so punctuation tokens end up detached ("Hi ,"); glue them back:
        # \s+ = the stray whitespace, ([,.!?;:…]) = the punctuation mark, kept via \1.
        turn.text = re.sub(r"\s+([,.!?;:…])", r"\1", turn.text)
    return turns


def group_dialogue(cues: list[Segment]) -> list[tuple[float, float, str]]:
    """Turn segments into SRT cues, pairing adjacent turns by different speakers into '- a' / '- b' cues.

    A pair is formed only if each line is at most DIALOGUE_LINE_LENGTH characters, the gap between them is at
    most MAX_DIALOGUE_GAP seconds and the pair spans at most MAX_DIALOGUE_DURATION seconds.
    Everything else becomes its own cue. Returns (start, end, text) tuples.
    """
    grouped = []
    index = 0
    while index < len(cues):
        cue = cues[index]
        following = cues[index + 1] if index + 1 < len(cues) else None
        if (
            cue.speaker is not None
            and following is not None
            and following.speaker not in (None, cue.speaker)
            and max(len(cue.text), len(following.text)) <= DIALOGUE_LINE_LENGTH
            and following.start - cue.end <= MAX_DIALOGUE_GAP
            and following.end - cue.start <= MAX_DIALOGUE_DURATION
        ):
            grouped.append((cue.start, following.end, f"- {cue.text}\n- {following.text}"))
            index += 2
        else:
            grouped.append((cue.start, cue.end, cue.text))
            index += 1
    return grouped


def split_sentences(text: str) -> list[str]:
    """Split text into sentences at SENTENCE_END matches.

    A boundary is kept only if the next character starts a sentence (uppercase, digit or opening quote/bracket)
    and the word before a period is not a known abbreviation or a single-letter initial.
    Decimals such as "3.5" never match because no whitespace follows the period.
    """
    sentences = []
    start = 0
    for match in SENTENCE_END.finditer(text):
        following = text[match.end() : match.end() + 1]
        if not following or not (following.isupper() or following.isdigit() or following in "\"'“‘(["):
            continue
        punctuation = match.group().strip()
        # (\S+)$ = the last whitespace-delimited word before the punctuation, e.g. "Dr" in "... Dr."
        word = re.search(r"(\S+)$", text[start : match.start()])
        token = word.group(1).lower().lstrip("\"'“‘([") if word else ""
        if punctuation.startswith(".") and (token in ABBREVIATIONS or (len(token) == 1 and token.isalpha())):
            continue
        sentences.append(text[start : match.start() + len(match.group().rstrip())].strip())
        start = match.end()
    sentences.append(text[start:].strip())
    return [sentence for sentence in sentences if sentence]


def split_cue(start: float, end: float, text: str) -> list[tuple[float, float, str]]:
    """Split a cue into one cue per sentence, dividing [start, end] in proportion to character length.

    Pieces shorter than MIN_SPLIT_DURATION are merged into their neighbour. The pieces are contiguous.
    Returns (start, end, text) tuples; a single-sentence cue is returned unchanged.
    """
    sentences = split_sentences(text=text)
    if len(sentences) < 2:
        return [(start, end, text)]
    total = sum(len(sentence) for sentence in sentences)
    pieces: list[list] = []
    cursor = start
    for sentence in sentences:
        piece_end = cursor + (end - start) * len(sentence) / total
        pieces.append([cursor, piece_end, sentence])
        cursor = piece_end
    pieces[-1][1] = end
    merged: list[list] = []
    for piece in pieces:
        if merged and (piece[1] - piece[0] < MIN_SPLIT_DURATION or merged[-1][1] - merged[-1][0] < MIN_SPLIT_DURATION):
            merged[-1][1] = piece[1]
            merged[-1][2] += f" {piece[2]}"
        else:
            merged.append(piece)
    return [(piece_start, piece_end, piece_text) for piece_start, piece_end, piece_text in merged]


def wrap_text(text: str) -> str:
    """Wrap text into balanced lines of at most LINE_LENGTH characters (e.g. 60 chars -> 2 lines of 30)."""
    if len(text) <= LINE_LENGTH:
        return text
    lines_needed = -(-len(text) // LINE_LENGTH)  # ceiling division
    # widen the target width until the words fit in the needed number of lines
    for width in range(-(-len(text) // lines_needed), LINE_LENGTH + 1):
        lines = textwrap.wrap(text, width=width)
        if len(lines) <= lines_needed:
            break
    return "\n".join(lines)


def split_long(start: float, end: float, text: str) -> list[tuple[float, float, str]]:
    """Split a cue longer than MAX_LINES lines into several cues, dividing [start, end] by character length.

    Splits near the middle, preferring a position right after a comma, semicolon, colon or sentence end
    (within the middle half of the text), else any space. Repeats until every piece fits.
    """
    if len(text) <= LINE_LENGTH * MAX_LINES:
        return [(start, end, text)]
    middle = len(text) / 2
    spaces = [match.start() for match in re.finditer(" ", text)]
    if not spaces:
        return [(start, end, text)]
    # a space preceded by , ; : . ! ? is a clause break
    clause_breaks = [position for position in spaces if text[position - 1] in ",;:.!?" and abs(position - middle) <= len(text) / 4]
    position = min(clause_breaks or spaces, key=lambda candidate: abs(candidate - middle))
    left, right = text[:position].strip(), text[position:].strip()
    boundary = start + (end - start) * len(left) / (len(left) + len(right))
    return split_long(start=start, end=boundary, text=left) + split_long(start=boundary, end=end, text=right)


def format_cues(cues: list[tuple[float, float, str]]) -> list[tuple[float, float, str]]:
    """Split single-speaker cues that are too long and wrap their text into lines; dialogue cues pass through."""
    formatted = []
    for start, end, text in cues:
        if text.startswith("- "):
            formatted.append((start, end, text))
        else:
            formatted.extend(
                (piece_start, piece_end, wrap_text(text=piece_text)) for piece_start, piece_end, piece_text in split_long(start=start, end=end, text=text)
            )
    return formatted


def format_timestamp(seconds: float) -> str:
    """Format seconds as an SRT timestamp, HH:MM:SS,mmm."""
    total_milliseconds = round(seconds * 1000)
    hours, remainder = divmod(total_milliseconds, 3_600_000)
    minutes, remainder = divmod(remainder, 60_000)
    secs, milliseconds = divmod(remainder, 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"


def build_srt(cues: list[tuple[float, float, str]]) -> str:
    """Render (start, end, text) cues as SRT, keeping them non-overlapping with a minimum length of 0.3 s."""
    blocks = []
    previous_end = 0.0
    for index, (start, end, text) in enumerate(cues, start=1):
        start = max(start, previous_end)
        end = max(end, start + 0.3)
        previous_end = end
        blocks.append(f"{index}\n{format_timestamp(start)} --> {format_timestamp(end)}\n{text}\n")
    return "\n".join(blocks)


def main() -> None:
    """Extract audio, transcribe it chunk by chunk, rescale timestamps by the speed-up and write the SRT."""
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
    args = parse_args()
    if not args.api_key:
        raise SystemExit("Provide --api-key or set OPENROUTER_API_KEY")
    if not args.video.is_file():
        raise SystemExit(f"Input not found: {args.video}")
    if not 1.0 <= args.speed_up <= 2.0:
        raise SystemExit("--speed-up must be between 1.0 and 2.0")

    output = args.output or args.video.with_suffix(".srt")
    cues: list[Segment] = []

    with (
        tempfile.TemporaryDirectory() as temp_directory,
        httpx.Client(
            base_url="https://openrouter.ai/api/v1",
            headers={"Authorization": f"Bearer {args.api_key}"},
            timeout=300,
        ) as client,
    ):
        chunks = run_ffmpeg_extract(
            video=args.video,
            directory=Path(temp_directory),
            chunk_seconds=int(args.chunk_minutes * 60),
            speed_up=args.speed_up,
        )
        logger.info("Extracted %d audio chunk(s)", len(chunks))
        offset = 0.0
        for index, chunk in enumerate(chunks, start=1):
            duration = get_duration(path=chunk)
            if duration < MIN_CHUNK_DURATION:
                logger.info("Skipping empty chunk %d/%d", index, len(chunks))
                continue
            logger.info("Transcribing chunk %d/%d with %s", index, len(chunks), args.model)
            segments, has_timestamps = transcribe_chunk(
                client=client,
                path=chunk,
                duration=duration,
                model=args.model,
                language=args.language,
                diarize=args.diarize,
            )
            for segment in segments:
                start = segment.start * args.speed_up + offset
                end = segment.end * args.speed_up + offset
                if args.speed_up > 1.0 or args.diarize or not has_timestamps:
                    pieces = split_cue(start=start, end=end, text=segment.text)
                else:
                    pieces = [(start, end, segment.text)]
                cues.extend(Segment(start=piece[0], end=piece[1], text=piece[2], speaker=segment.speaker) for piece in pieces)
            offset += duration * args.speed_up

    if not cues:
        raise SystemExit("Transcription returned no text")
    srt_cues = format_cues(cues=group_dialogue(cues=cues))
    output.write_text(build_srt(cues=srt_cues), encoding="utf-8")
    logger.info("Wrote %d cues to %s", len(srt_cues), output)


if __name__ == "__main__":
    main()