import gradio as gr
from transformers import pipeline
import tempfile, os, zipfile, traceback

translator_cache = {}

MODEL_MAP = {
    ("en", "zh"): "Helsinki-NLP/opus-mt-en-zh",
    ("zh", "en"): "Helsinki-NLP/opus-mt-zh-en",
    ("en", "ja"): "Helsinki-NLP/opus-mt-en-ja",
    ("ja", "en"): "Helsinki-NLP/opus-mt-ja-en",
}

def get_translator(src_lang, tgt_lang):
    key = (src_lang, tgt_lang)
    if key not in translator_cache:
        if key in MODEL_MAP:
            translator_cache[key] = pipeline("translation", model=MODEL_MAP[key])
        else:
            raise ValueError(f"No model for {src_lang} to {tgt_lang}")
    return translator_cache[key]

def safe_translate(text, src, tgt):
    try:
        if (src, tgt) in MODEL_MAP:
            translator = get_translator(src, tgt)
            return translator(text, max_length=512)[0]["translation_text"]
        elif (src, tgt) == ("ja", "zh") or (src, tgt) == ("zh", "ja"):
            mid = safe_translate(text, src, "en")
            return safe_translate(mid, "en", tgt)
        else:
            return f"[Unsupported: {src}->{tgt}]"
    except Exception as e:
        return f"[Translation error: {str(e)}]"

def parse_srt(srt_text):
    blocks = srt_text.strip().split("\n\n")
    subtitles = []
    for block in blocks:
        lines = block.splitlines()
        if len(lines) >= 3:
            idx = lines[0]
            timestamp = lines[1]
            text = " ".join(lines[2:])
            subtitles.append((idx, timestamp, text))
    return subtitles

def reassemble_srt(subtitles):
    return "\n\n".join(f"{idx}\n{ts}\n{txt}" for idx, ts, txt in subtitles)

def process_file(file_obj, src_lang, tgt_lang, output_dir, error_log):
    try:
        with open(file_obj.name, "r", encoding="utf-8", errors="ignore") as f:
            raw_text = f.read()
        subtitles = parse_srt(raw_text)
        translated_subs = []

        for idx, ts, txt in subtitles:
            translated = safe_translate(txt, src_lang, tgt_lang)
            bilingual = f"{txt}\n{translated}"
            translated_subs.append((idx, ts, bilingual))

        output_path = os.path.join(output_dir, os.path.basename(file_obj.name))
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(reassemble_srt(translated_subs))

    except Exception as e:
        error_log.append(f"File {file_obj.name} failed: {str(e)}\n{traceback.format_exc()}")

def batch_translate(files, src_lang, tgt_lang):
    tmp_dir = tempfile.mkdtemp()
    error_log = []

    for file_obj in files:
        process_file(file_obj, src_lang, tgt_lang, tmp_dir, error_log)

    zip_path = os.path.join(tmp_dir, "translated_srt.zip")
    try:
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for name in os.listdir(tmp_dir):
                path = os.path.join(tmp_dir, name)
                if os.path.isfile(path) and name.endswith(".srt"):
                    zipf.write(path, arcname=name)
            if error_log:
                log_path = os.path.join(tmp_dir, "log.txt")
                with open(log_path, "w") as logf:
                    logf.write("\n".join(error_log))
                zipf.write(log_path, arcname="log.txt")
        return zip_path
    except Exception as e:
        fail_zip = os.path.join(tmp_dir, "fail.zip")
        with zipfile.ZipFile(fail_zip, 'w') as zipf:
            with open(os.path.join(tmp_dir, "log.txt"), "w") as logf:
                logf.write(f"ZIP error: {str(e)}\n\n{traceback.format_exc()}")
            zipf.write(os.path.join(tmp_dir, "log.txt"), arcname="log.txt")
        return fail_zip

gr.Interface(
    fn=batch_translate,
    inputs=[
        gr.File(file_types=[".srt"], label="Upload SRT files", file_count="multiple"),
        gr.Dropdown(["en", "zh", "ja"], label="Source Language", value="ja"),
        gr.Dropdown(["en", "zh", "ja"], label="Target Language", value="zh"),
    ],
    outputs=gr.File(label="Download Translated ZIP"),
    title="Batch SRT Translator (EN-ZH-JA)",
    description="Upload .srt subtitle files and translate between English, Chinese, and Japanese. Dual-language output with original + translation. ZIP output. Errors will be logged.",
).launch()