"""
Subtitle Generation Module using Faster-Whisper on IntelLLM.
Generates word-level synchronized .srt subtitle tracks in English and Spanish.
"""

import os
import time
import logging
from typing import Dict, Any
from .remote import RemotePC

logger = logging.getLogger("YTFactory.Subs")

WHISPER_WORKER_SCRIPT = r'''
import sys
from faster_whisper import WhisperModel

audio_path = sys.argv[1]
lang = sys.argv[2]
out_srt = sys.argv[3]

# Load faster-whisper on CPU using int8 quantization
model = WhisperModel("base", device="cpu", compute_type="int8")
segments, _ = model.transcribe(audio_path, language=lang, word_timestamps=True)

def format_timestamp(seconds):
    ms = int(round(seconds * 1000))
    hours, ms = divmod(ms, 3600000)
    minutes, ms = divmod(ms, 60000)
    secs, ms = divmod(ms, 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"

words = []
for seg in segments:
    words.extend(seg.words or [])

cues = []
current_words = []
start_time = None
last_end_time = None

for w in words:
    if current_words and (len(current_words) >= 6 or (w.start - start_time) > 3.2):
        cues.append((start_time, last_end_time, " ".join(x.word.strip() for x in current_words)))
        current_words = []
        
    if not current_words:
        start_time = w.start
        
    current_words.append(w)
    last_end_time = w.end

if current_words:
    cues.append((start_time, last_end_time, " ".join(x.word.strip() for x in current_words)))

with open(out_srt, "w", encoding="utf-8") as f:
    for i, (s, e, text) in enumerate(cues, 1):
        f.write(f"{i}\n{format_timestamp(s)} --> {format_timestamp(e)}\n{text}\n\n")

print(f"GENERATED_{len(cues)}_CUES")
'''

class SubGen:
    def __init__(self, pcs: Dict[str, RemotePC]):
        self.pcs = pcs

    def generate(self, audio_path: str, job_dir: str, lang: str = "en") -> str:
        """Generate accurate .srt subtitle file on IntelLLM."""
        logger.info(f"[SUBS:{lang.upper()}] Transcribing audio with Faster-Whisper on IntelLLM...")
        pc = self.pcs.get("intelllm")
        if not pc:
            raise RuntimeError("IntelLLM is not available for Whisper transcription")

        ts = int(time.time())
        remote_audio = f"/tmp/whisper_in_{lang}_{ts}.wav"
        remote_srt = f"/tmp/whisper_out_{lang}_{ts}.srt"
        remote_script = f"/tmp/run_whisper_{lang}_{ts}.py"

        pc.scp_to(audio_path, remote_audio)
        
        # Upload whisper worker script
        rc, _, err = pc.run(f"cat > {remote_script} << 'PYEOF'\n{WHISPER_WORKER_SCRIPT.strip()}\nPYEOF")
        if rc != 0:
            raise RuntimeError(f"Failed to deploy whisper worker script: {err}")

        # Execute in intel venv
        cmd = (
            f"source ~/venvs/youtube-factory-intel/bin/activate 2>/dev/null || true; "
            f"python3 {remote_script} {remote_audio} {lang} {remote_srt}"
        )
        rc, stdout, stderr = pc.run(cmd, timeout=900)
        if rc != 0:
            logger.error(f"[SUBS:{lang.upper()}] Whisper transcription failed: {stderr}")
            raise RuntimeError(f"Whisper failed: {stderr}")

        local_srt = os.path.join(job_dir, f"subtitles_{lang}.srt")
        pc.scp_from(remote_srt, local_srt)

        # Cleanup remote tmp files
        pc.run(f"rm -f {remote_audio} {remote_srt} {remote_script}")
        logger.info(f"[SUBS:{lang.upper()}] Subtitles generated: {local_srt} ({stdout.strip()})")
        return local_srt