#!/usr/bin/env python3
"""
MLX Whisper enhanced transcription with precise subtitle timing.
Uses mlx-whisper (Apple Silicon optimized) + word-level timestamps + smart split.

Usage:
  python3 whisper_transcribe.py <audio_file> --initial-prompt "<提示词>" [选项]

Options:
  --model MODEL         mlx-whisper model: turbo/small/medium (default: turbo)
  --language LANG       Language code (ja/en/zh/...)
  --output-dir DIR      Output directory (default: same as input)
  --output-format FMT   Output format: txt/srt/json/all (default: txt)
  --no-clean            Disable hallucination cleanup
  --no-word-ts          Disable word-level timestamps (revert to segment-level)
  --no-smart-split      Disable smart split on pauses (default: on)
  --no-vad-filter       Disable silence/VAD filtering (default: on)
  --max-line-duration   Max duration per subtitle line in seconds (default: 5.0)
  --initial-prompt TEXT  Initial prompt for context guidance (required)
  --verbose             Show progress
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Optional

import tempfile as tf_module
import shutil

# mlx-whisper model repo mapping
MLX_MODELS = {
    "turbo": "mlx-community/whisper-large-v3-turbo",
    "large-v3-turbo": "mlx-community/whisper-large-v3-turbo",
    "small": "mlx-community/whisper-small-mlx",
    "medium": "mlx-community/whisper-medium-mlx",
}


# ── Audio Preprocessing ──────────────────────────────────────────────────────

def preprocess_audio(audio_path: str, output_dir: str) -> str:
    """Preprocess audio: downsample to 16kHz mono + loudness normalization."""
    out_path = os.path.join(output_dir, "_preprocessed.wav")
    cmd = [
        "ffmpeg", "-y", "-i", str(audio_path),
        "-af", "loudnorm=I=-16:TP=-1.5:LRA=11",
        "-ar", "16000", "-ac", "1",
        "-f", "wav", out_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg preprocessing failed: {result.stderr}")
    return out_path


def get_audio_duration(audio_path: str) -> float:
    """Get audio duration in seconds using ffprobe."""
    cmd = [
        "ffprobe", "-v", "quiet", "-print_format", "json",
        "-show_format", str(audio_path)
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffprobe failed: {result.stderr}")
    info = json.loads(result.stdout)
    return float(info["format"]["duration"])


def detect_silence(audio_path: str, noise_db: int = -40, min_duration: float = 0.3,
                   audio_duration: float = None) -> list[tuple[float, float]]:
    """Detect silent regions using ffmpeg silencedetect.
    Returns list of (start, end) tuples for silent segments.
    If audio ends in silence, ffmpeg emits silence_start without silence_end;
    we use audio_duration to close the last unclosed silence region."""
    cmd = [
        "ffmpeg", "-i", str(audio_path),
        "-af", f"silencedetect=noise={noise_db}dB:d={min_duration}",
        "-f", "null", "-"
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    starts = re.findall(r"silence_start: ([\d.]+)", result.stderr)
    ends = re.findall(r"silence_end: ([\d.]+)", result.stderr)

    starts = [float(s) for s in starts]
    ends = [float(e) for e in ends]

    silences = []
    # Pair starts with ends; if there's an unmatched start at the end,
    # close it with audio_duration
    for i, s in enumerate(starts):
        if i < len(ends):
            silences.append((s, ends[i]))
        elif audio_duration:
            # Unclosed silence at end of audio
            silences.append((s, audio_duration))

    return silences


# ── Transcription ────────────────────────────────────────────────────────────

def transcribe_file(audio_path: str, model_repo: str, language: str,
                    initial_prompt: Optional[str] = None,
                    word_timestamps: bool = True,
                    verbose: bool = False) -> dict:
    """Transcribe a single audio file with mlx-whisper."""
    from mlx_whisper import transcribe

    options = {
        "path_or_hf_repo": model_repo,
        "language": language,
        "verbose": verbose,
        "word_timestamps": word_timestamps,
    }
    if initial_prompt:
        options["initial_prompt"] = initial_prompt

    return transcribe(audio_path, **options)


# ── Post-processing: Smart Split & Glue Fix ──────────────────────────────────

def smart_split_segments(segments: list[dict], max_duration: float = 5.0,
                         min_pause: float = 0.3) -> list[dict]:
    """Split long segments on natural pauses using word-level timestamps.
    Fixes subtitle 'glue' where one line sticks on screen too long."""
    result = []

    for seg in segments:
        words = seg.get("words", [])
        if not words:
            result.append(seg)
            continue

        # If segment is short enough, keep as-is
        if seg["end"] - seg["start"] <= max_duration:
            result.append(seg)
            continue

        # Split on pauses between words
        current_words = []
        current_start = words[0]["start"]

        for i, w in enumerate(words):
            current_words.append(w)

            # Check for pause after this word
            pause = 0.0
            if i < len(words) - 1:
                pause = words[i + 1]["start"] - w["end"]

            # Split condition: pause long enough OR current chunk exceeds max_duration
            chunk_duration = w["end"] - current_start
            should_split = (pause >= min_pause and chunk_duration >= 1.5) or \
                           chunk_duration >= max_duration

            if should_split and current_words:
                text = "".join(wd["word"] for wd in current_words).strip()
                if text:
                    result.append({
                        "start": current_start,
                        "end": w["end"],
                        "text": text,
                        "words": current_words[:],
                    })
                current_words = []
                if i < len(words) - 1:
                    current_start = words[i + 1]["start"]

        # Flush remaining words
        if current_words:
            text = "".join(wd["word"] for wd in current_words).strip()
            if text:
                result.append({
                    "start": current_start,
                    "end": current_words[-1]["end"],
                    "text": text,
                    "words": current_words[:],
                })

    return result


def remove_vad_gaps(segments: list[dict], silences: list[tuple[float, float]],
                    min_gap: float = 0.5) -> list[dict]:
    """Remove segments that fall entirely within detected silent regions.
    Prevents 'ghost subtitles' during music-only sections."""
    if not silences:
        return segments

    def in_silence(start: float, end: float) -> bool:
        for s_start, s_end in silences:
            # If segment is mostly inside a silence region
            overlap_start = max(start, s_start)
            overlap_end = min(end, s_end)
            if overlap_end > overlap_start:
                overlap = overlap_end - overlap_start
                seg_len = end - start
                if seg_len <= 0:
                    continue  # Skip zero-length segments
                if overlap / seg_len > 0.7:  # 70% or more in silence
                    return True
        return False

    filtered = []
    for seg in segments:
        if not in_silence(seg["start"], seg["end"]):
            filtered.append(seg)
    return filtered


def clean_hallucinations(segments: list[dict]) -> list[dict]:
    """Remove trailing hallucinated repetitions."""
    if not segments:
        return segments

    texts = [s["text"] for s in segments]
    cut_point = len(texts)
    seen_repeats = 0
    max_repeats = 5

    for i in range(len(texts) - 1, 0, -1):
        if texts[i] == texts[i - 1] and texts[i].strip():
            seen_repeats += 1
            if seen_repeats >= max_repeats:
                cut_point = i
                break
        else:
            seen_repeats = 0

    filler_pattern = re.compile(r'^(あ+|啊+|um+|uh+|er+|ah+)$', re.IGNORECASE)
    filler_count = 0
    for i in range(len(texts) - 1, max(cut_point - 10, 0), -1):
        if filler_pattern.match(texts[i].strip()):
            filler_count += 1
        else:
            break

    if filler_count > 3:
        cut_point = len(texts) - filler_count

    return segments[:cut_point]


def is_phrase_repetition(text: str) -> bool:
    """检测短语/词组级别的幻听死循环。
    只对大段重复敏感：精确周期需 ≥10 次重复，n-gram 需 ≥30 字符。
    短文本如「うんうんうん」「うふふふ」正常放行。"""
    t = text.strip()
    if len(t) < 6:
        return False
    # 精确周期性检测: 如果 t 由某个子串重复多次构成
    pos = (t + t).find(t, 1)
    if pos != -1 and pos < len(t):
        # 至少重复 10 次以上才算循环（避免 "うん"×3 这种正常相槌）
        if len(t) // pos >= 10:
            return True
    # 宽松检测：前后两半完全相同或包含关系（仅对较长文本）
    if len(t) >= 16:
        mid = len(t) // 2
        if t[:mid] == t[mid:] or (mid >= 3 and t[:mid] in t[mid:]):
            return True
    # 近似周期检测：对付长度不对齐的循环（如 "ABC"*55 + "AB"）
    # n-gram 去重率极低 → 文本由少数短语反复拼接
    if len(t) >= 30:
        for n in (3, 4, 5):
            grams = [t[i:i + n] for i in range(len(t) - n + 1)]
            unique_ratio = len(set(grams)) / len(grams)
            if unique_ratio < 0.15:
                return True
    return False


def fix_repetition_loops(segments: list[dict], audio_path: str,
                         model_repo: str, language: str,
                         audio_duration: float,
                         initial_prompt: Optional[str] = None,
                         char_threshold: float = 0.7,
                         silences: Optional[list[tuple[float, float]]] = None) -> list[dict]:
    """检测并重新转写 Whisper 陷入死循环或长空段的异常片段。

    Detection modes:
    1. Single-char loop — one character dominates (e.g. '秀秀秀秀秀...')
    2. Phrase loop — text built from a repeating substring
    3. Empty segment — text empty AND segment ≥15s AND not in VAD silence

    If retry produces no valid text, the segment range is deleted.
    Traverses ranges in REVERSE to avoid index shifting from list mutations.
    """
    # Helper: check if a time range overlaps any VAD silence region
    def _in_silence(seg_start: float, seg_end: float) -> bool:
        if not silences:
            return False
        for s_start, s_end in silences:
            if seg_start < s_end and seg_end > s_start:
                return True
        return False

    # 1. Detect problem segments
    bad_indices = []
    for i, seg in enumerate(segments):
        text = seg.get("text", "").strip()

        # Mode A1: single-char loop (only for very long repetitions)
        if len(text) >= 20:
            most_common = max(set(text), key=text.count)
            if text.count(most_common) / len(text) > char_threshold:
                bad_indices.append(i)
                continue

        # Mode A2: phrase-level loop
        if is_phrase_repetition(text):
            bad_indices.append(i)
            continue

        # Mode B: long empty segment not explainable by VAD silence
        seg_len = seg["end"] - seg["start"]
        if not text and seg_len >= 15.0 and not _in_silence(seg["start"], seg["end"]):
            bad_indices.append(i)

    if not bad_indices:
        return segments

    # 2. Group consecutive indices into continuous ranges
    ranges = []
    start_i = bad_indices[0]
    end_i = bad_indices[0]
    for i in bad_indices[1:]:
        if i == end_i + 1:
            end_i = i
        else:
            ranges.append((start_i, end_i))
            start_i = end_i = i
    ranges.append((start_i, end_i))

    # 3. Re-transcribe each range with a cold start
    # CRITICAL: reversed() prevents index shifting when list length changes
    from mlx_whisper import transcribe

    result = list(segments)
    total_replaced = 0

    for seg_start_idx, seg_end_idx in reversed(ranges):
        t_start = max(0, result[seg_start_idx]["start"] - 0.3)
        t_end = min(result[seg_end_idx]["end"] + 0.3, audio_duration)

        print(f"  Problem segment at {t_start:.1f}s-{t_end:.1f}s, "
              f"re-transcribing ({seg_end_idx - seg_start_idx + 1} segs)...")

        tmp_dir = tf_module.mkdtemp()
        try:
            seg_path = os.path.join(tmp_dir, "_retry.wav")
            subprocess.run([
                "ffmpeg", "-y", "-i", audio_path,
                "-ss", str(t_start),
                "-to", str(t_end),
                "-af", "loudnorm=I=-16:TP=-1.5:LRA=11",
                "-ar", "16000", "-ac", "1",
                "-f", "wav", seg_path,
            ], capture_output=True, text=True, check=True)

            # Cold-start transcribe (no preceding context)
            retry_opts = {
                "path_or_hf_repo": model_repo,
                "language": language,
                "verbose": False,
                "word_timestamps": True,
            }
            if initial_prompt:
                retry_opts["initial_prompt"] = initial_prompt

            retry_result = transcribe(seg_path, **retry_opts)

            # Build replacement segments with absolute timestamps
            new_segs = []
            for s in retry_result.get("segments", []):
                text_clean = s["text"].strip()
                # Filter: no empty text, no repeated loops in retry
                if not text_clean:
                    continue
                if len(text_clean) >= 5:
                    mc = max(set(text_clean), key=text_clean.count)
                    if text_clean.count(mc) / len(text_clean) > char_threshold:
                        continue
                if is_phrase_repetition(text_clean):
                    continue
                # Density check: impossible char rate → hallucination
                # Normal Japanese speech is ~7 chars/s; 20 is very generous
                seg_dur = s["end"] - s["start"]
                if seg_dur > 0 and len(text_clean) > seg_dur * 20:
                    continue
                new_segs.append({
                    "start": s["start"] + t_start,
                    "end": s["end"] + t_start,
                    "text": text_clean,
                    "words": [
                        {"start": w["start"] + t_start,
                         "end": w["end"] + t_start,
                         "word": w["word"]}
                        for w in s.get("words", [])
                    ]
                })

            if new_segs:
                result[seg_start_idx:seg_end_idx + 1] = new_segs
                total_replaced += (seg_end_idx - seg_start_idx + 1)
                print(f"    Replaced with {len(new_segs)} corrected segments")
            else:
                # Retry produced nothing useful — delete (silence/noise/hallucination)
                del result[seg_start_idx:seg_end_idx + 1]
                print(f"    Retry empty, removed {seg_end_idx - seg_start_idx + 1} segment(s)")

        except Exception as e:
            print(f"    Warning: retry failed ({e}), keeping original segments")
        finally:
            shutil.rmtree(tmp_dir, ignore_errors=True)

    if total_replaced:
        print(f"  Repetition fix: replaced {total_replaced} hallucinated segments")
    return result


# ── Output Formatting ────────────────────────────────────────────────────────

def format_timestamp_srt(seconds: float) -> str:
    """Format seconds to SRT timestamp: HH:MM:SS,mmm
    Uses round() to avoid floating-point precision errors where
    1.001 % 1 = 0.000999... → int() truncates to 000 instead of 001.
    Clamps negative values to 0 (Whisper shouldn't return them, but be safe)."""
    if seconds < 0:
        seconds = 0.0
    total_ms = int(round(seconds * 1000))
    h = total_ms // 3600000
    m = (total_ms % 3600000) // 60000
    s = (total_ms % 60000) // 1000
    ms = total_ms % 1000
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def write_output(segments: list[dict], output_path: str, fmt: str):
    """Write transcription output in specified format."""
    base = Path(output_path).with_suffix("")

    if fmt in ("txt", "all"):
        txt_path = str(base) + ".txt"
        with open(txt_path, "w", encoding="utf-8") as f:
            for seg in segments:
                f.write(seg["text"] + "\n")
        print(f"  Written: {txt_path}")

    if fmt in ("srt", "all"):
        srt_path = str(base) + ".srt"
        with open(srt_path, "w", encoding="utf-8") as f:
            for i, seg in enumerate(segments, 1):
                f.write(f"{i}\n")
                f.write(f"{format_timestamp_srt(seg['start'])} --> "
                        f"{format_timestamp_srt(seg['end'])}\n")
                f.write(f"{seg['text']}\n\n")
        print(f"  Written: {srt_path}")

    if fmt in ("json", "all"):
        json_path = str(base) + ".json"
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(segments, f, ensure_ascii=False, indent=2)
        print(f"  Written: {json_path}")


# ── Main ─────────────────────────────────────────────────────────────────────

@contextmanager
def tempfile_directory():
    d = tf_module.mkdtemp()
    try:
        yield d
    finally:
        shutil.rmtree(d, ignore_errors=True)


def main():
    parser = argparse.ArgumentParser(
        description="MLX Whisper enhanced transcription with precise subtitles")
    parser.add_argument("audio", help="Input audio file path")
    parser.add_argument("--model", default="turbo",
                        choices=["turbo", "large-v3-turbo", "small", "medium"],
                        help="mlx-whisper model (default: turbo)")
    parser.add_argument("--language", default="ja",
                        help="Language code (default: ja)")
    parser.add_argument("--output-dir", default=None,
                        help="Output directory (default: same as input)")
    parser.add_argument("--output-format", default="txt",
                        choices=["txt", "srt", "json", "all"],
                        help="Output format (default: txt)")
    parser.add_argument("--no-clean", action="store_true",
                        help="Disable hallucination cleanup")
    parser.add_argument("--no-word-ts", action="store_true",
                        help="Disable word-level timestamps")
    parser.add_argument("--no-smart-split", action="store_true",
                        help="Disable smart split on pauses")
    parser.add_argument("--no-vad-filter", action="store_true",
                        help="Disable silence/VAD filtering")
    parser.add_argument("--max-line-duration", type=float, default=5.0,
                        help="Max seconds per subtitle line (default: 5.0)")
    parser.add_argument("--initial-prompt", required=True,
                        help="Initial prompt for context guidance (required)")
    parser.add_argument("--verbose", action="store_true",
                        help="Show progress")

    args = parser.parse_args()

    audio_path = Path(args.audio).resolve()
    if not audio_path.exists():
        print(f"Error: File not found: {audio_path}", file=sys.stderr)
        sys.exit(1)

    output_dir = Path(args.output_dir) if args.output_dir else audio_path.parent
    output_dir.mkdir(parents=True, exist_ok=True)

    output_name = audio_path.stem
    model_repo = MLX_MODELS[args.model]

    # Get duration
    duration = get_audio_duration(str(audio_path))
    duration_min = duration / 60
    print(f"Audio: {audio_path.name} ({duration_min:.1f} min)")
    print(f"Model: {args.model} ({model_repo})")

    # ── Preprocessing ──────────────────────────────────────────────────────
    print("Preprocessing: 16kHz mono + loudness normalization...")
    with tempfile_directory() as preprocess_dir:
        preprocessed = preprocess_audio(str(audio_path), preprocess_dir)
        print(f"  Preprocessed: {Path(preprocessed).name}")

        # Optional: detect silence regions for VAD filtering
        silences = []
        if not args.no_vad_filter:
            print("Detecting silent regions...")
            silences = detect_silence(preprocessed, noise_db=-40, min_duration=0.5,
                                      audio_duration=duration)
            print(f"  Found {len(silences)} silent regions")

        # ── Transcribe (mlx-whisper handles long audio internally) ──────────
        print(f"Transcribing with mlx-whisper ({args.model})...")
        t0 = time.time()
        result = transcribe_file(
            preprocessed, model_repo, args.language,
            args.initial_prompt,
            word_timestamps=not args.no_word_ts,
            verbose=args.verbose)
        elapsed = time.time() - t0

        segments_data = result.get("segments", [])
        merged = [
            {"start": s["start"], "end": s["end"],
             "text": s["text"].strip(),
             "words": [
                 {"start": w["start"], "end": w["end"], "word": w["word"]}
                 for w in s.get("words", [])
             ]}
            for s in segments_data
        ]
        print(f"  Transcribed {len(merged)} segments in {elapsed:.1f}s")

    # ── Post-processing ────────────────────────────────────────────────────

    # 1. VAD filter: remove ghost subtitles in silent regions
    if not args.no_vad_filter and silences:
        before = len(merged)
        merged = remove_vad_gaps(merged, silences)
        removed = before - len(merged)
        if removed > 0:
            print(f"  VAD filter: removed {removed} segments in silent regions")

    # 2. Smart split: break long glued segments on natural pauses
    if not args.no_smart_split and not args.no_word_ts:
        before = len(merged)
        merged = smart_split_segments(merged, max_duration=args.max_line_duration)
        if len(merged) > before:
            print(f"  Smart split: {before} → {len(merged)} segments")

    # 3. Hallucination cleanup
    if not args.no_clean:
        before = len(merged)
        merged = clean_hallucinations(merged)
        removed = before - len(merged)
        if removed > 0:
            print(f"  Cleaned: removed {removed} hallucinated segments")

    # 4. Repetition loop detection & local re-transcription
    if not args.no_clean:
        merged = fix_repetition_loops(
            merged, str(audio_path), model_repo,
            args.language, duration, args.initial_prompt,
            silences=silences)

    # ── Output ─────────────────────────────────────────────────────────────
    out_path = str(output_dir / output_name)
    print("Output:")
    write_output(merged, out_path, args.output_format)

    print("Done!")


if __name__ == "__main__":
    main()