#!/usr/bin/env python3
import os
import shutil
import subprocess
from pathlib import Path
import wave
from faster_whisper import WhisperModel

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

CHUNK_DURATION_SEC = 300 

def extract_audio_and_chunk(video_path: str, output_dir: Path) -> list[tuple[Path, float]]:
    """Returns list of (chunk_path, start_offset_seconds)"""
    print("🚀 Extracting and chunking audio with FFmpeg...")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    chunk_pattern = str(output_dir / "chunk_%03d.wav")
    
    command = [
        "ffmpeg", "-y", "-i", video_path,
        "-vn", "-ac", "1", "-ar", "16000",
        "-acodec", "pcm_s16le",
        "-f", "segment", "-segment_time", str(CHUNK_DURATION_SEC),
        "-reset_timestamps", "0", # keep each chunk starting at 0
        chunk_pattern
    ]
    
    result = subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        print(f"❌ FFmpeg Error:\n{result.stderr}")
        raise RuntimeError("FFmpeg failed")

    chunks = sorted(output_dir.glob("chunk_*.wav"))
    
    # FIX: Calculate real offset from actual file durations, not idx*300
    offsets = []
    current_offset = 0.0
    for chunk_path in chunks:
        offsets.append((chunk_path, current_offset))
        with wave.open(str(chunk_path), 'rb') as wf:
            duration = wf.getnframes() / wf.getframerate()
            current_offset += duration
            
    return offsets

def format_srt_time(seconds: float):
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int((seconds % 1) * 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"

def transcribe_to_srt(video_path: str, force: bool = False):
    if not hasattr(transcribe_to_srt, "model"):
        print("🔧 Initializing C++ Engine...")
        transcribe_to_srt.model = WhisperModel(
            "distil-small.en",
            device="cpu",
            compute_type="int8",
            cpu_threads=4, # 4 = physical cores on E3-1280 v5 = fastest
            num_workers=1
        )
        print("✅ Model loaded.")

    video_path_obj = Path(video_path)
    srt_path = video_path_obj.with_suffix(".srt")
    temp_dir = video_path_obj.parent / f"temp_chunks_{video_path_obj.stem}"

    if srt_path.exists() and not force:
        print(f"❌ SRT already exists: {srt_path}")
        return
    if srt_path.exists():
        srt_path.unlink()

    try:
        audio_chunks = extract_audio_and_chunk(str(video_path_obj), temp_dir)
        print(f"📦 {len(audio_chunks)} chunks.")
        print("🎙 Transcribing...")

        global_idx = 1
        with open(srt_path, "w", encoding="utf-8") as srt_file:
            for chunk_idx, (chunk_path, time_offset) in enumerate(audio_chunks):
                print(f"\n⏳ Chunk {chunk_idx+1}/{len(audio_chunks)} offset={time_offset:.1f}s")
                
                segments, info = transcribe_to_srt.model.transcribe(
                    str(chunk_path), 
                    language="en",
                    beam_size=1,
                    vad_filter=True,
                    vad_parameters=dict(min_silence_duration_ms=500),
                    condition_on_previous_text=False # important for chunked mode
                )
                
                for seg in segments:
                    actual_start = seg.start + time_offset
                    actual_end = seg.end + time_offset
                    
                    # Use comma for SRT, not dot
                    srt_file.write(f"{global_idx}\n{format_srt_time(actual_start)} --> {format_srt_time(actual_end)}\n{seg.text.strip()}\n\n")
                    global_idx += 1
                    
                chunk_path.unlink()

        print(f"\n✅ SRT saved: {srt_path}")

    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback; traceback.print_exc()
    finally:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)

if __name__ == "__main__":
    target_video = "download/videos/2813112936/2813112936.mp4"
    if not os.path.exists(target_video):
        print(f"❌ Not found: {target_video}")
    else:
        transcribe_to_srt(target_video, force=True)