#!/usr/bin/env -S uv run --script --quiet
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "faster-whisper>=1.0.0",
# ]
# ///
"""
Transcription with subtitles. Hebrew uses ivrit.ai, English uses whisper large-v3.

Usage: 
    ./transcribe.py video.mp4                      # Plain transcript
    ./transcribe.py video.mp4 --srt                # Generate SRT file
    ./transcribe.py video.mp4 --srt --embed        # Burn subtitles into video
    ./transcribe.py video.mp4 --srt --translate en # Translate to English
"""

import sys
import os
import argparse
import subprocess
import tempfile
from pathlib import Path
from dataclasses import dataclass


@dataclass
class Subtitle:
    index: int
    start: float
    end: float
    text: str

    def to_srt(self) -> str:
        return f"{self.index}\n{format_srt_timestamp(self.start)} --> {format_srt_timestamp(self.end)}\n{self.text}\n"


def format_srt_timestamp(seconds: float) -> str:
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"


def chunk_text_naturally(text: str, max_chars: int = 42) -> list[str]:
    text = text.strip()
    if len(text) <= max_chars:
        return [text]
    
    break_points = ['. ', ', ', '? ', '! ', ': ', '; ', ' - ', ' – ', ' — ']
    lines = []
    remaining = text
    
    while remaining:
        if len(remaining) <= max_chars:
            lines.append(remaining)
            break
        
        best_break = -1
        for bp in break_points:
            idx = remaining[:max_chars].rfind(bp)
            if idx > best_break:
                best_break = idx + len(bp) - 1
        
        if best_break <= 0:
            best_break = remaining[:max_chars].rfind(' ')
        if best_break <= 0:
            best_break = max_chars
        
        lines.append(remaining[:best_break].strip())
        remaining = remaining[best_break:].strip()
    
    if len(lines) > 2:
        lines = [lines[0], ' '.join(lines[1:])]
    
    return lines


def merge_into_subtitles(segments: list, min_duration: float = 1.0, max_duration: float = 7.0, max_chars: int = 84) -> list[Subtitle]:
    if not segments:
        return []
    
    subtitles = []
    current_text = ""
    current_start = None
    current_end = None
    
    for seg in segments:
        seg_text = seg.text.strip()
        seg_start = seg.start
        seg_end = seg.end
        
        if current_start is None:
            current_start = seg_start
            current_end = seg_end
            current_text = seg_text
            continue
        
        potential_text = current_text + " " + seg_text
        potential_duration = seg_end - current_start
        gap = seg_start - current_end
        
        should_merge = (
            (potential_duration <= max_duration and len(potential_text) <= max_chars and (current_end - current_start) < min_duration) or
            (gap < 0.3 and potential_duration <= max_duration and len(potential_text) <= max_chars)
        )
        
        if should_merge:
            current_text = potential_text
            current_end = seg_end
        else:
            lines = chunk_text_naturally(current_text)
            subtitles.append(Subtitle(len(subtitles) + 1, current_start, current_end, '\n'.join(lines)))
            current_start = seg_start
            current_end = seg_end
            current_text = seg_text
    
    if current_text:
        lines = chunk_text_naturally(current_text)
        subtitles.append(Subtitle(len(subtitles) + 1, current_start, current_end, '\n'.join(lines)))
    
    return subtitles


def transcribe(file_path: str, language: str | None = None, use_turbo: bool = True, 
               generate_srt: bool = False, translate_to: str | None = None):
    """Transcribe audio/video file. Auto-detects language, uses ivrit.ai for Hebrew."""
    from faster_whisper import WhisperModel
    
    # Determine task: transcribe or translate
    task = "translate" if translate_to == "en" else "transcribe"
    
    # Select model based on language and task
    if task == "translate":
        # Translation requires large-v3 (turbo doesn't translate well)
        model_name = "large-v3"
        # Use int8 quantization on CPU to save memory
        print(f"📦 Loading model: {model_name} (int8 for translation)...", file=sys.stderr)
        model = WhisperModel(model_name, device="cpu", compute_type="int8")
    elif language == "he":
        model_name = "ivrit-ai/whisper-large-v3-turbo-ct2" if use_turbo else "ivrit-ai/whisper-large-v3-ct2"
        print(f"📦 Loading model: {model_name}...", file=sys.stderr)
        model = WhisperModel(model_name, device="auto", compute_type="auto")
    else:
        model_name = "large-v3-turbo" if use_turbo else "large-v3"
        print(f"📦 Loading model: {model_name}...", file=sys.stderr)
        model = WhisperModel(model_name, device="auto", compute_type="auto")
    
    print(f"🎤 Transcribing: {file_path}...", file=sys.stderr)
    if task == "translate":
        print(f"🌐 Translating to English...", file=sys.stderr)
    
    segments, info = model.transcribe(
        file_path, 
        language=language,
        task=task,
        word_timestamps=generate_srt,
        vad_filter=generate_srt,
    )
    
    detected_lang = info.language
    print(f"✓ Detected: {detected_lang} (confidence: {info.language_probability:.0%})", file=sys.stderr)
    
    # If Hebrew detected but we used standard model, re-run with ivrit.ai (unless translating)
    if detected_lang == "he" and language is None and "ivrit-ai" not in model_name and task != "translate":
        print("🔄 Hebrew detected, switching to ivrit.ai model...", file=sys.stderr)
        return transcribe(file_path, language="he", use_turbo=use_turbo, 
                         generate_srt=generate_srt, translate_to=translate_to)
    
    raw_segments = list(segments)
    
    if generate_srt:
        subtitles = merge_into_subtitles(raw_segments)
        print(f"✓ Created {len(subtitles)} subtitles", file=sys.stderr)
        return '\n'.join(sub.to_srt() for sub in subtitles), subtitles[-1].end if subtitles else 0, detected_lang
    
    return " ".join(seg.text.strip() for seg in raw_segments), None, detected_lang


def embed_subtitles(video_path: str, srt_content: str, output_path: str, burn: bool = False):
    """Embed subtitles into video using ffmpeg."""
    # Use ffmpeg-full if available (has libass for burn-in)
    ffmpeg_bin = "/opt/homebrew/opt/ffmpeg-full/bin/ffmpeg"
    if not os.path.exists(ffmpeg_bin):
        ffmpeg_bin = "ffmpeg"
    
    # Write SRT to temp file
    srt_path = "/tmp/subtitles_temp.srt"
    with open(srt_path, 'w', encoding='utf-8') as f:
        f.write(srt_content)
    
    try:
        if burn:
            print(f"🔥 Burning subtitles into video...", file=sys.stderr)
            # Hard-code (burn) subtitles into video using libass
            # Style: movie-style - smaller text at bottom with outline
            escaped_srt = srt_path.replace(":", "\\:")
            filter_str = f"subtitles={escaped_srt}:force_style='FontSize=12,PrimaryColour=&Hffffff,OutlineColour=&H000000,BorderStyle=1,Outline=1,Shadow=0,MarginV=12,Alignment=2'"
            
            cmd = [
                ffmpeg_bin, '-y',
                '-i', video_path,
                '-vf', filter_str,
                '-c:a', 'copy',
                output_path
            ]
        else:
            print(f"🎬 Embedding soft subtitles...", file=sys.stderr)
            # Soft subtitles (selectable in player)
            cmd = [
                ffmpeg_bin, '-y',
                '-i', video_path,
                '-i', srt_path,
                '-c', 'copy',
                '-c:s', 'mov_text',
                '-metadata:s:s:0', 'language=heb',
                output_path
            ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print(f"ffmpeg error: {result.stderr}", file=sys.stderr)
            raise RuntimeError("ffmpeg failed")
        
        print(f"✓ Video saved: {output_path}", file=sys.stderr)
        if not burn:
            print(f"  (Soft subs - enable in player with V key)", file=sys.stderr)
    finally:
        if os.path.exists(srt_path):
            os.unlink(srt_path)


def main():
    parser = argparse.ArgumentParser(
        description="Transcribe audio/video (Hebrew via ivrit.ai, English via whisper)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s video.mp4                        # Transcribe (auto-detect language)
  %(prog)s video.mp4 --lang he              # Force Hebrew
  %(prog)s video.mp4 --srt                  # Generate SRT subtitles
  %(prog)s video.mp4 --srt --embed          # Burn subtitles into video
  %(prog)s video.mp4 --srt --translate en   # Translate to English subtitles
        """
    )
    parser.add_argument("file", help="Audio or video file")
    parser.add_argument("--lang", "-l", choices=["he", "en"],
                       help="Force language (auto-detect if not specified)")
    parser.add_argument("--turbo", action="store_true", default=True, 
                       help="Use turbo model for Hebrew (faster, default)")
    parser.add_argument("--accurate", action="store_true",
                       help="Use accurate model for Hebrew (slower but better)")
    parser.add_argument("--timestamps", "-t", action="store_true",
                       help="Include timestamps in plain text output")
    parser.add_argument("--srt", action="store_true",
                       help="Generate SRT subtitle file")
    parser.add_argument("--embed", action="store_true",
                       help="Embed soft subtitles into video (toggle in player)")
    parser.add_argument("--burn", action="store_true",
                       help="Burn subtitles into video (always visible, for WhatsApp)")
    parser.add_argument("--translate", metavar="LANG", choices=["en"],
                       help="Translate subtitles to language (currently: en)")
    parser.add_argument("--output", "-o", help="Output file path")
    
    args = parser.parse_args()
    
    input_path = Path(args.file)
    if not input_path.exists():
        print(f"❌ File not found: {args.file}", file=sys.stderr)
        sys.exit(1)
    
    if (args.embed or args.burn) and not args.srt:
        print("❌ --embed/--burn requires --srt", file=sys.stderr)
        sys.exit(1)
    
    use_turbo = not args.accurate
    result, duration, detected_lang = transcribe(
        args.file, 
        language=args.lang,
        use_turbo=use_turbo, 
        generate_srt=args.srt,
        translate_to=args.translate
    )
    
    # Determine output path
    if args.output:
        output_path = Path(args.output)
    elif args.embed or args.burn:
        output_path = input_path.with_stem(input_path.stem + "_subtitled").with_suffix('.mp4')
    elif args.srt:
        output_path = input_path.with_suffix('.srt')
    else:
        output_path = None
    
    # Handle embedding
    if args.embed or args.burn:
        embed_subtitles(str(input_path), result, str(output_path), burn=args.burn)
    elif output_path:
        output_path.write_text(result, encoding="utf-8")
        print(f"✓ Saved: {output_path}", file=sys.stderr)
        if duration:
            print(f"  Duration: {format_srt_timestamp(duration)}", file=sys.stderr)
    else:
        print(result)


if __name__ == "__main__":
    main()