translation.py

"""
translation.py — Chunked Translation with Checkpoint / Resume

Checkpoint system
-----------------
After each chunk succeeds, results are written to a checkpoint dict in
st.session_state["tx_checkpoint"].  If the session is interrupted (rate
limit, network drop, user refresh), the next run can call
load_checkpoint() and resume from the last completed chunk.

Error handling per chunk
------------------------
RateLimitError  → wait + retry (api_handler handles this internally)
TokenLimitError → split chunk in half, retry each half separately
QuotaError      → pause everything, surface to user with clear message
AuthError       → surface to user immediately
NetworkError    → retry (handled in api_handler)
ParseError      → retry with simplified prompt, then fallback to originals

Timestamp safety — see translation.py header for full guarantee doc.
"""

import json
import logging
import re
import time
import threading
from typing import Callable

from api_handler import (
    call_ai,
    RateLimitError, TokenLimitError, QuotaError, AuthError,
)

logger = logging.getLogger(__name__)

# Adaptive chunk size per model RPM limit
# Small chunk = fewer merge errors but more API calls
# Large chunk = fewer calls but higher merge risk
BASE_CHUNK_SIZE = 50   # default (gemini-2.0-flash free: 15 RPM → ok)

def _get_chunk_size(model: str) -> int:
    """
    Return safe chunk size based on model RPM limit.
    Pro models have lower RPM on free tier → use larger chunks
    to avoid hitting rate limits with too many API calls.
    """
    m = (model or "").lower()
    # gemini-1.5-pro / 2.5-pro: free tier = 2 RPM → use 100
    if "pro" in m and "flash" not in m:
        return 100
    # gemini-1.5-flash-8b: slightly higher limit but slower response
    if "8b" in m or "lite" in m:
        return 80
    # gemini-2.0-flash / 1.5-flash / openrouter: 15+ RPM → use 50
    return 50


def _cooldown_seconds(model: str) -> float:
    """
    Return inter-request sleep in seconds to stay safely under RPM limits.
    Free tier limits (April 2026):
      Pro       →  5 RPM → 12s gap (+1s buffer = 13s)
      Flash     → 10 RPM →  6s gap (+0.5s buffer = 6.5s)
      Flash-Lite / 8b → 15 RPM → 4s gap (+0.5s buffer = 4.5s)
      OpenRouter / unknown → 3s (usually higher RPM)
    """
    m = (model or "").lower()
    if "pro" in m and "flash" not in m:
        return 13.0
    if "lite" in m or "8b" in m:
        return 4.5
    if "flash" in m:
        return 6.5
    return 3.0  # openrouter / unknown


# Glossary uses full file — Gemini 1M context handles it
# For very large files (5000+ lines) we spread-sample to stay within limits
GLOSSARY_MAX_LINES = 3000   # hard cap: beyond this we spread-sample
MAX_PARSE_RETRY = 2     # parse-only retries per chunk


# ─────────────────────────────────────────────────────────────────────────────
# Checkpoint helpers
# ─────────────────────────────────────────────────────────────────────────────

def save_checkpoint(state: dict, chunk_idx: int, translated: list,
                    offset: int) -> None:
    """Save per-chunk progress to session_state."""
    cp = state.setdefault("tx_checkpoint", {
        "chunks_done": [],
        "lines": {},          # global_index → translated_text
        "total_lines": 0,
        "failed_chunks": [],
    })
    cp["chunks_done"].append(chunk_idx)
    for local_i, text in enumerate(translated):
        cp["lines"][offset + local_i] = text


def load_checkpoint(state: dict) -> dict:
    """Return checkpoint dict or empty dict if none."""
    return state.get("tx_checkpoint", {})


def clear_checkpoint(state: dict) -> None:
    state.pop("tx_checkpoint", None)


def checkpoint_progress(state: dict, total_lines: int) -> tuple[int, int]:
    """Return (n_done, total) from checkpoint."""
    cp = load_checkpoint(state)
    done = len(cp.get("lines", {}))
    return done, total_lines


# ─────────────────────────────────────────────────────────────────────────────
# JSON / text extraction — 5 strategies
# ─────────────────────────────────────────────────────────────────────────────

def _parse_array(raw: str, n: int) -> list | None:
    text = raw.strip()

    # S1 — markdown fence strip then json.loads
    clean = re.sub(r"^```(?:json)?\s*", "", text, flags=re.IGNORECASE)
    clean = re.sub(r"\s*```$", "", clean).strip()
    try:
        r = json.loads(clean)
        if isinstance(r, list):
            return r
    except Exception:
        pass

    # S2 — first [...] block
    m = re.search(r"\[.*?\]", text, re.DOTALL)
    if m:
        try:
            r = json.loads(m.group())
            if isinstance(r, list):
                return r
        except Exception:
            pass

    # S3 — all "quoted strings"
    quoted = [q for q in re.findall(r'"((?:[^"\\]|\\.)*)"', text) if q.strip()]
    if quoted:
        return quoted

    # S4 — numbered/bulleted lines
    cleaned = []
    for ln in text.splitlines():
        ln = ln.strip()
        if not ln:
            continue
        ln = re.sub(r"^\d+[\.\)\-\:]\s*", "", ln)
        ln = re.sub(r"^[•\*\-]\s*", "", ln)
        ln = ln.strip("\"'")
        if ln:
            cleaned.append(ln)
    if cleaned:
        return cleaned

    # S5 — Myanmar Unicode lines
    myanmar_re = re.compile(r'[\u1000-\u109f]')
    my_lines = [l.strip() for l in text.splitlines()
                if myanmar_re.search(l) and l.strip()]
    if my_lines:
        return my_lines

    return None


def _parse_dict(raw: str) -> dict | None:
    text = raw.strip()
    clean = re.sub(r"^```(?:json)?\s*", "", text, flags=re.IGNORECASE)
    clean = re.sub(r"\s*```$", "", clean).strip()
    try:
        r = json.loads(clean)
        if isinstance(r, dict):
            return r
    except Exception:
        pass
    m = re.search(r"\{.*?\}", text, re.DOTALL)
    if m:
        try:
            r = json.loads(m.group())
            if isinstance(r, dict):
                return r
        except Exception:
            pass
    return None


def _enforce_length(lst: list, n: int, originals: list) -> list:
    """Always return exactly n items. Pad with originals if short."""
    result = [str(x) for x in lst]
    if len(result) > n:
        return result[:n]
    if len(result) < n:
        result += list(originals[len(result):n])
    return result


def _chunk(lst: list, size: int) -> list:
    return [lst[i: i + size] for i in range(0, len(lst), size)]


# ─────────────────────────────────────────────────────────────────────────────
# Prompts — Multi-Preset System (zero redundancy, each rule said once)
#
# Architecture: build_system_prompt() assembles 4 non-overlapping layers:
#   1. STYLE   — tone / register / swearing (preset-specific)
#   2. TERMS   — special terms to preserve (culture-specific)
#   3. NOUNS   — proper noun handling (noun_style × culture, generated once)
#   4. TECH    — count rule / tags / glossary / output / examples (always identical)
# ─────────────────────────────────────────────────────────────────────────────

# ── Core philosophy (shared, referenced only by TECH) ─────────────────────────
# Written once here for maintainability — injected into _TECH below.

# ── Layer 1: STYLE blocks (preset-specific) ───────────────────────────────────
# Covers: who is speaking, what register, how to handle strong language.
# Does NOT repeat: "don't add/remove" (that's TECH), nouns (that's NOUNS).
#
# DEFAULT PRONOUN RULE (applies to all presets):
#   ငါ/မင်း is the natural default for most dialogue.
#   ကျွန်တော်/ကျွန်မ only when the character is genuinely being servile/formal
#   (e.g. servant to emperor, soldier to commanding officer, student to professor).
#   Unknown gender → ငါ/မင်း. Ambiguous context → ငါ/မင်း.
#   မဆုံးဖြတ်နိုင်ရင် ငါ/မင်း သုံး — ကျွန်တော်/ကျွန်မ မသုံး။

_STYLE_STANDARD = """\
You are a Myanmar subtitle translator. Output natural spoken Myanmar — the way a real person talks, not the way a textbook translates.

VOICE: Reproduce the character's actual voice. Ask yourself: "မြန်မာလူတစ်ယောက် ဒီ situation မှာ တကယ် ဘာပြောမလဲ?"
DEFAULT pronouns — ငါ/မင်း for most characters. Only use ကျွန်တော်/ကျွန်မ when the character is genuinely being servile or very formal (servant → master, cadet → commander). Unknown gender → ငါ/မင်း.
• Rough / villain → ငါ/မင်း/ကောင် + ကွ/ဗျ/ကွာ (မဆဲရ — family-safe)
• Female casual → ငါ/ကျမ + ကွာ/ဟေ့/နော်
• Gentle / romantic → ငါ/မင်း + နူးညံ့သောဟန်
• Comic → ပြောပြောဆိုဆို ပေါ့ပေါ့ပါးပါး
Strong language → source intensity ကို family-safe မြန်မာ equivalent ဖြင့် ထိန်းသိမ်း\
"""

_STYLE_ADULT = """\
You are a Myanmar subtitle translator. Output natural spoken Myanmar — the way a real person talks, not the way a textbook translates.

VOICE: Reproduce the character's actual voice. Ask yourself: "မြန်မာလူတစ်ယောက် ဒီ situation မှာ တကယ် ဘာပြောမလဲ?"
DEFAULT pronouns — ငါ/မင်း for most characters. Only use ကျွန်တော်/ကျွန်မ when the character is genuinely being servile or formally subordinate. Unknown gender → ငါ/မင်း.
• Rough / gang / villain → ငါ/မင်း/ကောင်/ဟေ့ကောင် + ကွ/ဗျ/ကွာ
• Female casual / hostile → ငါ/နင် + ကွာ/ဟေ့/မ
• Dominant / boss → ငါ + ပြင်းထန်ဟန်
• Light / comic → ပြောပြောဆိုဆို ပေါ့ပေါ့ပါးပါး
Strong language — keep the exact weight of the source:
  English source → English as-is (Fuck / Shit / Bitch / Bastard / Asshole)
  Other language → nearest English equivalent\
"""

_STYLE_ANIME = """\
You are a Myanmar subtitle translator. Output natural spoken Myanmar — the way a real person talks, not the way a textbook translates.

VOICE: Reproduce the character's vocal personality. Ask yourself: "မြန်မာလူတစ်ယောက် ဒီ character လို ဘာပြောမလဲ?"
DEFAULT pronouns — ငါ/မင်း. Only use ကျွန်တော်/ကျွန်မ for genuinely servile/formal moments (student → sensei in apology, servant to lord).
• Energetic / shonen → ငါ/မင်း + ကွ/ဗျ အလေ့ကျ
• Tsundere → ပြတ်သားပြီး ဒေါသပျောက်ဝင် — ငါ/မင်း တင်းတင်း
• Kuudere / cold → တိုတို ပြတ်သား ရိုးရိုး
• Cute / moe → ငါ/ကျမ + particle: နော်/ရှင် + သံချို
• Villain → ကြမ်း ငါ/မင်း/ကောင်
Strong language → English as-is (Fuck / Shit / Bitch)\
"""

_STYLE_CDRAMA = """\
You are a Myanmar subtitle translator. Output natural spoken Myanmar — the way a real person talks, not the way a textbook translates.

VOICE: Preserve the dramatic weight and emotional layering of Chinese drama. Ask yourself: "မြန်မာလူတစ်ယောက် ဒီ scene မှာ ဘာပြောမလဲ?"
DEFAULT pronouns — ငါ/မင်း. ကျွန်တော်/ကျွန်မ only for genuine servile situations (servant before emperor, disciple before master in formal address).
• Commoner / young → ငါ/မင်း colloquial
• Villain / schemer → ငါ/မင်း + ကြမ်းတမ်းဟန်
• Noble in private → ငါ/မင်း + dignified
• Noble in court ceremony → ကျွန်တော် (genuine formal ritual only)
Poetic / dramatic lines → preserve emotional weight in natural Myanmar; ဟန်ဆန်ဆန် formal Myanmar မသုံး\
"""

_STYLE_KDRAMA = """\
You are a Myanmar subtitle translator. Output natural spoken Myanmar — the way a real person talks, not the way a textbook translates.

VOICE: Preserve K-drama emotional intensity and sincerity. Ask yourself: "မြန်မာလူတစ်ယောက် ဒီ feeling မှာ ဘာပြောမလဲ?"
DEFAULT pronouns — ငါ/မင်း. ကျွန်တော်/ကျွန်မ only for genuine formal subordinate moments (employee to CEO in official setting, child formal apology to elder).
• Young adult / peers → ငါ/မင်း + ကွာ/ဟေ့/ဗျ
• Emotional / romantic → ငါ/မင်း + weight ကိုသိမ်း
• Chaebol / CEO (private) → ငါ/မင်း + measured
• Elder / authority (official) → ကျွန်တော်/ကျမ (genuine only)
Strong language → English as-is\
"""

_STYLE_FORMAL = """\
You are a Myanmar subtitle translator for documentary, news, and educational content. Output clear, objective, educated standard Myanmar.

Register: Formal written Myanmar — ပညာတတ် standard. No colloquial ငါ/မင်း — use သင်/ကျွန်ုပ် or avoid pronouns where natural.
• Narration → formal, neutral, third-person where possible
• Expert dialogue → semi-formal educated register
• Interview subject → natural but measured
Technical terms → precise Myanmar equivalent; if none exists, keep (English term) in parentheses.
Numbers / dates / statistics → exact.\
"""

# ── Layer 2: TERMS blocks (culture-specific special terms) ────────────────────

_TERMS_WESTERN    = ""

_TERMS_JAPANESE   = """\
Japanese terms (keep as-is — these are NOT proper nouns, noun-style rule does not apply):
• Honorific suffixes: -san / -kun / -chan / -senpai / -sensei / -sama / -dono / -nii / -nee → keep or use natural Myanmar equivalent (Glossary wins)
• Anime/otaku: baka, kawaii, nakama, sugoi, yare yare, ara ara → English as-is
• Game/isekai system: HP, MP, Level, Skill, Quest, Guild, Status, Rank → English as-is
• Cultural: onsen, matsuri, bento, senpai-kouhai → English/Romaji as-is\
"""

_TERMS_CHINESE    = """\
Chinese terms (keep as-is — these are NOT proper nouns, noun-style rule does not apply):
• Relationship/title: Shifu, Shidi/Shixiong/Shijie, Gege/Jiejie/Meimei/Didi, A'niang, Wangye, Gongzi, Guniang → Pinyin as-is
• Wuxia/Xianxia: Golden Core, Nascent Soul, cultivation, qi, meridians, sect, immortal, spiritual root → English/Pinyin as-is
• Court titles: Emperor, Empress, Crown Prince, Noble Consort, Imperial Guard → English as-is\
"""

_TERMS_KOREAN     = """\
Korean terms (keep as-is — these are NOT proper nouns, noun-style rule does not apply):
• Relationship/honorifics: Oppa, Unnie, Hyung, Noona, Sunbae, Hoobae, Sajangnim, Seonsaengnim → Romanized as-is
• Speech level: 존댓말 → formal Myanmar register | 반말 → colloquial Myanmar register
• Cultural: Ramyeon, Tteokbokki, Hanbok, Jjimjilbang, Aegyo → Romanized as-is\
"""

_TERMS_THAI       = """\
Thai/SE Asian terms (keep as-is — these are NOT proper nouns, noun-style rule does not apply):
• Honorifics: Khun, Phi, Nong → Romanized as-is or natural Myanmar equivalent
• Cultural: Wai, Songkran, Muay Thai, Tuk-tuk → Romanized as-is\
"""

_TERMS_UNIVERSAL  = ""

# ── Layer 3: NOUNS block (generated from noun_style × culture) ────────────────

def _noun_block(style: str, culture_key: str) -> str:
    script = {
        "japanese":  "Romaji",
        "chinese":   "Pinyin",
        "korean":    "Romanized Korean",
        "thai":      "Romanized Thai",
        "western":   "English",
        "universal": "romanized",
    }.get(culture_key, "romanized")

    examples = {
        "western":   [("Victor","ဗစ်တာ"),            ("Hope Street","ဟုတ်ပ်စထရိတ်"), ("London","လန်ဒန်")],
        "japanese":  [("Sasuke","ဆဆူကေ"),             ("Tokyo","တိုကျို"),            ("Konoha","ကိုနိုဟာ")],
        "chinese":   [("Wei Wuxian","ဝေဝူရှျား"),     ("Beijing","ပေကျင်း"),          ("Yunmeng","ယွန်မိန်")],
        "korean":    [("Park Seo-jun","ပါ့ဆေဂျွန်"),  ("Seoul","ဆိုးလ်"),             ("Busan","ပူဆမ်")],
        "thai":      [("Prayuth","ပြာရတ်"),            ("Chiang Mai","ချင်းမိုင်"),    ("Bangkok","ဘန်ကောက်")],
        "universal": [("Victor","ဗစ်တာ"),             ("Sakura","ဆာကူရာ"),            ("Madison","မတ်ဒစ်ဆင်")],
    }.get(culture_key, [("Victor","ဗစ်တာ"), ("Hope Street","ဟုတ်ပ်စထရိတ်"), ("London","လန်ဒန်")])

    if style == "phonetic":
        ex_lines = "\n".join(f"  {e} → {p}" for e, p in examples)
        return (
            f"PROPER NOUNS — Phonetic Myanmar:\n"
            f"Names / places / organizations → Myanmar phonetic transliteration only (meaning-translation ห้าม)\n"
            f"{ex_lines}\n"
            f"Glossary ပါ → Glossary | မပါ → phonetic"
        )
    else:
        ex_lines = "\n".join(f'  "{e}" not "{p}"' for e, p in examples)
        return (
            f"PROPER NOUNS — {script} spelling:\n"
            f"Names / places / organizations → keep {script} original spelling unchanged\n"
            f"{ex_lines}\n"
            f"Glossary ပါ → Glossary | မပါ → {script} မူရင်း"
        )

# ── Layer 4: TECH block (always identical) ────────────────────────────────────

_TECH = """\
ABSOLUTE RULES:
1. မပိုထည့် မနုတ် — source မပါသောအရာ မဖြည့်ရ၊ ပါသောအရာ မပယ်ရ
2. မြန်မာ word order — source language sentence structure copy မကူးရ
3. ကိုယ်ရေး consistency — same character တစ်ကျပ်တည်း pronoun + style သုံး
4. ASS tags ({\\an8} {\\pos} {\\be5} {\\fad} etc.) + ♪ + ⏎ → as-is ထား
5. Glossary term ပါ → Glossary value ကိုသာ သုံး၊ ကွဲပြားသောဘာသာပြန် မသုံးရ

COUNT: Input N → output EXACTLY N. Merge / split ห้าม။

NATURAL vs AI — ဒီ pattern တေကို ရှောင်:
❌ "ဒါသည် ကောင်းမွန်သောရလဒ်တစ်ခုဖြစ်သည်"     ✅ "ဒါကတော့ ကောင်းတဲ့ outcome ပဲ"
❌ "သင်သည် ဤနေရာသို့ ဘာကြောင့်ရောက်လာသနည်း"  ✅ "မင်း ဒီမှာ ဘာလာလုပ်တာလဲ"
❌ "ကျွန်ုပ်ကို ကူညီပေးပါ ကျေးဇူးပြု၍"          ✅ "ကူညီပါ" / "ကူညီကွာ"
❌ "ကျွန်တော် သင့်ကို ချစ်ပါသည်"               ✅ "ငါ မင်းကို ချစ်တယ်"
❌ "ထိုနည်းအတိုင်း ကျွန်ုပ်တို့ ဆောင်ရွက်ရမည်"  ✅ "အဲ့လိုဆိုရင် ငါတို့ လုပ်ရမှာပဲ"

Output: JSON array, EXACTLY n items. No markdown, no preamble.
["ဘာသာပြန် ၁", "ဘာသာပြန် ၂", ..., "ဘာသာပြန် n"]\
"""

# ── Registry ──────────────────────────────────────────────────────────────────

PROMPT_PRESETS: dict = {
    "standard": {"label": "🎬 Standard",           "desc": "General / family-safe content",               "style": _STYLE_STANDARD},
    "adult_18": {"label": "🔞 18+ Adult",           "desc": "Mature content — strong language in English", "style": _STYLE_ADULT},
    "anime":    {"label": "⛩ Anime / J-Drama",     "desc": "Japanese anime, manga, light novels",         "style": _STYLE_ANIME},
    "cdrama":   {"label": "🏮 C-Drama / Movie",     "desc": "Chinese drama, wuxia, xianxia",               "style": _STYLE_CDRAMA},
    "kdrama":   {"label": "🌸 K-Drama / Movie",     "desc": "Korean drama, webtoon, Korean movies",        "style": _STYLE_KDRAMA},
    "formal":   {"label": "📰 Documentary / Formal","desc": "News, documentary, educational narration",    "style": _STYLE_FORMAL},
    "custom":   {"label": "✏️ Custom",              "desc": "ကိုယ်တိုင် system prompt ရေးထည့်ရတာ",       "style": ""},
}

SOURCE_CULTURES: dict = {
    "western":   {"label": "🌍 Western",          "desc": "English/European productions",       "terms": _TERMS_WESTERN},
    "japanese":  {"label": "🇯🇵 Japanese",         "desc": "Anime, J-Drama, Japanese movies",    "terms": _TERMS_JAPANESE},
    "chinese":   {"label": "🇨🇳 Chinese",          "desc": "C-Drama, Chinese movies, Donghua",   "terms": _TERMS_CHINESE},
    "korean":    {"label": "🇰🇷 Korean",           "desc": "K-Drama, Korean movies",             "terms": _TERMS_KOREAN},
    "thai":      {"label": "🇹🇭 Thai / SE Asian",  "desc": "Thai drama, SE Asian content",       "terms": _TERMS_THAI},
    "universal": {"label": "🌐 Unknown / Mixed",   "desc": "Multi-cultural or unknown origin",   "terms": _TERMS_UNIVERSAL},
}

PROPER_NOUN_STYLES: dict = {
    "english":  {
        "label": "🔤 English / Romanized (မူရင်း spelling)",
        "desc":  "Victor, Sakura, Hope Street — မြန်မာ phonetic မပြောင်း (default)",
    },
    "phonetic": {
        "label": "🔊 Myanmar Phonetic (အသံထွက်မြန်မာ)",
        "desc":  "Victor → ဗစ်တာ, Sakura → ဆာကူရာ, Hope Street → ဟုတ်ပ်စထရိတ်",
    },
}


def build_system_prompt(
    preset_key: str        = "adult_18",
    culture_key: str       = "western",
    custom_text: str       = "",
    proper_noun_style: str = "english",
) -> str:
    """
    Assemble final system prompt from 4 non-overlapping layers:
      STYLE → TERMS → NOUNS → TECH
    Each rule appears exactly once. No redundancy.
    """
    if preset_key == "custom":
        # Custom: user writes the full prompt; only TECH appended
        body = custom_text.strip() or PROMPT_PRESETS["adult_18"]["style"]
        return body + "\n\n" + _TECH

    style = PROMPT_PRESETS.get(preset_key, PROMPT_PRESETS["adult_18"])["style"]
    terms = SOURCE_CULTURES.get(culture_key, SOURCE_CULTURES["western"])["terms"]
    nouns = _noun_block(proper_noun_style, culture_key)

    parts = [style]
    if terms:
        parts.append(terms)
    parts.append(nouns)
    parts.append(_TECH)
    return "\n\n".join(parts)


def _sys_glossary_for_culture(
    culture_key: str       = "western",
    proper_noun_style: str = "english",
) -> str:
    """Culture + noun-style aware glossary extraction prompt."""
    noun_rule = _noun_block(proper_noun_style, culture_key)
    return (
        "Extract a Myanmar subtitle translation glossary.\n"
        "Include: character names, place names, organization names, "
        "recurring special terms / slang.\n\n"
        + noun_rule + "\n\n"
        "For recurring special terms / slang: Myanmar မြန်မာ translation as value.\n"
        'Return ONLY JSON: {"source term": "value"}'
    )


# Backward-compatible defaults
_SYS_TRANSLATE = build_system_prompt("adult_18", "western", proper_noun_style="english")
_SYS_GLOSSARY  = _sys_glossary_for_culture("western", "english")

def _glossary_block(glossary: dict) -> str:
    if not glossary:
        return ""
    lines = ["=== GLOSSARY (use exactly) ==="]
    for o, t in glossary.items():
        lines.append(f'  "{o}" → "{t}"')
    lines.append("=" * 30)
    return "\n".join(lines)


# ─────────────────────────────────────────────────────────────────────────────
# Single chunk translation — with error-type handling
# ─────────────────────────────────────────────────────────────────────────────

class TranslationPaused(Exception):
    """Raised when translation must stop (quota exhausted, auth error)."""
    def __init__(self, reason: str, is_resumable: bool = False):
        super().__init__(reason)
        self.is_resumable = is_resumable


# Placeholder used to encode subtitle-internal newlines in the numbered prompt.
# Prevents multi-line subtitles from confusing the AI's line count.
_NL = "⏎"


def _encode_nl(text: str) -> str:
    """Replace internal \n with placeholder so each subtitle = exactly 1 line."""
    return text.replace("\n", _NL)


def _decode_nl(text: str) -> str:
    """Restore placeholder to \n in translated text."""
    return text.replace(_NL, "\n")


def _strip_html(text: str) -> str:
    """Strip HTML tags (e.g. <i>, </i>, <b>, </b>) from subtitle text."""
    return re.sub(r'<[^>]+>', '', text)


# Pronoun substitution table — applied after translation.
# Each entry: (pattern_string, replacement)
# Skipped for "formal" / "custom" presets.
_PRONOUN_SUBS: list[tuple] = [
    # 1st person formal → natural
    (r'ကျွန်တော်',  'ငါ'),
    (r'ကျွန်မ',     'ငါ'),
    # 2nd person formal → natural
    (r'ခင်ဗျား',    'မင်း'),
]
# Compile once
_PRONOUN_RE: list[tuple] = [
    (re.compile(p), r) for p, r in _PRONOUN_SUBS
]

def _naturalize_pronouns(text: str) -> str:
    """
    Replace formal pronouns with natural spoken equivalents.
    ကျွန်တော် → ငါ  |  ကျွန်မ → ကျမ  |  ခင်ဗျား → မင်း
    Applied after translation for all non-formal presets.
    Does NOT touch ASS tags ({...}) or music markers.
    """
    # Don't touch content inside ASS override tag blocks
    # Split on { } boundaries to protect tags
    parts = re.split(r'(\{[^}]*\})', text)
    out = []
    for part in parts:
        if part.startswith('{'):
            out.append(part)   # ASS tag — untouched
        else:
            for pattern, replacement in _PRONOUN_RE:
                part = pattern.sub(replacement, part)
            out.append(part)
    return ''.join(out)


# Thread-local override for count-mismatch retry feedback
# _translate_chunk_safe sets [0] = error message; _call_chunk_once injects it
_chunk_once_override: list = [None]


def _call_chunk_once(chunk: list, chunk_idx: int, total: int,
                     offset: int, glossary: dict,
                     provider: str, model: str, api_key: str,
                     system_prompt: str | None = None) -> str:
    """
    Single API call for a chunk. Returns raw response string.
    system_prompt: pre-built prompt from build_system_prompt(); falls back to _SYS_TRANSLATE.
    """
    n      = len(chunk)
    system = system_prompt if system_prompt else _SYS_TRANSLATE
    gloss  = _glossary_block(glossary)
    if gloss:
        system = gloss + "\n\n" + system

    # Strip HTML tags, encode internal newlines → placeholder (one subtitle = one numbered line)
    encoded = [_encode_nl(_strip_html(line)) for line in chunk]
    numbered = "\n".join(f"{i+1}. {line}" for i, line in enumerate(encoded))

    # Build glossary reminder for user prompt (ensures AI sees it right before input)
    gloss_reminder = ""
    if glossary:
        gloss_reminder = (
            "GLOSSARY REMINDER — use these EXACT terms, no alternatives:\n"
            + "\n".join(f'  "{k}" → "{v}"' for k, v in glossary.items())
            + "\n  (Proper nouns NOT in glossary → keep original English spelling)\n\n"
        )

    prompt = (
        f"Translate {n} subtitle lines to Myanmar.\n"
        f"[Chunk {chunk_idx+1}/{total} | Lines {offset+1}–{offset+n}]\n\n"
        f"{gloss_reminder}"
        f"INPUT ({n} lines):\n{numbered}\n\n"
        f"JSON array of EXACTLY {n} Myanmar strings:\n"
        f'["မြန်မာ ၁", ..., "မြန်မာ {n}"]\nOUTPUT:'
    )
    # Inject count-error feedback from previous attempt if set
    override_msg = _chunk_once_override[0]
    if override_msg:
        prompt = override_msg + "\n\n" + prompt
        _chunk_once_override[0] = None   # consume once

    return call_ai(prompt, provider=provider, model=model,
                   api_key=api_key, system_prompt=system, timeout=180)


def _translate_chunk_safe(
    chunk: list,
    chunk_idx: int,
    total_chunks: int,
    offset: int,
    glossary: dict,
    *,
    provider: str,
    model: str,
    api_key: str,
    debug_log: list,
    system_prompt: str | None = None,
    chunk_size_override: int | None = None,
) -> list:
    """
    Translate one chunk with full error handling.
    Returns exactly len(chunk) strings.
    Raises TranslationPaused for unrecoverable errors (quota, auth).
    """
    n = len(chunk)

    # ── Try normal call ────────────────────────────────────────────────
    raw       = None
    last_err  = None

    for parse_attempt in range(1, MAX_PARSE_RETRY + 1):
        try:
            raw = _call_chunk_once(chunk, chunk_idx, total_chunks,
                                   offset, glossary, provider, model, api_key,
                                   system_prompt=system_prompt)
        except QuotaError as e:
            # Daily quota done — must stop, checkpoint already saved
            debug_log.append({"chunk": chunk_idx+1, "status": "QUOTA_EXHAUSTED",
                               "error": str(e)})
            raise TranslationPaused(
                f"Quota ကုန်ပြီ — {e}\n\n"
                "ကျန်သော chunks များကို checkpoint မှ ဆက်ပြန်နိုင်သည်",
                is_resumable=True
            ) from e

        except AuthError as e:
            debug_log.append({"chunk": chunk_idx+1, "status": "AUTH_ERROR",
                               "error": str(e)})
            raise TranslationPaused(f"API key မမှန်ပါ — {e}",
                                    is_resumable=False) from e

        except TokenLimitError:
            # Chunk too big — split in half and retry each half
            logger.warning("Chunk %d: token limit — splitting in half", chunk_idx+1)
            debug_log.append({"chunk": chunk_idx+1, "status": "TOKEN_LIMIT_SPLIT"})
            mid  = n // 2
            half1 = _translate_chunk_safe(
                chunk[:mid], chunk_idx, total_chunks, offset,
                glossary, provider=provider, model=model,
                api_key=api_key, debug_log=debug_log,
                system_prompt=system_prompt,
            )
            half2 = _translate_chunk_safe(
                chunk[mid:], chunk_idx, total_chunks, offset + mid,
                glossary, provider=provider, model=model,
                api_key=api_key, debug_log=debug_log,
                system_prompt=system_prompt,
            )
            result = half1 + half2
            return _enforce_length(result, n, chunk)

        except RateLimitError as e:
            # api_handler already retried — if still raised, it's exhausted
            debug_log.append({"chunk": chunk_idx+1, "status": "RATE_LIMIT_EXHAUSTED",
                               "error": str(e)})
            raise TranslationPaused(
                f"Rate limit ကုန်ပြီ — {e}\n\nCheckpoint မှ ဆက်ပြန်နိုင်သည်",
                is_resumable=True
            ) from e

        except Exception as e:
            last_err = e
            debug_log.append({"chunk": chunk_idx+1, "status": "API_ERROR",
                               "attempt": parse_attempt, "error": str(e)})
            if parse_attempt < MAX_PARSE_RETRY:
                time.sleep(3)
                continue
            break

        # ── Parse the raw response ─────────────────────────────────────
        parsed = _parse_array(raw, n)
        if parsed is not None:
            # Decode ⏎ placeholder back to \n (multi-line subtitle restore)
            parsed = [_decode_nl(str(p)) for p in parsed]

            # ── COUNT VALIDATION (root cause of timestamp shift) ──────────
            # If AI returned wrong count, retry with explicit error feedback
            if len(parsed) != n and parse_attempt < MAX_PARSE_RETRY:
                wrong = len(parsed)
                debug_log.append({
                    "chunk": chunk_idx+1, "status": "COUNT_MISMATCH",
                    "expected": n, "got": wrong, "attempt": parse_attempt,
                })
                # Rebuild prompt with explicit error message
                if wrong < n:
                    count_err = (
                        f"CORRECTION NEEDED: You returned {wrong} items "
                        f"but I sent {n}. You MERGED some lines. "
                        f"Please re-translate keeping EXACTLY {n} separate items. "
                        f"Even short lines must remain separate."
                    )
                else:
                    count_err = (
                        f"CORRECTION NEEDED: You returned {wrong} items "
                        f"but I sent {n}. You SPLIT some lines. "
                        f"Please re-translate keeping EXACTLY {n} separate items. "
                        f"Do not split any single line into multiple items."
                    )
                # Inject error into next attempt's prompt
                _chunk_once_override[0] = count_err
                time.sleep(1)
                continue

            result = _enforce_length(parsed, n, chunk)
            assert len(result) == n
            changed = sum(1 for o, t in zip(chunk, result) if o != t)
            debug_log.append({
                "chunk": chunk_idx+1, "status": "OK",
                "translated": changed, "total": n,
                "raw_preview": raw[:200],
            })
            if changed == 0 and parse_attempt < MAX_PARSE_RETRY:
                debug_log[-1]["status"] = "RETRY_ZERO_CHANGE"
                time.sleep(2)
                continue
            return result
        else:
            debug_log.append({
                "chunk": chunk_idx+1, "status": "PARSE_FAIL",
                "attempt": parse_attempt, "raw_preview": (raw or "")[:400],
            })
            if parse_attempt < MAX_PARSE_RETRY:
                time.sleep(2)
                continue

    # All attempts failed — fallback to originals (length-safe)
    logger.error("Chunk %d: all attempts failed. Keeping originals.", chunk_idx+1)
    debug_log.append({"chunk": chunk_idx+1, "status": "FALLBACK_ORIGINAL"})
    return list(chunk[:n])


# ─────────────────────────────────────────────────────────────────────────────
# Glossary extraction
# ─────────────────────────────────────────────────────────────────────────────

def extract_glossary(lines: list, *, provider, model, api_key,
                     culture_key: str = "western",
                     proper_noun_style: str = "english") -> dict:
    """
    Extract translation glossary from the FULL subtitle file.
    Sends all non-empty lines so terms appearing late in the file
    (episode-specific locations, late-introduced characters, slang) are captured.

    For very large files (>GLOSSARY_MAX_LINES), we spread-sample:
      - First 1/3 + Middle 1/3 + Last 1/3
    This preserves coverage across the whole file while staying within token limits.
    Gemini 1.5/2.0 handles 1M tokens so most files fit entirely.
    """
    non_empty = [l for l in lines if l.strip()]
    total     = len(non_empty)

    if total <= GLOSSARY_MAX_LINES:
        # ── Full file ──────────────────────────────────────────────────
        sample = non_empty
    else:
        # ── Spread sample: beginning + middle + end ────────────────────
        third   = GLOSSARY_MAX_LINES // 3
        mid_s   = (total // 2) - (third // 2)
        sample  = (
            non_empty[:third] +
            non_empty[mid_s: mid_s + third] +
            non_empty[-third:]
        )
        logger.info(
            "File has %d lines > %d limit — spread sampling %d lines",
            total, GLOSSARY_MAX_LINES, len(sample)
        )

    numbered = "\n".join(f"{i+1}. {l}" for i, l in enumerate(sample))
    prompt   = (
        f"Extract Myanmar translation glossary from these {len(sample)} subtitle lines "
        f"(full file: {total} lines total):\n\n{numbered}"
    )

    try:
        raw = call_ai(prompt, provider=provider, model=model, api_key=api_key,
                      system_prompt=_sys_glossary_for_culture(culture_key, proper_noun_style), timeout=120)
        r = _parse_dict(raw)
        if r:
            return {str(k): str(v) for k, v in r.items()}
    except Exception as exc:
        logger.warning("Glossary extraction failed: %s", exc)
    return {}


# ─────────────────────────────────────────────────────────────────────────────
# translate_all — main entry point with checkpoint/resume
# ─────────────────────────────────────────────────────────────────────────────

def translate_all(
    lines: list,
    glossary: dict,
    *,
    provider: str,
    model: str,
    api_key: str,
    progress_cb: Callable | None = None,
    debug_log: list | None = None,
    session_state: dict | None = None,
    resume: bool = False,
    cooldown: float | None = None,
    system_prompt: str | None = None,
    preset_key: str = "adult_18",   # used to decide pronoun post-processing
) -> list:
    """
    Translate subtitle text lines in chunks.

    EMPTY LINE HANDLING (critical for timestamp alignment):
    - Empty strings "" (Comment events, skip-style events) are filtered OUT
      before sending to AI. Their positions are tracked separately.
    - After translation, translated text is re-inserted at original positions.
    - This prevents AI from receiving empty lines, which could cause it to
      return fewer items → padding → position shift → wrong timestamp mapping.

    Returns list[str] of EXACTLY len(lines) strings.
    """
    if debug_log is None:
        debug_log = []
    if session_state is None:
        session_state = {}

    if not lines:
        if progress_cb:
            progress_cb(1.0, "ပြီးစီးသည်")
        return []

    # ── Separate translatable lines from empty/skip lines ─────────────────
    # Build two parallel structures:
    #   translatable_lines  : list of (original_index, text) for non-empty lines
    #   result              : final output, pre-filled with originals
    result = list(lines)  # start with originals everywhere

    translatable = []   # [(original_index, text), ...]
    for idx, line in enumerate(lines):
        if line.strip():   # non-empty = needs translation
            translatable.append((idx, line))
        # empty = keep original (already in result)

    if not translatable:
        if progress_cb:
            progress_cb(1.0, "ပြီးစီးသည် (ဘာသာပြန်ရမည့် lines မရှိ)")
        return result

    logger.info(
        "translate_all: %d total lines, %d translatable, %d skipped (empty/comment)",
        len(lines), len(translatable), len(lines) - len(translatable)
    )

    # ── Extract just the text for chunking ────────────────────────────────
    trans_texts   = [text for _, text in translatable]
    trans_indices = [idx  for idx, _ in translatable]

    # ── Load checkpoint ───────────────────────────────────────────────────
    cp = load_checkpoint(session_state)
    done_chunks  = set(cp.get("chunks_done", []))
    saved_lines  = cp.get("lines", {})

    # Apply saved translations
    for global_idx, text in saved_lines.items():
        if isinstance(global_idx, int) and global_idx < len(result):
            result[global_idx] = text

    if resume and done_chunks:
        logger.info("Resuming from checkpoint: %d chunks done", len(done_chunks))

    # Init checkpoint
    if "tx_checkpoint" not in session_state:
        session_state["tx_checkpoint"] = {
            "chunks_done":  [],
            "lines":        {},
            "total_lines":  len(lines),
            "failed_chunks": [],
            "glossary":     dict(glossary),
            "chunk_size":   BASE_CHUNK_SIZE,
        }
    elif "glossary" not in session_state["tx_checkpoint"]:
        session_state["tx_checkpoint"]["glossary"] = dict(glossary)

    # ── Chunk only the translatable lines ─────────────────────────────────
    chunk_size = _get_chunk_size(model)
    chunks = _chunk(list(range(len(trans_texts))), chunk_size)
    total  = len(chunks)
    chunk_stats = []

    try:
        for ci, chunk_indices in enumerate(chunks):
            # chunk_indices = positions within trans_texts / trans_indices

            if ci in done_chunks:
                chunk_stats.append((ci+1, "SKIPPED", len(chunk_indices), len(chunk_indices)))
                if progress_cb:
                    progress_cb((ci+1)/total, f"Chunk {ci+1}/{total} ✓ (checkpoint)  {int((ci+1)/total*100)}%")
                continue

            chunk_texts   = [trans_texts[i]   for i in chunk_indices]
            chunk_orig_ix = [trans_indices[i]  for i in chunk_indices]  # original event positions
            n = len(chunk_texts)

            if progress_cb:
                first_line = chunk_orig_ix[0] + 1
                last_line  = chunk_orig_ix[-1] + 1
                progress_cb(
                    ci / total,
                    f"Chunk {ci+1}/{total}  ({n} lines | Events {first_line}–{last_line})  ဘာသာပြန်နေသည်…"
                )

            # ── Translate this chunk ───────────────────────────────────
            translated_chunk = _translate_chunk_safe(
                chunk_texts, ci, total, chunk_orig_ix[0], glossary,
                provider=provider, model=model, api_key=api_key,
                debug_log=debug_log,
                system_prompt=system_prompt,
            )

            # Length enforcement
            if len(translated_chunk) != n:
                translated_chunk = _enforce_length(translated_chunk, n, chunk_texts)
            assert len(translated_chunk) == n

            # ── Re-insert at original event positions ──────────────────
            # This is the key: trans_indices[chunk_indices[j]] = original event index
            for j, orig_idx in enumerate(chunk_orig_ix):
                result[orig_idx] = translated_chunk[j]

            # ── Checkpoint save ────────────────────────────────────────
            checkpoint_entries = {orig_idx: translated_chunk[j]
                                  for j, orig_idx in enumerate(chunk_orig_ix)}
            cp_state = session_state["tx_checkpoint"]
            cp_state["chunks_done"].append(ci)
            cp_state["lines"].update(checkpoint_entries)

            changed = sum(1 for o, t in zip(chunk_texts, translated_chunk) if o != t)
            chunk_stats.append((ci+1, "OK" if changed > 0 else "ORIGINAL", changed, n))

            if progress_cb:
                pct = int((ci+1)/total*100)
                st_label = "✓" if changed > 0 else "⚠ original"
                progress_cb((ci+1)/total,
                            f"Chunk {ci+1}/{total} {st_label}  {changed}/{n}  {pct}%")

            # ── Cooldown between chunks to avoid RPM rate limit ────────────
            if ci < total - 1:   # skip sleep after last chunk
                wait = cooldown if cooldown is not None else _cooldown_seconds(model)
                if wait > 0:
                    if progress_cb:
                        progress_cb(
                            (ci+1)/total,
                            f"Chunk {ci+1}/{total} ✓  ⏳ cooldown {wait:.0f}s…"
                        )
                    time.sleep(wait)

    except TranslationPaused:
        raise

    finally:
        if progress_cb:
            done_n    = sum(c for _, s, c, _ in chunk_stats if s != "SKIPPED")
            fallbacks = sum(1 for _, s, c, _ in chunk_stats if c == 0 and s != "SKIPPED")
            progress_cb(
                1.0,
                f"ပြီးစီးသည် ✓  {done_n}/{len(translatable)} translatable lines"
                + (f"  |  {fallbacks} chunks original" if fallbacks else "")
            )

    # ── Post-process: naturalize pronouns (skip for formal/custom presets) ──
    _SKIP_NATURALIZE = {"formal", "custom"}
    if preset_key not in _SKIP_NATURALIZE:
        result = [_naturalize_pronouns(line) for line in result]

    # Final length guarantee
    assert len(result) == len(lines), \
        f"Length invariant broken: {len(result)} != {len(lines)}"

    return result



def run_translation_bg(
    job_id: str,
    lines: list,
    glossary: dict,
    *,
    provider: str,
    model: str,
    api_key: str,
    fmt: str,
    ssafile,
    filename: str,
    system_prompt: str | None = None,
    preset_key: str = "adult_18",
) -> threading.Thread:
    """
    Spawn background thread for translation.
    Survives browser disconnect — progress saved to job_store.
    On limit-hit (TranslationPaused): saves checkpoint, sets status=paused.
    User can resume from History page.
    """
    from job_store import update_job, save_result, load_job
    from subtitle_parser import replace_lines, write_subtitle, output_filename

    def _worker():
        debug_log  = []
        # Load existing checkpoint if resuming
        job        = load_job(job_id)
        cp_init    = job.get("checkpoint", {}) if job else {}
        session_st = {"tx_checkpoint": cp_init} if cp_init else {}

        def pcb(frac: float, msg: str) -> None:
            update_job(job_id, progress=frac, status_msg=msg)
            # Save checkpoint periodically
            cp = session_st.get("tx_checkpoint", {})
            if cp:
                update_job(job_id, checkpoint=cp)

        update_job(job_id, status="running", progress=0.0)

        try:
            translated = translate_all(
                lines, glossary,
                provider=provider, model=model, api_key=api_key,
                progress_cb=pcb,
                debug_log=debug_log,
                session_state=session_st,
                resume=bool(cp_init),
                system_prompt=system_prompt,
                preset_key=preset_key,
            )

            # Build output file
            out_subs  = replace_lines(ssafile, translated)
            out_bytes = write_subtitle(out_subs, fmt)
            save_result(job_id, out_bytes)   # sets status=completed

        except TranslationPaused as e:
            cp = session_st.get("tx_checkpoint", {})
            update_job(job_id,
                       status="paused",
                       checkpoint=cp,
                       status_msg=f"⏸ Limit hit — {str(e)[:200]}",
                       error=str(e)[:300])

        except Exception as e:
            update_job(job_id,
                       status="failed",
                       error=str(e)[:400],
                       status_msg=f"❌ {str(e)[:200]}")

    t = threading.Thread(
        target=_worker,
        daemon=True,
        name=f"translate-{job_id}",
    )
    t.start()
    return t