"""
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
from typing import Callable

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

logger = logging.getLogger(__name__)

BASE_CHUNK_SIZE = 150   # default — may shrink on TokenLimitError
# 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
# ─────────────────────────────────────────────────────────────────────────────

_SYS_TRANSLATE = (
    "သင်သည် ရုပ်ရှင်/စီးရီး subtitle ဘာသာပြန်ကျွမ်းကျင်သူဖြစ်သည်။ "
    "ဘာသာပြန်ချက်များသည် မြန်မာလူများ နေ့စဉ်ပြောသကဲ့သို့ သဘာဝကျရမည်။\n\n"
    "RULES:\n"
    "1. Register — ဇာတ်ကောင်ပေါ်မူတည်၍ မင်း/ငါ/သင်/ကျွန်တော် ခွဲသုံးပါ\n"
    "2. Idioms — literal မဟုတ်ဘဲ မြန်မာ colloquial ဖြင့် ပြောင်းပါ\n"
    "3. Genre tone — Action(ကြမ်း) / Comedy(ဟာသ) / Romance(နွေးထွေး)\n"
    "4. Inline tags {\\an8}{\\b1}♪ — မပြောင်းရ\n"
    "5. Glossary terms — ပေးထားသည့်အတိုင်း တိကျစွာ သုံးပါ\n"
    "6. Empty/♪-only lines — ထိုအတိုင်း ပြန်ပါ\n\n"
    "OUTPUT: JSON array of EXACTLY n strings. No other text.\n"
    'Example: ["မြန်မာ ၁", "မြန်မာ ၂"]'
)

_SYS_GLOSSARY = (
    "Extract Myanmar subtitle translation glossary: "
    "character names, places, orgs, recurring slang.\n"
    'Return ONLY JSON: {"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


def _call_chunk_once(chunk: list, chunk_idx: int, total: int,
                     offset: int, glossary: dict,
                     provider: str, model: str, api_key: str) -> str:
    """Single API call for a chunk. Returns raw response string."""
    n      = len(chunk)
    system = _SYS_TRANSLATE
    gloss  = _glossary_block(glossary)
    if gloss:
        system = gloss + "\n\n" + system

    numbered = "\n".join(f"{i+1}. {line}" for i, line in enumerate(chunk))
    prompt = (
        f"Translate {n} subtitle lines to Myanmar.\n"
        f"[Chunk {chunk_idx+1}/{total} | Lines {offset+1}–{offset+n}]\n\n"
        f"INPUT ({n} lines):\n{numbered}\n\n"
        f"JSON array of EXACTLY {n} Myanmar strings:\n"
        f'["မြန်မာ ၁", ..., "မြန်မာ {n}"]\nOUTPUT:'
    )
    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,
    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)
        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,
            )
            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,
            )
            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:
            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:
                # Nothing translated — retry once more
                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) -> 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, 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,   # st.session_state for checkpointing
    resume: bool = False,                 # True = resume from checkpoint
) -> list:
    """
    Translate all subtitle text lines in chunks.
    Saves checkpoint after each chunk so translation can be resumed
    if interrupted by rate limits, quota, or network failure.

    Returns list[str] of EXACTLY len(lines) strings.
    Raises TranslationPaused if an unrecoverable limit is hit (quota/auth).
    """
    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 []

    chunks = _chunk(lines, BASE_CHUNK_SIZE)
    total  = len(chunks)

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

    # Pre-fill result from checkpoint
    result = list(lines)   # start with originals everywhere
    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:
        skipped = sum(len(chunks[i]) for i in done_chunks if i < len(chunks))
        logger.info("Resuming: %d chunks done (%d lines), %d remaining",
                    len(done_chunks), skipped, total - len(done_chunks))

    # ── Init checkpoint metadata ───────────────────────────────────────
    if "tx_checkpoint" not in session_state:
        session_state["tx_checkpoint"] = {
            "chunks_done":  [],
            "lines":        {},        # global_idx -> translated text
            "total_lines":  len(lines),
            "failed_chunks":[],
            "glossary":     dict(glossary),  # store for consistency on resume
            "chunk_size":   BASE_CHUNK_SIZE,
        }
    elif "glossary" not in session_state["tx_checkpoint"]:
        # Back-compat: add glossary if missing
        session_state["tx_checkpoint"]["glossary"] = dict(glossary)

    offset      = 0
    chunk_stats = []

    try:
        for i, chunk in enumerate(chunks):
            n = len(chunk)

            # Skip already-done chunks (resume mode)
            if i in done_chunks:
                offset += n
                chunk_stats.append((i+1, "SKIPPED", n, n))
                if progress_cb:
                    pct = int((i+1)/total*100)
                    progress_cb((i+1)/total,
                                f"Chunk {i+1}/{total} ✓ (checkpoint)  {pct}%")
                continue

            if progress_cb:
                progress_cb(
                    i / total,
                    f"Chunk {i+1}/{total}  ({n} lines | "
                    f"Lines {offset+1}–{offset+n})  ဘာသာပြန်နေသည်…"
                )

            # ── Translate chunk ────────────────────────────────────────
            translated = _translate_chunk_safe(
                chunk, i, total, offset, glossary,
                provider=provider, model=model, api_key=api_key,
                debug_log=debug_log,
            )

            # ── Length safety checkpoint ───────────────────────────────
            if len(translated) != n:
                translated = _enforce_length(translated, n, chunk)
            assert len(translated) == n

            # ── Write into result ──────────────────────────────────────
            for local_i, text in enumerate(translated):
                result[offset + local_i] = text

            # ── Save checkpoint ────────────────────────────────────────
            save_checkpoint(session_state, i, translated, offset)

            changed = sum(1 for o, t in zip(chunk, translated) if o != t)
            chunk_stats.append((i+1, "OK" if changed > 0 else "ORIGINAL", changed, n))
            offset += n

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

    except TranslationPaused:
        # Re-raise so app.py can show the resume UI
        raise

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

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

    return result


# ─────────────────────────────────────────────────────────────────────────────
# Background translation thread
# ─────────────────────────────────────────────────────────────────────────────

import threading as _threading

def run_translation_bg(
    job_id: str,
    lines: list,
    glossary: dict,
    *,
    provider: str,
    model: str,
    api_key: str,
    fmt: str,
    ssafile,            # pysubs2.SSAFile
    filename: str,
) -> _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),
            )

            # 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