import os
import re
import json
import sys
import time
from typing import NamedTuple
from pptx import Presentation
from pptx.enum.shapes import MSO_SHAPE_TYPE
import anthropic


# Patterns to never translate
SKIP_PATTERNS = [
    re.compile(r'^\s*$'),                          # empty / whitespace-only
    re.compile(r'^\d+$'),                          # pure slide numbers
    re.compile(r'^https?://'),                     # URLs
    re.compile(r'^[➔+\-•·…©®™°×÷±≈≤≥≠∞∑∫√π←→↑↓↔⇒⇔]$'),  # single symbols
    re.compile(r'^[a-zA-Z0-9\s]$'),               # single char/digit
]

OOXML_NS = {
    'p':   'http://schemas.openxmlformats.org/presentationml/2006/main',
    'a':   'http://schemas.openxmlformats.org/drawingml/2006/main',
    'mc':  'http://schemas.openxmlformats.org/markup-compatibility/2006',
    'a14': 'http://schemas.microsoft.com/office/drawing/2010/main',
}
A_NS = '{http://schemas.openxmlformats.org/drawingml/2006/main}'


class TranslationConfig(NamedTuple):
    """Configuration for translation settings."""
    source_language: str
    target_language: str
    skip_exact: set
    spacing_fix_prefixes: list
    system_prompt: str


def load_config(path: str) -> dict:
    """Read and validate the JSON config file. Fails fast with clear error messages."""
    try:
        with open(path) as f:
            cfg = json.load(f)
    except FileNotFoundError:
        print(f"ERROR: Config file not found: {path}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"ERROR: Invalid JSON in config file: {e}")
        sys.exit(1)

    errors = []
    if "lecture" not in cfg:
        errors.append("Missing required section: 'lecture'")
    else:
        for key in ("source_language", "target_language", "description"):
            if key not in cfg["lecture"]:
                errors.append(f"Missing required key: lecture.{key}")

    if "files" not in cfg or not isinstance(cfg["files"], list):
        errors.append("Missing or invalid 'files' section (must be a list)")
    else:
        for i, entry in enumerate(cfg["files"]):
            if "source" not in entry or "target" not in entry:
                errors.append(f"files[{i}] must have 'source' and 'target' keys")

    if "translation" not in cfg:
        errors.append("Missing required section: 'translation'")

    if errors:
        print("ERROR: Config validation failed:")
        for e in errors:
            print(f"  - {e}")
        sys.exit(1)

    return cfg


def build_translation_system_prompt(cfg: dict) -> tuple:
    """Build system prompt dynamically from config values."""
    lecture = cfg["lecture"]
    translation = cfg["translation"]

    src = lecture["source_language"]
    tgt = lecture["target_language"]
    desc = lecture["description"]

    keep_terms = translation.get("keep_in_source_language", [])
    forced = translation.get("forced_translations", {})

    keep_block = ", ".join(keep_terms) if keep_terms else "(none)"
    forced_lines = "\n".join(f"   - {k} → {v}" for k, v in forced.items())
    if not forced_lines:
        forced_lines = "   (none)"

    system_prompt = f"""You are a professional academic translator specializing in {desc}.
Translate the given {src} texts to formal academic {tgt} suitable for a university lecture.

Translation rules:
1. Keep the following terms in {src} (standard in {tgt} literature):
   {keep_block}

2. Translate these terms:
{forced_lines}

3. Keep author names, paper titles, and citations unchanged.
4. Keep acronyms as-is unless spelled out.
5. Use formal constructions appropriate for lectures.
6. Preserve capitalization style (e.g., if original is ALL CAPS, keep ALL CAPS in {tgt}).
7. Preserve any leading/trailing whitespace in the text.
8. SPACING: Always preserve spaces between words. Never merge two separate words into a
   compound without a hyphen. When an {src} term is combined with a {tgt} word, use a hyphen.
9. Every text MUST have a {tgt} translation. If most words are already the same in {tgt},
   at least translate the non-technical {src} words (e.g., articles, prepositions, verbs).
   Do NOT return the original {src} text unchanged unless it is a pure proper noun or acronym.
"""

    return system_prompt


def build_init_prompt(pptx_paths: list, target_language: str, text_dump: str) -> str:
    filenames = [os.path.basename(p) for p in pptx_paths]
    filenames_str = ", ".join(filenames)
    schema_example = '''{
  "lecture": {
    "description": "<5-10 word subject summary>",
    "source_language": "<detected language>",
    "target_language": "''' + target_language + '''"
  },
  "files": [
    {"source": "lecture.pptx", "target": "lecture_XX.pptx"}
  ],
  "translation": {
    "skip_exact": ["Acronym1", "AuthorName", "0", "1", "k", "n"],
    "keep_in_source_language": ["TechnicalTerm1", "TechnicalTerm2"],
    "forced_translations": {
      "SourceTerm": "TargetTranslation"
    },
    "spacing_fix_prefixes": ["Prefix1", "Prefix2"]
  }
}'''

    return f"""Analyze the following lecture slide texts extracted from: {filenames_str}

Your task is to produce a JSON configuration file for a translation tool that will translate these \
slides from their source language to {target_language}.

Instructions:
1. Detect the source language from the slide content.
2. Write a 5-10 word description of the lecture subject for "lecture.description".
3. For "files": derive each entry from the input filenames. Replace the .pptx extension \
with _XX.pptx for the target, where XX is the two-letter ISO 639-1 code for {target_language} \
(e.g. DE for German, FR for French, ES for Spanish, IT for Italian, PT for Portuguese, \
NL for Dutch, PL for Polish, SV for Swedish).
4. For "skip_exact": list strings that must NEVER be translated — acronyms found in the slides \
(e.g. CNN, SVM, DBSCAN), proper nouns and author names, single letters used as variables \
(k, n, d, m, i, j), single digits (0-9), and operators (+, -, ×).
5. For "keep_in_source_language": list multi-word technical terms that are conventionally kept \
in the source language even in {target_language} academic writing. Include terms you find in \
the slides that are standard in the field.
6. For "forced_translations": map source-language terms that have a specific, well-known \
{target_language} equivalent. Only include terms where the translation is unambiguous \
and widely used.
7. For "spacing_fix_prefixes": list source-language terms that commonly get merged with \
{target_language} words without a separator when the model translates (e.g. "Clustering" \
becoming "ClusteringAlgorithmen" instead of "Clustering-Algorithmen").

Return ONLY the JSON object — no explanation, no markdown fences, no preamble.

Schema to follow exactly:
{schema_example}

Slide texts to analyze:
{text_dump}"""


def should_skip(text: str, skip_exact: set) -> bool:
    """Determine if a text should be skipped based on exact matches and regex patterns."""
    if text in skip_exact:
        return True
    for pat in SKIP_PATTERNS:
        if pat.match(text):
            return True
    return False


def extract_from_text_frame(tf, location: tuple, slide_idx: int, file: str, records: list):
    """Extract paragraphs from a text frame.

    Paragraphs WITHOUT \x0b (no soft line-breaks): translate as whole, collapse runs.
    Paragraphs WITH \x0b (<a:br/> elements): translate each run independently so the
    existing <a:br/> XML elements stay in place and run text never contains \x0b.
    """
    for para_idx, para in enumerate(tf.paragraphs):
        runs = para.runs
        if not runs:
            continue  # no editable runs

        has_linebreaks = '\x0b' in para.text

        if has_linebreaks:
            # Per-run mode: keep <a:br/> structure, translate each run separately
            for run_idx, run in enumerate(runs):
                text = run.text
                if not text or not text.strip():
                    continue
                records.append({
                    "mode": "per_run",
                    "file": file,
                    "slide_idx": slide_idx,
                    "location": location,
                    "para_idx": para_idx,
                    "run_idx": run_idx,
                    "text": text,
                })
        else:
            # Whole-paragraph mode: translate combined run text, then collapse runs
            text = para.text  # == ''.join(r.text for r in runs) when no \x0b
            if not text or not text.strip():
                continue
            records.append({
                "mode": "whole",
                "file": file,
                "slide_idx": slide_idx,
                "location": location,
                "para_idx": para_idx,
                "text": text,
            })


def extract_from_shapes(shapes, slide_idx: int, file: str, records: list, shape_prefix=()):
    """Recursively extract text from shapes, handling groups and tables."""
    for shape_idx, shape in enumerate(shapes):
        current_path = shape_prefix + (shape_idx,)

        if shape.shape_type == MSO_SHAPE_TYPE.GROUP:
            extract_from_shapes(shape.shapes, slide_idx, file, records, current_path)
            continue

        if shape.has_text_frame:
            location = ("tf", current_path)
            extract_from_text_frame(shape.text_frame, location, slide_idx, file, records)

        if shape.has_table:
            table = shape.table
            for row_idx, row in enumerate(table.rows):
                for col_idx, cell in enumerate(row.cells):
                    location = ("table", current_path, row_idx, col_idx)
                    extract_from_text_frame(cell.text_frame, location, slide_idx, file, records)


def extract_altcontent_shapes(slide, slide_idx: int, file: str, records: list):
    """Extract text from shapes hidden inside <mc:AlternateContent> wrappers.

    python-pptx's slide.shapes iterator skips these elements entirely.
    We walk the raw lxml tree to find them and store direct element references
    so write-back can modify them in-place without a complex index path.
    """
    sp_tree = slide._element.find('.//p:spTree', OOXML_NS)
    if sp_tree is None:
        return
    for alt in sp_tree.findall('mc:AlternateContent', OOXML_NS):
        choice = alt.find('mc:Choice', OOXML_NS)
        if choice is None:
            continue
        for sp in choice.findall('p:sp', OOXML_NS):
            txBody = sp.find('p:txBody', OOXML_NS)
            if txBody is None:
                continue
            for para_elem in txBody.findall('a:p', OOXML_NS):
                runs = para_elem.findall('a:r', OOXML_NS)
                if not runs:
                    continue
                has_br = para_elem.find('a:br', OOXML_NS) is not None
                has_eq = para_elem.find('a14:m', OOXML_NS) is not None
                if has_br or has_eq:
                    for run_idx, run in enumerate(runs):
                        t = run.find('a:t', OOXML_NS)
                        text = t.text if t is not None else ''
                        if text and text.strip():
                            records.append({
                                "mode": "xml_per_run",
                                "file": file,
                                "slide_idx": slide_idx,
                                "para_elem": para_elem,
                                "run_idx": run_idx,
                                "text": text,
                            })
                else:
                    text = ''.join(
                        (r.find('a:t', OOXML_NS).text or '')
                        for r in runs
                        if r.find('a:t', OOXML_NS) is not None
                    )
                    if text and text.strip():
                        records.append({
                            "mode": "xml_whole",
                            "file": file,
                            "slide_idx": slide_idx,
                            "para_elem": para_elem,
                            "text": text,
                        })


def extract_records(prs: Presentation, file: str) -> list:
    records = []
    for slide_idx, slide in enumerate(prs.slides):
        extract_from_shapes(slide.shapes, slide_idx, file, records)
        extract_altcontent_shapes(slide, slide_idx, file, records)
        try:
            notes_tf = slide.notes_slide.notes_text_frame
            location = ("notes",)
            extract_from_text_frame(notes_tf, location, slide_idx, file, records)
        except Exception:
            pass

    return records


def api_call_with_backoff(client: anthropic.Anthropic, **kwargs) -> anthropic.types.Message:
    """Call the Anthropic API with exponential backoff on overloaded/rate-limit errors."""
    delays = [5, 15, 30, 60]
    for attempt, delay in enumerate(delays + [None]):
        try:
            return client.messages.create(**kwargs)
        except (anthropic.APIStatusError, anthropic.RateLimitError) as e:
            if hasattr(e, 'status_code') and e.status_code not in (429, 529):
                raise
            if delay is None:
                raise
            print(f"    API overloaded/rate-limited (attempt {attempt+1}), retrying in {delay}s...")
            time.sleep(delay)


def translate_batch(client: anthropic.Anthropic, texts: list[str],
                    model: str, system: str, src_lang: str, tgt_lang: str,
                    prefixes: list) -> dict[str, str]:
    """Translate a list of strings using XML-tagged numbered format.

    Per-run texts never contain \\x0b (handled at extraction time), so no placeholder needed.
    """
    numbered = '\n'.join(f"{i+1}. {t}" for i, t in enumerate(texts))
    user_msg = (
        f"Translate these {len(texts)} {src_lang} texts to {tgt_lang}.\n"
        "Return ONLY XML tags, one per line, in this exact format:\n"
        '<t id="1">translation 1</t>\n'
        '<t id="2">translation 2</t>\n'
        "...\n\n"
        f"{src_lang} texts:\n"
        f"{numbered}"
    )

    response = api_call_with_backoff(
        client,
        model=model,
        max_tokens=8192,
        system=system,
        messages=[{"role": "user", "content": user_msg}],
    )

    raw = response.content[0].text

    result = {}
    # Parse the XML-like response to extract translations by ID, then fix spacing issues.
    # Could be replaced by actual XML parsing if we enforce stricter formatting,
    # but regex is simpler for this controlled format. Would lead to failures if slides
    # contained literal <t id="..."> tags.
    for m in re.finditer(r'<t id="(\d+)">(.*?)</t>', raw, re.DOTALL):
        idx = int(m.group(1)) - 1
        translation = m.group(2)
        if 0 <= idx < len(texts):
            result[texts[idx]] = fix_spacing(texts[idx], translation, prefixes)

    return result


def translate_with_retry(client: anthropic.Anthropic, model: str, texts: list[str], cache: dict,
                         tcfg: TranslationConfig):
    """Translate texts with automatic fallback to smaller batches on failure."""
    if not texts:
        return

    result = translate_batch(
        client, texts,
        model, tcfg.system_prompt,
        tcfg.source_language, tcfg.target_language,
        tcfg.spacing_fix_prefixes,
    )
    cache.update(result)

    missing = [t for t in texts if t not in result]
    if not missing:
        return

    if len(texts) == 1:
        # Single text that still failed — keep original
        print(f"    Keeping original (translation failed): {repr(texts[0])[:60]}")
        cache[texts[0]] = texts[0]
        return

    if len(missing) < len(texts):
        # Partial success — retry only missing ones individually
        print(f"    {len(missing)} texts missing, retrying individually...")
        for t in missing:
            translate_with_retry(client, model, [t], cache, tcfg)
    else:
        # Complete failure — split batch in half
        mid = len(texts) // 2
        print(f"    All {len(texts)} missing, splitting into {mid} + {len(texts)-mid}...")
        translate_with_retry(client, model, texts[:mid], cache, tcfg)
        translate_with_retry(client, model, texts[mid:], cache, tcfg)


def fix_spacing(original: str, translation: str, prefixes: list[str]) -> str:
    """Post-process a translation to fix common merged-word artifacts.

    When the model merges English+target-language words without a separator (e.g.
    'ClusteringAlgorithmen'), insert a hyphen at the boundary.
    """
    for prefix in prefixes:
        # Pattern: PrefixTargetWord (capital after the prefix, no separator)
        pattern = re.compile(r'\b' + re.escape(prefix) + r'([A-ZÄÖÜ][a-zäöüß])')
        replacement = prefix + r'-\1'
        translation = pattern.sub(replacement, translation)

    return translation


def build_translation_cache(records: list, client: anthropic.Anthropic, model: str,
                            batch_size: int, tcfg: TranslationConfig) -> dict[str, str]:
    """Collect unique translatable texts, send to API in batches, return cache."""
    unique_texts = []
    seen = set()
    for rec in records:
        t = rec["text"]
        if t not in seen and not should_skip(t, tcfg.skip_exact):
            unique_texts.append(t)
            seen.add(t)

    print(f"  Unique texts to translate: {len(unique_texts)}")

    cache = {}

    for i in range(0, len(unique_texts), batch_size):
        batch = unique_texts[i : i + batch_size]
        batch_num = i // batch_size + 1
        total_batches = (len(unique_texts) + batch_size - 1) // batch_size
        print(f"  Translating batch {batch_num}/{total_batches} ({len(batch)} texts)...")
        translate_with_retry(client, model, batch, cache, tcfg)

    return cache


def get_text_frame_at(prs: Presentation, slide_idx: int, location: tuple):
    """Retrieve the text frame at the given location."""
    loc_type = location[0]

    if loc_type == "notes":
        return prs.slides[slide_idx].notes_slide.notes_text_frame

    if loc_type == "tf":
        shape_path = location[1]
        shapes = prs.slides[slide_idx].shapes
        shape = navigate_shapes(shapes, shape_path)
        return shape.text_frame

    if loc_type == "table":
        shape_path = location[1]
        row_idx = location[2]
        col_idx = location[3]
        shapes = prs.slides[slide_idx].shapes
        shape = navigate_shapes(shapes, shape_path)
        return shape.table.rows[row_idx].cells[col_idx].text_frame

    raise ValueError(f"Unknown location type: {loc_type}")


def navigate_shapes(shapes, path: tuple):
    """Navigate nested shape groups using an index path."""
    shape = shapes[path[0]]
    for idx in path[1:]:
        shape = shape.shapes[idx]
    return shape


def write_whole(tf, para_idx: int, translation: str):
    """Whole-paragraph write: set runs[0] to full translation, delete the rest."""
    para = tf.paragraphs[para_idx]
    runs = para.runs
    if not runs:
        return
    runs[0].text = translation
    for run in runs[1:]:
        para._p.remove(run._r)


def write_per_run(tf, para_idx: int, run_idx: int, translation: str):
    """Per-run write: update a single run in-place (no collapsing, no \x0b written)."""
    para = tf.paragraphs[para_idx]
    runs = para.runs
    if run_idx >= len(runs):
        return
    runs[run_idx].text = translation


def write_whole_xml(para_elem, translation: str):
    """Whole-paragraph write for AlternateContent shapes via lxml."""
    runs = para_elem.findall(f'{A_NS}r')
    if not runs:
        return
    t = runs[0].find(f'{A_NS}t')
    if t is not None:
        t.text = translation
    for run in runs[1:]:
        para_elem.remove(run)


def write_per_run_xml(para_elem, run_idx: int, translation: str):
    """Per-run write for AlternateContent shapes via lxml."""
    runs = para_elem.findall(f'{A_NS}r')
    if run_idx >= len(runs):
        return
    t = runs[run_idx].find(f'{A_NS}t')
    if t is not None:
        t.text = translation


def restore_surrounding_whitespace(original: str, translation: str) -> str:
    """Re-attach leading/trailing spaces that the model stripped from a translation.

    Text runs adjacent to equations (e.g. 'A feature map ') carry their spacing
    as literal spaces in the run text.  The translation model consistently strips
    these, breaking the visual gap between text and the inline equation symbol.
    """
    leading  = len(original) - len(original.lstrip(' '))
    trailing = len(original) - len(original.rstrip(' '))
    result   = translation.strip(' ')
    if leading:
        result = original[:leading] + result
    if trailing:
        result = result + original[len(original) - trailing:]
    return result


def apply_translations(prs: Presentation, records: list, cache: dict,
                       tcfg: TranslationConfig):
    """Apply all translations back to the presentation."""
    applied = 0
    skipped = 0

    for rec in records:
        text = rec["text"]
        if should_skip(text, tcfg.skip_exact):
            skipped += 1
            continue

        translation = cache.get(text)
        if translation is None:
            print(f"  WARNING: no translation for: {repr(text)[:60]}")
            skipped += 1
            continue

        translation = restore_surrounding_whitespace(text, translation)

        try:
            if rec["mode"] == "xml_whole":
                write_whole_xml(rec["para_elem"], translation)
            elif rec["mode"] == "xml_per_run":
                write_per_run_xml(rec["para_elem"], rec["run_idx"], translation)
            else:
                tf = get_text_frame_at(prs, rec["slide_idx"], rec["location"])
                if rec["mode"] == "per_run":
                    write_per_run(tf, rec["para_idx"], rec["run_idx"], translation)
                else:
                    write_whole(tf, rec["para_idx"], translation)
            applied += 1
        except Exception as e:
            print(f"  ERROR writing translation for {repr(text)[:40]}: {e}")
            skipped += 1

    print(f"  Applied: {applied}, Skipped/Kept: {skipped}")


def generate_config(client: anthropic.Anthropic, model: str, pptx_paths: list,
                    target_language: str) -> dict:
    """Extract text from PPTX files and ask Claude to generate an initial config."""
    all_texts = []
    for path in pptx_paths:
        print(f"  Extracting text from {os.path.basename(path)}...")
        prs = Presentation(path)
        records = extract_records(prs, os.path.basename(path))
        for rec in records:
            text = rec["text"].strip()
            if text:
                all_texts.append(text)

    unique_texts = list(dict.fromkeys(all_texts))  # deduplicate, preserve order
    text_dump = "\n".join(unique_texts)
    print(f"  {len(unique_texts)} unique text strings extracted.")

    print(f"  Calling {model} to generate config...")
    prompt = build_init_prompt(pptx_paths, target_language, text_dump)
    response = client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{"role": "user", "content": prompt}],
    )
    raw = response.content[0].text

    # Strip markdown fences if present
    raw = re.sub(r'^```(?:json)?\s*', '', raw.strip(), flags=re.MULTILINE)
    raw = re.sub(r'\s*```$', '', raw.strip(), flags=re.MULTILINE)

    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        print("ERROR: Claude returned invalid JSON. Raw response:")
        print(raw)
        print("\nTry re-running or use --model claude-sonnet-4-6 for better results.")
        sys.exit(1)


def audit_files(files: list, config_dir: str):
    """Print all extracted texts grouped by slide without making API calls."""
    for entry in files:
        src_path = os.path.join(config_dir, entry["source"])
        if not os.path.exists(src_path):
            print(f"WARNING: source file not found: {src_path}")
            continue

        prs = Presentation(src_path)
        file_label = os.path.basename(src_path)
        records = extract_records(prs, file_label)

        print(f"\n=== {file_label} — extracted texts ===")
        for rec in records:
            slide_num = rec["slide_idx"] + 1
            source = "xml " if rec["mode"].startswith("xml") else "pptx"
            print(f"  Slide {slide_num:2d} [source: {source}]:  {rec['text']}")

    print()
