import sys
from unittest.mock import MagicMock

# Comprehensive workaround for missing _lzma in some Python builds
try:
    import lzma
except ImportError:
    mock_lzma = MagicMock()
    mock_lzma.FORMAT_XZ = 1
    mock_lzma.FORMAT_ALONE = 2
    mock_lzma.FORMAT_RAW = 3
    mock_lzma.CHECK_NONE = 0
    mock_lzma.CHECK_CRC32 = 1
    mock_lzma.CHECK_CRC64 = 4
    mock_lzma.CHECK_SHA256 = 10
    sys.modules["_lzma"] = MagicMock()
    sys.modules["lzma"] = mock_lzma

import torch
from transformers import MarianMTModel, MarianTokenizer
import re
import argparse
import multiprocessing
import time
import requests

TARGET_LANGS = {
    "es": "Helsinki-NLP/opus-mt-en-es",
    "fr": "Helsinki-NLP/opus-mt-en-fr",
    "ar": "Helsinki-NLP/opus-mt-en-ar"
}

def split_text_for_translation(text, max_chars=250):
    normalized = " ".join(text.split()).strip()
    if not normalized or len(normalized) <= max_chars:
        return [normalized] if normalized else []
    chunks = []
    current = ""
    sentences = [s for s in re.split(r"(?<=[.!?])\s+", normalized) if s]
    for sentence in sentences:
        if len(sentence) > max_chars:
            words = sentence.split()
            word_chunk = ""
            for word in words:
                candidate = f"{word_chunk} {word}".strip()
                if len(candidate) <= max_chars: word_chunk = candidate
                else:
                    if word_chunk: chunks.append(word_chunk)
                    word_chunk = word
            if word_chunk: chunks.append(word_chunk)
            continue
        candidate = f"{current} {sentence}".strip()
        if candidate and len(candidate) <= max_chars: current = candidate
        else:
            if current: chunks.append(current)
            current = sentence
    if current: chunks.append(current)
    return chunks

def call_ollama_generate(args, prompt):
    payload = {"model": args.post_correct_model, "prompt": prompt, "stream": False}
    try:
        resp = requests.post("http://127.0.0.1:11434/api/generate", json=payload, timeout=10)
        return resp.json().get("response", "").strip()
    except: return ""

def structure_paragraph_with_llm(text, args):
    prompt = f"Format the following transcribed segments into a coherent paragraph. Output ONLY the paragraph:\n{text}"
    return call_ollama_generate(args, prompt) or text

def post_correct_with_llm(text, args):
    prompt = f"Correct grammar and typos in this caption. Output ONLY the corrected text:\n{text}"
    return call_ollama_generate(args, prompt) or text

def run_translation(in_queue, out_queue, args):
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    translation_engines = {}
    for lang_key, model_id in TARGET_LANGS.items():
        if getattr(args, lang_key, False):
            print(f"[Translate] Loading {lang_key} model...")
            translation_engines[lang_key] = (
                MarianMTModel.from_pretrained(model_id).to(device),
                MarianTokenizer.from_pretrained(model_id)
            )

    paragraph_buffer = []
    
    while True:
        try:
            item = in_queue.get()
            if item is None: break
            
            # Draft support
            if "draft" in item:
                out_queue.put({"draft": item["draft"]})
                continue

            text = item.get("original", "").strip()
            detected_lang = item.get("detected_lang", "en")
            
            if args.post_correct and args.post_correct_llm:
                text = post_correct_with_llm(text, args)

            paragraph_buffer.append(text)
            
            # Paragraph logic
            if not args.llm_paragraph or text.endswith((".", "?", "!")):
                full_text = " ".join(paragraph_buffer)
                
                if args.llm_paragraph:
                    full_text = structure_paragraph_with_llm(full_text, args)

                payload = {"original": full_text, "ts": item.get("ts", time.time())}
                if args.en: payload["en"] = full_text
                
                for lang_key, (model, tokenizer) in translation_engines.items():
                    chunks = split_text_for_translation(full_text, args.mt_max_chars)
                    translated_parts = []
                    for chunk in chunks:
                        inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
                        with torch.no_grad():
                            translated_tokens = model.generate(**inputs, max_new_tokens=args.mt_max_new_tokens, num_beams=args.mt_num_beams)
                        translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
                    payload[lang_key] = " ".join(translated_parts)

                out_queue.put(payload)
                paragraph_buffer = []
                
        except Exception as e: print(f"[Translate] Error: {e}")

if __name__ == "__main__":
    import multiprocessing
    run_translation(multiprocessing.Queue(), multiprocessing.Queue(), argparse.Namespace())