#!/usr/bin/env python3
"""Write an SRT with faster-whisper on CPU, int8.

Models: small.en (default) or base.en.
Docs: https://github.com/SYSTRAN/faster-whisper

Word timestamps are on. Decoding is one fixed model on one CPU thread
(cpu_threads=1, num_workers=1, OMP_NUM_THREADS=1, MKL_NUM_THREADS=1),
language en, temperature=0 (a single value, not a fallback list),
beam_size=5, best_of=1, condition_on_previous_text=False, and
vad_filter=False. Every transcript
word is kept, in order, in exactly one cue. A cue closes at 6 words,
2.5 seconds, or about 32 characters, whichever comes first. A cue stays
at least 4 words when that still fits. Cue text is one line.

ai-editor-edit.sh waits for memory, then runs this file under
flock /tmp/epic-heavy-job.lock. This process waits again until
MemAvailable is at least 3GB before it loads the model, so the
lock holder starts the heavy job only after that check.

The word list is also written to words.json beside the SRT.
"""

from __future__ import annotations

import os

# Set before faster-whisper (and its OpenMP / MKL libs) load.
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

import argparse
import json
import sys
import time
from pathlib import Path
from typing import NamedTuple

try:
    from faster_whisper import WhisperModel
except ImportError:
    WhisperModel = None

# 3 GiB of MemAvailable. /proc/meminfo reports KiB. 3 GiB is past 3GB.
NEED_KIB = 3 * 1024 * 1024
MODELS = ("base.en", "small.en")
MAX_WORDS = 6
MIN_WORDS = 4
MAX_SECONDS = 2.5
MAX_CHARS = 32


class Word(NamedTuple):
    text: str
    start: float
    end: float


class Cue(NamedTuple):
    start: float
    end: float
    text: str
    words: tuple[Word, ...]


def format_timestamp(seconds: float) -> str:
    if seconds < 0:
        seconds = 0.0
    total_ms = int(round(seconds * 1000.0))
    hours, rem = divmod(total_ms, 3_600_000)
    minutes, rem = divmod(rem, 60_000)
    secs, millis = divmod(rem, 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"


def normalize_word(text: str) -> str:
    return " ".join(str(text).replace("\r", " ").replace("\n", " ").split())


def cue_text(words: list[Word]) -> str:
    parts: list[str] = []
    for word in words:
        cleaned = " ".join(word.text.replace("\r", " ").replace("\n", " ").split())
        if cleaned:
            parts.append(cleaned)
    return " ".join(parts)


def fits(words: list[Word]) -> bool:
    if not words:
        return False
    if len(words) > MAX_WORDS:
        return False
    if words[-1].end - words[0].start > MAX_SECONDS:
        return False
    if len(cue_text(words)) > MAX_CHARS:
        return False
    return True


def split_cues(words: list[Word]) -> list[Cue]:
    """Group words into one-line cues. Words stay in order and none are dropped."""
    groups: list[list[Word]] = []
    current: list[Word] = []
    for word in words:
        if not current:
            current = [word]
            continue
        trial = current + [word]
        if fits(trial):
            current = trial
            continue
        groups.append(current)
        current = [word]
    if current:
        groups.append(current)
    groups = _borrow_for_minimum(groups)
    return [_make_cue(group) for group in groups]


def _borrow_for_minimum(groups: list[list[Word]]) -> list[list[Word]]:
    """Pull words into a short final cue when the previous cue can spare them."""
    if len(groups) < 2:
        return groups
    while len(groups[-1]) < MIN_WORDS and len(groups[-2]) > MIN_WORDS:
        moved = groups[-2][-1]
        previous = groups[-2][:-1]
        last = [moved, *groups[-1]]
        if len(previous) < MIN_WORDS or not fits(previous) or not fits(last):
            break
        groups[-2] = previous
        groups[-1] = last
    return groups


def _make_cue(group: list[Word]) -> Cue:
    start = float(group[0].start)
    end = float(group[-1].end)
    if end <= start:
        end = start + 0.20
    return Cue(start, end, cue_text(group), tuple(group))


def words_from_segments(segments) -> list[Word]:
    collected: list[Word] = []
    for segment in segments:
        raw_words = getattr(segment, "words", None) or []
        if raw_words:
            for raw in raw_words:
                text = normalize_word(getattr(raw, "word", ""))
                if not text:
                    continue
                start = float(raw.start)
                end = float(raw.end)
                if end < start:
                    end = start
                collected.append(Word(text, start, end))
            continue
        text = normalize_word(getattr(segment, "text", ""))
        if not text:
            continue
        parts = text.split(" ")
        start = float(segment.start)
        end = float(segment.end)
        if end < start:
            end = start
        span = (end - start) / len(parts)
        for index, part in enumerate(parts):
            collected.append(Word(part, start + index * span, start + (index + 1) * span))
    return collected


def render_srt(cues: list[Cue]) -> str:
    blocks: list[str] = []
    for index, cue in enumerate(cues, start=1):
        blocks.append(
            f"{index}\n{format_timestamp(cue.start)} --> {format_timestamp(cue.end)}\n{cue.text}\n"
        )
    body = "\n".join(blocks)
    if body:
        body += "\n"
    return body


def words_json_path(srt_path: Path) -> Path:
    return srt_path.with_name("words.json")


def dump_words(words: list[Word], path: Path) -> None:
    payload = {
        "words": [
            {"text": word.text, "start": word.start, "end": word.end} for word in words
        ]
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


def load_words(path: Path) -> list[Word]:
    data = json.loads(path.read_text(encoding="utf-8"))
    return [
        Word(str(item["text"]), float(item["start"]), float(item["end"]))
        for item in data["words"]
    ]


def sentence_ender(text: str) -> bool:
    token = text.strip().rstrip("\"'”’")
    return token.endswith((".", "?", "!"))


def last_sentence_end(words: list[Word]) -> float | None:
    found: float | None = None
    for word in words:
        if sentence_ender(word.text):
            found = float(word.end)
    return found


def trim_words_for_export(words: list[Word]) -> tuple[list[Word], float | None]:
    """Drop words that start at or after the last sentence, and anything after it."""
    export_end = last_sentence_end(words)
    if export_end is None:
        return list(words), None
    last_index = max(index for index, word in enumerate(words) if sentence_ender(word.text))
    kept = [
        word
        for index, word in enumerate(words)
        if index <= last_index and (index == last_index or word.start < export_end)
    ]
    return kept, export_end


def cues_trimmed_to_sentence(words: list[Word]) -> tuple[list[Cue], float | None]:
    """Rebuild cues so the last one ends on the sentence word, at the export end."""
    kept, export_end = trim_words_for_export(words)
    if export_end is None:
        return split_cues(words), None
    cues = split_cues(kept)
    if not cues:
        return cues, export_end
    last = cues[-1]
    cues[-1] = last._replace(end=min(float(last.end), float(export_end)))
    return cues, export_end


def write_outputs(words: list[Word], srt_path: Path) -> int:
    cues = split_cues(words)
    srt_path.parent.mkdir(parents=True, exist_ok=True)
    srt_path.write_text(render_srt(cues), encoding="utf-8")
    dump_words(words, words_json_path(srt_path))
    return len(cues)


def mem_available_kib() -> int:
    text = Path("/proc/meminfo").read_text(encoding="ascii")
    for line in text.splitlines():
        if line.startswith("MemAvailable:"):
            return int(line.split()[1])
    raise RuntimeError("MemAvailable is missing from /proc/meminfo")


def wait_for_free_memory() -> None:
    if not Path("/proc/meminfo").exists():
        sys.stderr.write("no /proc/meminfo (not Linux); skipping the free-memory wait\n")
        return
    while True:
        try:
            avail = mem_available_kib()
        except OSError as err:
            sys.stderr.write(f"cannot read /proc/meminfo: {err}\n")
            raise SystemExit(1) from err
        except RuntimeError as err:
            sys.stderr.write(f"{err}\n")
            raise SystemExit(1) from err
        if avail >= NEED_KIB:
            sys.stderr.write(f"free memory ok: MemAvailable {avail} kB\n")
            return
        sys.stderr.write(
            "waiting for >=3GB free memory "
            f"(MemAvailable {avail} kB, need {NEED_KIB} kB)\n"
        )
        time.sleep(10)


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Transcribe with faster-whisper (CPU, int8) and write an SRT."
    )
    parser.add_argument("source", type=Path, help="Audio or video file.")
    parser.add_argument(
        "output",
        type=Path,
        nargs="?",
        help="SRT path. Default: the input path with a .srt suffix.",
    )
    parser.add_argument("--model", default="small.en", choices=MODELS)
    return parser.parse_args(argv)


def apply_sentence_trim_main(argv: list[str]) -> int:
    if len(argv) != 3:
        sys.stderr.write(
            "usage: whisper-srt.py --apply-sentence-trim words.json captions.srt\n"
        )
        return 2
    words_path = Path(argv[1])
    srt_path = Path(argv[2])
    if not words_path.is_file():
        sys.stderr.write(f"words.json not found: {words_path}\n")
        return 1
    cues, export_end = cues_trimmed_to_sentence(load_words(words_path))
    if export_end is None or not cues:
        return 0
    kept = [word for cue in cues for word in cue.words]
    srt_path.parent.mkdir(parents=True, exist_ok=True)
    srt_path.write_text(render_srt(cues), encoding="utf-8")
    dump_words(kept, words_path)
    sys.stdout.write(f"{export_end:.6f}\n")
    return 0


def sentence_end_main(argv: list[str]) -> int:
    if len(argv) != 2:
        sys.stderr.write("usage: whisper-srt.py --sentence-end words.json\n")
        return 2
    path = Path(argv[1])
    if not path.is_file():
        sys.stderr.write(f"words.json not found: {path}\n")
        return 1
    end = last_sentence_end(load_words(path))
    if end is not None:
        sys.stdout.write(f"{end:.6f}\n")
    return 0


def main(argv: list[str]) -> int:
    if argv and argv[0] == "--sentence-end":
        return sentence_end_main(argv)
    if argv and argv[0] == "--apply-sentence-trim":
        return apply_sentence_trim_main(argv)
    args = parse_args(argv)
    if WhisperModel is None:
        sys.stderr.write(
            "faster-whisper is not installed. "
            "Install with: python3 -m pip install faster-whisper\n"
            "Docs: https://github.com/SYSTRAN/faster-whisper\n"
        )
        return 1
    if not args.source.is_file():
        sys.stderr.write(f"input not found: {args.source}\n")
        return 1
    output = args.output if args.output is not None else args.source.with_suffix(".srt")
    wait_for_free_memory()
    # One checkpoint, one thread. temperature is a single value, not a list,
    # so faster-whisper does not walk its fallback schedule.
    model_name = args.model
    model = WhisperModel(
        model_name,
        device="cpu",
        compute_type="int8",
        cpu_threads=1,
        num_workers=1,
    )
    segments, info = model.transcribe(
        str(args.source),
        language="en",
        beam_size=5,
        best_of=1,
        temperature=0,
        condition_on_previous_text=False,
        vad_filter=False,
        word_timestamps=True,
    )
    words = words_from_segments(segments)
    count = write_outputs(words, output)
    sys.stderr.write(
        f"{model_name} language={info.language} cues={count} "
        f"words={len(words)} -> {output}\n"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))