#!/usr/bin/env python3
import os
import sys
import shutil
import subprocess
from pathlib import Path
from faster_whisper import WhisperModel
from faster_whisper.utils import format_timestamp

# Prevent OpenMP thread conflicts from crashing the script
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 5 minutes per chunk (300 seconds) keeps RAM usage low and stable
CHUNK_DURATION_SEC = 300 

print("🔧 Initializing C++ Engine...")
model = WhisperModel(
    "base", 
    device="cpu", 
    compute_type="int8", 
    cpu_threads=0, # Let CTranslate2 auto-detect safe core counts
    num_workers=1         
)
print("✅ C++ Model loaded successfully.")


def extract_audio_and_chunk(video_path: str, output_dir: Path) -> list:
    """Extracts and splits audio into 5-minute chunks using a single FFmpeg pass."""
    print("🚀 Extracting and chunking audio with FFmpeg...")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Segment format output: chunk_000.wav, chunk_001.wav, etc.
    chunk_pattern = str(output_dir / "chunk_%03d.wav")
    
    command = [
        "ffmpeg", "-y", "-i", video_path,
        "-vn", "-ac", "1", "-ar", "16000",
        "-acodec", "pcm_s16le", "-sn", "-map_chapters", "-1",
        "-f", "segment", "-segment_time", str(CHUNK_DURATION_SEC),
        chunk_pattern
    ]
    
    result = subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        print(f"❌ FFmpeg Error Output:\n{result.stderr}")
        raise RuntimeError("FFmpeg extraction and chunking failed.")
        
    # Return sorted list of generated chunk files
    return sorted(list(output_dir.glob("chunk_*.wav")))


def transcribe_to_srt(video_path: str, force: bool = False):
    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():
        if not force:
            print(f"❌ Error: Transcription file already exists: {srt_path}")
            return
        else:
            srt_path.unlink()

    try:
        # Step 1: Split audio into bite-sized pieces
        audio_chunks = extract_audio_and_chunk(str(video_path_obj), temp_dir)
        if not audio_chunks:
            print("❌ Error: No audio chunks were generated.")
            return

        print(f"📦 Successfully split audio into {len(audio_chunks)} chunks.")
        print("🎙️ Starting safe chunk-by-chunk transcription...")

        global_segment_index = 1
        
        with open(srt_path, "w", encoding="utf-8") as srt_file:
            for chunk_idx, chunk_path in enumerate(audio_chunks):
                # Calculate the time offset for the current chunk
                time_offset = chunk_idx * CHUNK_DURATION_SEC
                print(f"\n⏳ Processing chunk {chunk_idx + 1}/{len(audio_chunks)} ({chunk_path.name})...")
                
                segments_generator, info = model.transcribe(
                    str(chunk_path), 
                    beam_size=1,
                    vad_filter=True,
                    temperature=0.0
                )
                
                # Consume chunk generator and shift timestamps instantly
                for segment in segments_generator:
                    # Shift timestamps relative to the original video timeline
                    actual_start = segment.start + time_offset
                    actual_end = segment.end + time_offset
                    
                    start_str = format_timestamp(actual_start, always_include_hours=True)
                    end_str = format_timestamp(actual_end, always_include_hours=True)
                    
                    srt_file.write(f"{global_segment_index}\n{start_str} --> {end_str}\n{segment.text.strip()}\n\n")
                    global_segment_index += 1
                    
                # Free up space as we go by deleting the processed chunk
                chunk_path.unlink()

        print(f"\n✅ All chunks combined! SRT subtitle file saved in: {srt_path}")

    except Exception as e:
        print(f"\n❌ Execution Error: {e}")
    finally:
        # Clean up the temporary folder entirely
        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"❌ System Error: Target video file does not exist at path: {target_video}")
    else:
        transcribe_to_srt(target_video, force=True)