#!/usr/bin/env python3
"""
Balochi Latin translation API server for Syáhag.

Priority (understanding → English → target):
  1) dataset.xlsx / dataset.csv — sentence search
  2) words_balochi_to_english.xlsx — word senses
  3) Balochi_grammar_rules.pdf — grammar understanding
  4) AI (Groq) Balochi→English grounded by (1)–(3)
  5) AI / Google Translate English→selected language

TranslateGemma adapter is NOT used.
"""

from __future__ import annotations

import asyncio
import os
import re
import time
import unicodedata
from collections import Counter, defaultdict
from heapq import nlargest
from pathlib import Path

import pandas as pd
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from google import genai
from google.genai import types
from google.genai.errors import APIError
from pydantic import BaseModel, Field
from pypdf import PdfReader

# Understanding stack — no TranslateGemma adapter.
GENAI_SUPPORT = [
    "dataset",
    "lexicon",
    "grammar",
]
PRIORITIES = list(GENAI_SUPPORT)
DEFAULT_GROQ_MODEL = os.getenv("GROQ_MODEL", "openai/gpt-oss-120b")
GROQ_SUPPORT = [
    f"groq:{DEFAULT_GROQ_MODEL}",
    "google_translate",
]
# Adapter disabled for labzank pipeline (dataset/lexicon/grammar → EN → AI).
os.environ.setdefault("ENABLE_TRANSLATEGEMMA", "0")
# Skip LoRA drafts above these limits (long labzank posts time out on ZeroGPU).
TG_MAX_CHARS = int(os.getenv("TRANSLATEGEMMA_MAX_CHARS", "1200"))
TG_MAX_WORDS = int(os.getenv("TRANSLATEGEMMA_MAX_WORDS", "180"))
TG_MAX_LINES = int(os.getenv("TRANSLATEGEMMA_MAX_LINES", "6"))
# Chunk Gemini only for very long essays (many sequential calls can exceed 90s).
GENAI_CHUNK_CHARS = int(os.getenv("GENAI_CHUNK_CHARS", "5000"))
GENAI_CHUNK_PARAS = int(os.getenv("GENAI_CHUNK_PARAS", "10"))
TRANSLATEGEMMA_ADAPTER_DIR = os.getenv(
    "TRANSLATEGEMMA_ADAPTER",
    "Aynkader/translategemma-balochi-adapter",
)
TRANSLATEGEMMA_HUB = os.getenv(
    "TRANSLATEGEMMA_HUB",
    "Aynkader/translategemma-balochi-adapter",
)

# Optional .env support
try:
    from dotenv import load_dotenv

    load_dotenv(Path(__file__).resolve().parent / ".env")
except Exception:
    pass

ROOT = Path(__file__).resolve().parent
STATIC_DIR = ROOT / "static"
GRAMMAR_PDF = ROOT / "Balochi_grammar_rules.pdf"
GRAMMAR_CACHE = ROOT / "Balochi_grammar_rules.cache.txt"
DATASET_CSV = ROOT / "Datasets" / "dataset.csv"
DATASET_CSV_ROOT = ROOT / "dataset.csv"
DATASET_XLSX = ROOT / "dataset.xlsx"
WORDS_CSV = ROOT / "Datasets" / "words_balochi_to_english.csv"
WORDS_CSV_ROOT = ROOT / "words_balochi_to_english.csv"
WORDS_XLSX = ROOT / "words_balochi_to_english.xlsx"

app = FastAPI(title="Balochi Latin Translation API", version="2.0")

_cors_origins = [
    o.strip()
    for o in os.getenv(
        "CORS_ORIGINS",
        "https://gwachen.com,https://www.gwachen.com,https://api.gwachen.com,https://labzank.com,https://www.labzank.com,http://127.0.0.1:8000,http://localhost:8000",
    ).split(",")
    if o.strip()
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

PATH_PREFIX = os.getenv("PATH_PREFIX", "").rstrip("/")
if PATH_PREFIX:

    @app.middleware("http")
    async def _strip_path_prefix(request, call_next):
        path = request.scope.get("path", "")
        if path == PATH_PREFIX or path.startswith(PATH_PREFIX + "/"):
            new_path = path[len(PATH_PREFIX) :] or "/"
            request.scope["path"] = new_path
            raw = request.scope.get("raw_path")
            if isinstance(raw, (bytes, bytearray)):
                prefix_b = PATH_PREFIX.encode("utf-8")
                if raw == prefix_b or raw.startswith(prefix_b + b"/"):
                    request.scope["raw_path"] = raw[len(prefix_b) :] or b"/"
        return await call_next(request)

# Model configuration — match gwaschen production; override via .env
GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-3.5-flash-lite")
GEMINI_FALLBACKS = [
    m.strip()
    for m in os.getenv(
        "GEMINI_FALLBACKS",
        "gemini-2.5-flash-lite,gemini-3.5-flash,gemini-3.6-flash,gemini-2.5-flash",
    ).split(",")
    if m.strip()
]
GEMINI_RETRIES = int(os.getenv("GEMINI_RETRIES", "2"))
GEMINI_RETRY_BASE_SEC = float(os.getenv("GEMINI_RETRY_BASE_SEC", "1.5"))
GEMINI_THINKING_BUDGET = int(os.getenv("GEMINI_THINKING_BUDGET", "0"))
MAX_FEWSHOT = int(os.getenv("MAX_FEWSHOT", "12"))
MAX_LEXICON_HITS = int(os.getenv("MAX_LEXICON_HITS", "30"))
CORPUS_DIRECT_SCORE = float(os.getenv("CORPUS_DIRECT_SCORE", "0.92"))
GRAMMAR_CHARS = int(os.getenv("GRAMMAR_CHARS", "0"))
FEWSHOT_CANDIDATE_CAP = int(os.getenv("FEWSHOT_CANDIDATE_CAP", "400"))

# Lazy Gemini client initialization to prevent top-level import crashes
_gemini_client: genai.Client | None = None
# After free-tier 429, skip Gemini until this monotonic deadline (unix time).
_gemini_quota_until: float = 0.0


def get_gemini_client() -> genai.Client | None:
    global _gemini_client
    if _gemini_client is not None:
        return _gemini_client

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        print("[WARNING] GEMINI_API_KEY is not set. Gemini fallbacks will fail if invoked.")
        return None

    try:
        _gemini_client = genai.Client(api_key=api_key)
        return _gemini_client
    except Exception as err:
        print(f"[ERROR] Failed to initialize Gemini client: {err}")
        return None


def _mark_gemini_quota(err: Exception | None = None) -> None:
    global _gemini_quota_until
    wait = _quota_retry_seconds(err) if err else 60.0
    _gemini_quota_until = max(_gemini_quota_until, time.time() + wait)
    print(f"[gemini] quota cooldown until +{wait:.0f}s")


def gemini_allowed() -> bool:
    return time.time() >= _gemini_quota_until


BAL_FUNCTION = {
    "ant": "are",
    "ent": "is",
    "at": "was",
    "o": "and",
    "ke": "that",
    "pa": "for",
    "cha": "from",
    "man": "I",
    "taw": "you",
    "má": "we",
    "shomá": "you (pl)",
    "é": "this",
    "á": "that/he/she",
}

CORE_ORTHOGRAPHY = """
You MUST follow Balóchiay Látini Syáhag (Balochi_grammar_rules.pdf) for ALL Balochi output.

Alphabet (only these): A Á B Ch D Dh E É G H I J K L M N O Ó P R S Sh T Th U W Y Z Zh
- Digraphs (jotkáb): Ch, Dh, Sh, Th, Zh
- Never invent letters f, q, v, x (map f→p, q→k, v→w, x→ks)
- Digits: Western 0-9 only (never Arabic-Indic / Persian digit shapes)
- Never output Arabic/Perso-Arabic script for Balochi

Vowels (kassháb): short a e o · long á é i ó u
- Long-vowel shortening before suffixes: Ráh+Band→Rahband, Sháh+Zádag→Shahzádag

Grammar (mandatory for Balochi Latin):
- Word order often SOV; copula ent / ant / at…
- Verb endings: man→án, má→én, taw→ay, shomá→et, á sg→t/it/ent, á pl→ant
- Nonpast = present/future (Man rawán = I go / I will go)
- Imperative: be-/bo- affirmative, ma- prohibitive
- Possessive clitics: -on/-en/-et/-o/-y/-esh; adjectives often -én (sharrén)
- Gemination (dhakk): poll/porr; digraph gemination drops first h (Bachch→Bacch)

Output rules:
- Reply with the translation ONLY (no quotes, no commentary)
- Preserve structure: one source line → one output line; blank lines stay blank
- SUPPORT ORDER: (1) dataset.xlsx sentence/word search
  (2) words_balochi_to_english.xlsx lexicon glosses
  (3) Balochi_grammar_rules.pdf for understanding
  (4) Groq English, then Google Translate / Groq for the target language
- Match dataset.xlsx Latin Syáhag style (á/é/ó, digraphs, e.g. Thám)
"""

GRAMMAR_MUST_FOLLOW = """
MUST FOLLOW Balochi_grammar_rules.pdf for Balochi Latin output:
1) Latin Syáhag alphabet only (A Á B Ch D Dh E É G H I J K L M N O Ó P R S Sh T Th U W Y Z Zh)
2) No Arabic/Perso-Arabic letters; no f/q/v/x; Western digits 0-9
3) Prefer SOV; use correct endings án/én/ay/et/t|it|ent/ant and copula ent/ant/at
4) Imperatives with be-/bo- or prohibitive ma-
5) Prefer wording from dataset.xlsx examples when they match meaning
If a draft violates these rules, rewrite it so it obeys the PDF.
"""

ARABIC_SCRIPT_TARGETS = {
    "urdu",
    "arabic",
    "persian",
    "farsi",
    "persian (farsi)",
    "pashto",
    "sindhi",
    "dari",
    "kurdish",
    "sorani",
}


def norm(text: str) -> str:
    text = unicodedata.normalize("NFC", str(text)).strip()
    return re.sub(r"\s+", " ", text)


def norm_preserve_structure(text: str) -> str:
    text = unicodedata.normalize("NFC", str(text)).replace("\r\n", "\n").replace("\r", "\n")
    lines = [re.sub(r"[ \t]+", " ", ln).strip() for ln in text.split("\n")]
    while lines and not lines[0]:
        lines.pop(0)
    while lines and not lines[-1]:
        lines.pop()
    return "\n".join(lines)


def structure_hint(text: str, text_format: str | None = None) -> str:
    fmt = (text_format or "").strip().casefold()
    lines = text.split("\n")
    n = len(lines)
    blank = sum(1 for ln in lines if not ln.strip())

    if fmt == "poetry":
        return (
            "FORMAT: POETRY / VERSE. "
            f"Source has {n} lines ({blank} blank stanza breaks). "
            f"Your translation MUST have exactly {n} lines with the same blank-line pattern. "
            "Keep poem layout: one source line → one output line; blank lines stay blank between stanzas. "
            "Do NOT rewrite as a prose paragraph. Do NOT merge couplets into one sentence. "
            "Do NOT translate only the first line — translate EVERY non-blank line."
        )
    if fmt == "story":
        return (
            "FORMAT: STORY / PROSE. "
            "Translate as natural narrative paragraphs. "
            "Use a blank line between paragraphs where the source has paragraph breaks. "
            "Do NOT force a line-per-line verse layout."
        )
    if n <= 1:
        return "The source is a single line or paragraph — keep it as one unit."
    return (
        f"The source has {n} lines ({blank} blank). "
        f"Your translation MUST have exactly {n} lines with the same blank-line pattern. "
        "Translate line-by-line; do not merge lines."
    )


def wants_arabic_script(route: dict) -> bool:
    if route["key"] != "bal-any":
        return False
    label = f"{route.get('target', '')} {route.get('target_label', '')}".casefold()
    return any(name in label for name in ARABIC_SCRIPT_TARGETS)


def arabic_script_ratio(text: str) -> float:
    letters = [c for c in text if c.isalpha()]
    if not letters:
        return 1.0
    arabic = sum(1 for c in letters if "\u0600" <= c <= "\u06FF" or "\u0750" <= c <= "\u077F")
    return arabic / len(letters)


def align_line_structure(source: str, translated: str) -> str:
    src_lines = source.split("\n")
    out = translated.replace("\r\n", "\n").replace("\r", "\n").strip()
    out_lines = out.split("\n")
    if len(out_lines) == len(src_lines):
        aligned = "\n".join(ln.strip() for ln in out_lines)
        return preserve_untranslated(source, aligned)

    non_empty_src = [ln for ln in src_lines if ln.strip()]
    if len(src_lines) == len(non_empty_src) and len(out_lines) == 1 and len(non_empty_src) > 1:
        parts = re.split(r"(?<=[.!?۔؟؛])\s+", out)
        parts = [p.strip() for p in parts if p.strip()]
        if len(parts) == len(non_empty_src):
            return preserve_untranslated(source, "\n".join(parts))
    return preserve_untranslated(source, out)


def preserve_untranslated(source: str, translated: str | None) -> str:
    """Keep original Balochi for lines that failed or were left empty.

    Aligns non-empty lines; blank lines stay blank. If the model returned fewer
    lines, or a line looks like a refusal / empty, copy the source line as-is.
    """
    src = (source or "").replace("\r\n", "\n").replace("\r", "\n")
    dst = (translated or "").replace("\r\n", "\n").replace("\r", "\n")
    if not src.strip():
        return (dst or "").strip()
    if not dst.strip():
        return src

    src_lines = src.split("\n")
    dst_lines = dst.split("\n")

    # Fast path: same shape — only replace empty / refusal lines.
    if len(dst_lines) == len(src_lines):
        fixed: list[str] = []
        for s, d in zip(src_lines, dst_lines):
            if not s.strip():
                fixed.append("")
                continue
            d_st = d.strip()
            if not d_st or _looks_like_mt_refusal(d_st):
                fixed.append(s)
            else:
                fixed.append(d_st)
        return "\n".join(fixed)

    # Map non-empty → non-empty; rebuild with original blanks.
    src_ne = [ln for ln in src_lines if ln.strip()]
    dst_ne = [ln.strip() for ln in dst_lines if ln.strip()]
    mapped: list[str] = []
    for i, s in enumerate(src_ne):
        if i < len(dst_ne):
            d = dst_ne[i]
            if not d or _looks_like_mt_refusal(d):
                mapped.append(s)
            else:
                mapped.append(d)
        else:
            mapped.append(s)  # no translation for this line → keep original

    out: list[str] = []
    mi = 0
    for ln in src_lines:
        if not ln.strip():
            out.append("")
        else:
            out.append(mapped[mi] if mi < len(mapped) else ln)
            mi += 1
    return "\n".join(out)


def non_empty_line_count(text: str) -> int:
    return sum(1 for ln in text.replace("\r\n", "\n").replace("\r", "\n").split("\n") if ln.strip())


def should_run_translategemma(text: str, text_format: str | None = None) -> bool:
    """TranslateGemma adapter is disabled — Groq + dataset/lexicon/grammar only."""
    return False



def story_paragraphs(text: str) -> list[str]:
    parts = re.split(r"\n\s*\n+", (text or "").strip())
    return [p.strip() for p in parts if p.strip()]


def should_chunk_genai(text: str, text_format: str | None) -> bool:
    if (text_format or "") == "poetry":
        return False
    paras = story_paragraphs(text)
    return len(text) >= GENAI_CHUNK_CHARS or len(paras) >= GENAI_CHUNK_PARAS


def draft_covers_source(source: str, draft: str) -> bool:
    """Reject short one-liners when the source is a multi-line poem/story."""
    draft = (draft or "").strip()
    if not draft:
        return False
    src_n = non_empty_line_count(source)
    if src_n <= 1:
        return True
    dst_n = non_empty_line_count(draft)
    src_words = max(1, len(source.split()))
    dst_words = len(draft.split())
    # Need most lines, or enough words that it is clearly a full rendering.
    min_lines = max(2, int(src_n * 0.6 + 0.999))
    if dst_n >= min_lines:
        return True
    if dst_words >= max(12, int(src_words * 0.45)):
        return True
    return False


def norm_key(text: str) -> str:
    return norm(text).casefold()


def tokenize(text: str) -> set[str]:
    return {
        t
        for t in re.findall(
            r"[A-Za-zÁÉÓŪĪáéóūīōŌZhzhChchDhShThdhshth]+",
            text.casefold(),
        )
        if len(t) > 1
    }


def fold_accents(text: str) -> str:
    """á/é/ó/ū/ī/ō → a/e/o/u/i/o for false-friend key matching."""
    table = str.maketrans(
        {
            "á": "a",
            "à": "a",
            "ä": "a",
            "é": "e",
            "è": "e",
            "ë": "e",
            "í": "i",
            "ì": "i",
            "ī": "i",
            "ó": "o",
            "ò": "o",
            "ö": "o",
            "ō": "o",
            "ú": "u",
            "ù": "u",
            "ü": "u",
            "ū": "u",
        }
    )
    return (text or "").casefold().translate(table)


# LoRA false friends: folded Balochi lemma → English strings that mean the WRONG sense.
LEXICON_FALSE_FRIENDS: dict[str, list[str]] = {
    "hajjam": [
        "pilgrim",
        "pilgrims",
        "kaaba",
        "ka'ba",
        "ka'bah",
        "mecca",
        "makkah",
        "hajj",
        "pilgrimage",
    ],
}


def lexicon_draft_conflicts(source: str, draft: str, route_key: str) -> list[str]:
    """If the LoRA draft contradicts clear lexicon senses, list conflicts."""
    if route_key != "bal-en" or not draft:
        return []
    draft_l = f" {draft.casefold()} "
    conflicts: list[str] = []
    hits = lexicon_hits(source, "bal-en", limit=50)
    hit_by_fold = {fold_accents(tok): (tok, meaning) for tok, meaning in hits}

    # Also catch source tokens that map via folded false-friend keys
    # (lexicon stores Hajjám → key "hajjám"; false friends use "hajjam").
    for tok in tokenize(source):
        folded = fold_accents(tok)
        wrongs = LEXICON_FALSE_FRIENDS.get(folded, [])
        if not wrongs:
            continue
        meaning = hit_by_fold.get(folded, (tok, "Barber"))[1]
        if folded in hit_by_fold:
            tok, meaning = hit_by_fold[folded]
        elif fold_accents("hajjám") == folded:
            meaning = RESOURCES["bal_lex"].get("hajjám") or "Barber"
            tok = "hajjám"
        for w in wrongs:
            needle = w.casefold().strip()
            if not needle:
                continue
            # Word-ish match so "hajj" does not require trailing space.
            if re.search(rf"(?<![a-z]){re.escape(needle)}(?![a-z])", draft_l):
                conflicts.append(
                    f"{tok} means '{meaning}' (not '{needle}')"
                )
                break
    return conflicts


def draft_misses_lexicon_senses(source: str, draft: str, route_key: str) -> bool:
    """True when the draft ignores clear lexicon glosses from the source (hallucination)."""
    if route_key != "bal-en" or not draft:
        return False
    hits = lexicon_hits(source, "bal-en", limit=20)
    if len(hits) < 2:
        return False
    draft_l = draft.casefold()
    matched = 0
    for _tok, meaning in hits:
        head = re.split(r"[.;/]", meaning.casefold())[0].strip()
        words = [w for w in re.findall(r"[a-z]{4,}", head)]
        if any(w in draft_l for w in words):
            matched += 1
    return matched == 0


def load_grammar_text() -> str:
    if GRAMMAR_PDF.exists():
        reader = PdfReader(str(GRAMMAR_PDF))
        parts: list[str] = []
        for page in reader.pages:
            try:
                parts.append(page.extract_text() or "")
            except Exception:
                continue
        text = re.sub(r"\n{3,}", "\n\n", "\n".join(parts)).strip()
        if text:
            GRAMMAR_CACHE.write_text(text, encoding="utf-8")
            return text

    if GRAMMAR_CACHE.exists():
        return GRAMMAR_CACHE.read_text(encoding="utf-8", errors="ignore")
    return ""


def extract_grammar_rules(full: str, max_chars: int = 0) -> str:
    if not full:
        return "(grammar PDF unavailable)"

    start = 0
    m = re.search(r"A\s*Á\s*B\s*Ch", full)
    if m:
        start = m.start()
    else:
        for needle in ("Balóchiay Áb", "Kassháb", "Kasshab", "Jotkáb"):
            i = full.find(needle)
            if i >= 0:
                start = i
                break

    body = full[start:].strip()
    if not body:
        body = full.strip()
    if max_chars and max_chars > 0:
        body = body[:max_chars]
    return body


def violates_syahag(text: str) -> list[str]:
    issues: list[str] = []
    if re.search(r"[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDFF\uFE70-\uFEFF]", text):
        issues.append("contains_arabic_script")
    if re.search(r"[fqvxFQVX]", text):
        issues.append("forbidden_letters_fqvx")
    if re.search(r"[\u0660-\u0669\u06F0-\u06F9]", text):
        issues.append("non_western_digits")
    return issues


def _first_existing(*paths: Path) -> Path | None:
    for pth in paths:
        if pth.exists():
            return pth
    return None


def _read_table(*candidates: Path):
    path = _first_existing(*candidates)
    if path is None:
        names = ", ".join(pth.name for pth in candidates)
        raise FileNotFoundError(f"Missing table among: {names}")
    if path.suffix.lower() == ".csv":
        return pd.read_csv(path), path
    return pd.read_excel(path), path


def load_resources():
    dataset_df, dataset_path = _read_table(DATASET_CSV, DATASET_CSV_ROOT, DATASET_XLSX)

    pairs: list[dict] = []
    en_to_bal: dict[str, str] = {}
    bal_to_en: dict[str, str] = {}
    for en, bal in zip(dataset_df["English_sentences"], dataset_df["Balochi_sentences"]):
        if pd.isna(en) or pd.isna(bal):
            continue
        en_s, bal_s = norm(en), norm(bal)
        if not en_s or not bal_s:
            continue
        pairs.append({"en": en_s, "bal": bal_s, "en_toks": tokenize(en_s), "bal_toks": tokenize(bal_s)})
        en_to_bal.setdefault(norm_key(en_s), bal_s)
        bal_to_en.setdefault(norm_key(bal_s), en_s)

    # Lexicon (helps gwachen quality on CPU when TranslateGemma is unavailable).
    bal_lex: dict[str, str] = {}
    en_lex: dict[str, str] = {}
    words_path = None
    try:
        words_df, words_path = _read_table(WORDS_CSV, WORDS_CSV_ROOT, WORDS_XLSX)
        for bal, meaning in zip(words_df["balochi_word"], words_df["english_meaning"]):
            if pd.isna(bal) or pd.isna(meaning):
                continue
            bal_s, en_s = norm(bal), norm(meaning)
            if not bal_s or not en_s:
                continue
            bal_lex.setdefault(norm_key(bal_s), en_s)
            en_lex.setdefault(norm_key(en_s), bal_s)
    except FileNotFoundError:
        words_path = None

    grammar = load_grammar_text()
    grammar_excerpt = extract_grammar_rules(grammar, GRAMMAR_CHARS)

    bal_index: dict[str, list[int]] = defaultdict(list)
    en_index: dict[str, list[int]] = defaultdict(list)
    for i, pair in enumerate(pairs):
        for tok in pair["bal_toks"]:
            bal_index[tok].append(i)
        for tok in pair["en_toks"]:
            en_index[tok].append(i)

    print(
        f"[startup] dataset={dataset_path}  words={words_path}  "
        f"pairs={len(pairs):,}  lexicon={len(bal_lex):,}  "
        f"grammar_chars={len(grammar_excerpt):,}  model={GEMINI_MODEL}"
    )
    return {
        "pairs": pairs,
        "en_to_bal": en_to_bal,
        "bal_to_en": bal_to_en,
        "bal_lex": bal_lex,
        "en_lex": en_lex,
        "grammar_excerpt": grammar_excerpt,
        "bal_index": bal_index,
        "en_index": en_index,
        "dataset_path": str(dataset_path),
        "words_path": str(words_path) if words_path else None,
    }


RESOURCES = load_resources()

BAL_ALIASES = {
    "balochi",
    "baluchi",
    "bal",
    "bal_latn",
    "balochi-latin",
    "latin-balochi",
    "balochi latin",
}
EN_ALIASES = {"english", "en", "eng"}

TARGET_LANGUAGES = [
    ("english", "English"),
    ("urdu", "Urdu"),
    ("persian", "Persian (Farsi)"),
    ("arabic", "Arabic"),
    ("hindi", "Hindi"),
    ("pashto", "Pashto"),
    ("sindhi", "Sindhi"),
    ("brahui", "Brahui"),
    ("turkish", "Turkish"),
    ("spanish", "Spanish"),
    ("french", "French"),
    ("german", "German"),
    ("portuguese", "Portuguese"),
    ("russian", "Russian"),
    ("chinese", "Chinese"),
    ("japanese", "Japanese"),
    ("korean", "Korean"),
    ("italian", "Italian"),
    ("dutch", "Dutch"),
    ("swedish", "Swedish"),
    ("indonesian", "Indonesian"),
    ("bengali", "Bengali"),
    ("punjabi", "Punjabi"),
    ("gujarati", "Gujarati"),
    ("marathi", "Marathi"),
    ("tamil", "Tamil"),
    ("telugu", "Telugu"),
    ("malayalam", "Malayalam"),
    ("thai", "Thai"),
    ("vietnamese", "Vietnamese"),
    ("polish", "Polish"),
    ("ukrainian", "Ukrainian"),
    ("greek", "Greek"),
    ("hebrew", "Hebrew"),
    ("swahili", "Swahili"),
]
LANG_BY_NAME = {name: code for code, name in TARGET_LANGUAGES}
LANG_BY_NAME.update({code: code for code, _name in TARGET_LANGUAGES})
LANG_BY_NAME["Persian (Farsi)"] = "persian"
LANG_BY_NAME["Farsi"] = "persian"


def is_balochi(lang: str) -> bool:
    return lang.strip().casefold() in BAL_ALIASES


def is_english(lang: str) -> bool:
    return lang.strip().casefold() in EN_ALIASES


def parse_route(source_lang: str, target_lang: str) -> dict:
    src = source_lang.strip()
    tgt = target_lang.strip()
    if not src or not tgt:
        raise HTTPException(status_code=400, detail="source_lang and target_lang are required")

    if is_balochi(tgt):
        if is_english(src):
            return {
                "key": "en-bal",
                "source": "english",
                "target": "balochi",
                "target_label": "Balochi Latin Syáhag",
                "source_label": "English",
            }
        if is_balochi(src):
            raise HTTPException(
                status_code=400,
                detail="Source and target cannot both be Balochi. Choose another source language.",
            )
        label = src
        for code, name in TARGET_LANGUAGES:
            if src.casefold() in {code, name.casefold()}:
                label = name
                break
        return {
            "key": "any-bal",
            "source": src.casefold(),
            "target": "balochi",
            "target_label": "Balochi Latin Syáhag",
            "source_label": label,
        }

    if is_balochi(src) and is_english(tgt):
        return {
            "key": "bal-en",
            "source": "balochi",
            "target": "english",
            "target_label": "English",
        }
    if is_balochi(src):
        label = tgt
        for code, name in TARGET_LANGUAGES:
            if tgt.casefold() in {code, name.casefold()}:
                label = name
                break
        return {
            "key": "bal-any",
            "source": "balochi",
            "target": tgt.casefold(),
            "target_label": label,
        }

    raise HTTPException(
        status_code=400,
        detail=(
            "Supported: Balochi → any language, or any language → Balochi. "
            f"Got source={source_lang!r} target={target_lang!r}"
        ),
    )


def exact_lookup(text: str, route_key: str) -> str | None:
    key = norm_key(text)
    if route_key == "en-bal":
        return RESOURCES["en_to_bal"].get(key) or RESOURCES["en_lex"].get(key)
    if route_key in {"bal-en", "bal-any"}:
        return RESOURCES["bal_to_en"].get(key) or RESOURCES["bal_lex"].get(key)
    return None


def lexicon_hits(text: str, route_key: str, limit: int = MAX_LEXICON_HITS) -> list[tuple[str, str]]:
    toks = tokenize(text)
    if not toks:
        return []
    lex = RESOURCES["en_lex"] if route_key == "en-bal" else RESOURCES["bal_lex"]
    out: list[tuple[str, str]] = []
    seen: set[str] = set()
    for tok in toks:
        k = norm_key(tok)
        if k in seen:
            continue
        hit = lex.get(k)
        if hit:
            seen.add(k)
            out.append((tok, hit))
            if len(out) >= limit:
                break
    return out


def retrieve_fewshot(text: str, route_key: str, k: int = MAX_FEWSHOT) -> list[tuple[float, dict]]:
    q_toks = tokenize(text)
    if not q_toks:
        return []

    index = RESOURCES["en_index"] if route_key == "en-bal" else RESOURCES["bal_index"]
    pairs = RESOURCES["pairs"]

    hit_counts: dict[int, int] = defaultdict(int)
    for tok in sorted(q_toks, key=lambda t: len(index.get(t, []))):
        postings = index.get(tok)
        if not postings:
            continue
        for idx in postings:
            hit_counts[idx] += 1
        if len(hit_counts) >= FEWSHOT_CANDIDATE_CAP * 3:
            break

    if not hit_counts:
        return []

    cand_ids = [
        idx
        for idx, _ in nlargest(
            FEWSHOT_CANDIDATE_CAP,
            hit_counts.items(),
            key=lambda item: item[1],
        )
    ]

    scored: list[tuple[float, dict]] = []
    for idx in cand_ids:
        pair = pairs[idx]
        ref = pair["en_toks"] if route_key == "en-bal" else pair["bal_toks"]
        if not ref:
            continue
        inter = len(q_toks & ref)
        if inter == 0:
            continue
        union = len(q_toks | ref)
        score = inter / union
        src = pair["en"] if route_key == "en-bal" else pair["bal"]
        length_penalty = abs(len(src) - len(text)) / max(len(text), 1)
        score = score - 0.05 * min(length_penalty, 2.0)
        if score >= 0.08:
            scored.append((score, pair))

    scored.sort(key=lambda x: x[0], reverse=True)
    out: list[tuple[float, dict]] = []
    seen = set()
    for score, pair in scored:
        sig = (pair["en"], pair["bal"])
        if sig in seen:
            continue
        seen.add(sig)
        out.append((score, pair))
        if len(out) >= k:
            break
    return out


def offline_fallback(
    text: str,
    route_key: str,
    fewshot: list[tuple[float, dict]],
    target_label: str,
) -> tuple[str, dict] | None:
    if route_key == "bal-any":
        return None

    # Never replace a multi-line poem/story with a single dataset sentence.
    if non_empty_line_count(text) > 1:
        return None

    to_bal = route_key in {"en-bal", "any-bal"}

    if fewshot:
        score, pair = fewshot[0]
        if score >= 0.22:
            tgt = pair["bal"] if to_bal else pair["en"]
            return tgt, {
                "reason": "closest_dataset_example",
                "score": round(score, 3),
                "example": {"en": pair["en"], "bal": pair["bal"]},
            }

    if fewshot:
        score, pair = fewshot[0]
        tgt = pair["bal"] if to_bal else pair["en"]
        return tgt, {
            "reason": "weak_dataset_example",
            "score": round(score, 3),
            "example": {"en": pair["en"], "bal": pair["bal"]},
        }
    return None


def _is_retryable(err: Exception) -> bool:
    msg = str(err).casefold()
    needles = (
        "high demand",
        "resource_exhausted",
        "unavailable",
        "try again",
        "rate",
        "quota",
        "429",
        "503",
        "overloaded",
        "temporarily",
        "spikes in demand",
        "capacity",
        "server error",
        "internal error",
        "deadline",
        "timeout",
        "timed out",
        "free_tier",
    )
    return any(n in msg for n in needles)


def _is_quota_exhausted(err: Exception) -> bool:
    msg = str(err).casefold()
    return any(
        n in msg
        for n in (
            "exceeded your current quota",
            "free_tier",
            "quota exceeded",
            "resource_exhausted",
            "rate limit",
        )
    )


def _quota_retry_seconds(err: Exception) -> float:
    """Parse 'Please retry in 30.29s' from Gemini quota errors."""
    msg = str(err)
    m = re.search(r"retry in\s+([0-9]+(?:\.[0-9]+)?)\s*s", msg, flags=re.I)
    if m:
        return min(45.0, max(5.0, float(m.group(1)) + 1.0))
    return 20.0


def generate_with_retries(prompt: str, *, light: bool = False) -> tuple[str, str]:
    if not gemini_allowed():
        raise RuntimeError(
            f"Gemini free-tier cooldown active; retry in "
            f"{max(1, int(_gemini_quota_until - time.time()))}s"
        )
    client = get_gemini_client()
    if client is None:
        raise RuntimeError("GEMINI_API_KEY is not configured on the server.")

    models = [GEMINI_MODEL] + [m for m in GEMINI_FALLBACKS if m != GEMINI_MODEL]
    last_err: Exception | None = None
    system = (
        "You are a careful literary translator. "
        "Reply with the translation ONLY — no preface."
        if light
        else build_system_prompt()
    )
    config_kwargs: dict = {"system_instruction": system, "temperature": 0.1}

    if (not light) and GEMINI_THINKING_BUDGET >= 0:
        try:
            config_kwargs["thinking_config"] = types.ThinkingConfig(
                thinking_budget=GEMINI_THINKING_BUDGET
            )
        except Exception:
            pass

    gen_config = types.GenerateContentConfig(**config_kwargs)
    quota_hits = 0

    for model in models:
        for attempt in range(1, GEMINI_RETRIES + 1):
            try:
                response = client.models.generate_content(
                    model=model,
                    contents=prompt,
                    config=gen_config,
                )
                text = (response.text or "").strip()
                if not text:
                    raise RuntimeError("Empty model response")
                return text, model
            except Exception as err:
                last_err = err
                msg = str(err).casefold()
                bad_model = any(
                    n in msg
                    for n in (
                        "not found",
                        "not_found",
                        "invalid",
                        "404",
                        "unsupported",
                        "no longer available",
                        "update your code",
                        "deprecated",
                    )
                )
                # Free-tier quota: do NOT sleep 30–45s here (client already waits).
                # One quick retry only — cascading models burns the 20/min limit.
                if _is_quota_exhausted(err):
                    quota_hits += 1
                    _mark_gemini_quota(err)
                    print(f"[gemini] quota on {model} (hit={quota_hits})")
                    if quota_hits >= 2:
                        raise err
                    time.sleep(1.5)
                    continue
                high_demand = any(
                    n in msg
                    for n in (
                        "high demand",
                        "spikes in demand",
                        "overloaded",
                        "429",
                        "503",
                    )
                )
                if high_demand:
                    time.sleep(min(12.0, GEMINI_RETRY_BASE_SEC * (2 ** (attempt - 1))))
                    if attempt >= 2:
                        break
                    continue
                if attempt < GEMINI_RETRIES and _is_retryable(err):
                    time.sleep(GEMINI_RETRY_BASE_SEC * attempt)
                    continue
                if _is_retryable(err) or bad_model:
                    break
                raise
    assert last_err is not None
    if _is_quota_exhausted(last_err):
        _mark_gemini_quota(last_err)
    raise last_err


def build_pivot_prompt(
    english: str,
    target_label: str,
    *,
    text_format: str | None,
    arabic_script: bool,
) -> str:
    """Small EN→target prompt (avoids huge Balochi grammar system prompt)."""
    parts = [
        f"Translate from English to {target_label}.",
        "Keep meaning faithful. Reply with the translation ONLY.",
        structure_hint(english, text_format),
    ]
    if arabic_script:
        parts.append(
            f"CRITICAL: Write the entire {target_label} translation in Arabic / "
            "Perso-Arabic script only. No Latin transliteration."
        )
    else:
        parts.append(f"Use the standard writing system for {target_label}.")
    parts.append(f"English source:\n{english}\n\n{target_label} translation:")
    return "\n\n".join(parts)


# Free EN→target when Gemini quota is dead (Google Translate + Groq + HF).
GOOGLE_TL = {
    "english": "en",
    "urdu": "ur",
    "arabic": "ar",
    "persian": "fa",
    "farsi": "fa",
    "hindi": "hi",
    "pashto": "ps",
    "sindhi": "sd",
    "turkish": "tr",
    "spanish": "es",
    "french": "fr",
    "german": "de",
    "portuguese": "pt",
    "russian": "ru",
    "chinese": "zh-CN",
    "japanese": "ja",
    "korean": "ko",
    "italian": "it",
    "dutch": "nl",
    "swedish": "sv",
    "indonesian": "id",
    "bengali": "bn",
    "punjabi": "pa",
    "gujarati": "gu",
    "marathi": "mr",
    "tamil": "ta",
    "telugu": "te",
    "malayalam": "ml",
    "thai": "th",
    "vietnamese": "vi",
    "polish": "pl",
    "ukrainian": "uk",
    "greek": "el",
    "hebrew": "he",
    "swahili": "sw",
}

NLLB_TARGET = {
    "urdu": "urd_Arab",
    "arabic": "arb_Arab",
    "persian": "pes_Arab",
    "farsi": "pes_Arab",
    "hindi": "hin_Deva",
    "pashto": "pbt_Arab",
    "sindhi": "snd_Arab",
    "turkish": "tur_Latn",
    "spanish": "spa_Latn",
    "french": "fra_Latn",
    "german": "deu_Latn",
    "portuguese": "por_Latn",
    "russian": "rus_Cyrl",
    "chinese": "zho_Hans",
    "japanese": "jpn_Jpan",
    "korean": "kor_Hang",
    "italian": "ita_Latn",
    "dutch": "nld_Latn",
    "swedish": "swe_Latn",
    "indonesian": "ind_Latn",
    "bengali": "ben_Beng",
    "punjabi": "pan_Guru",
    "gujarati": "guj_Gujr",
    "marathi": "mar_Deva",
    "tamil": "tam_Taml",
    "telugu": "tel_Telu",
    "malayalam": "mal_Mlym",
    "thai": "tha_Thai",
    "vietnamese": "vie_Latn",
    "polish": "pol_Latn",
    "ukrainian": "ukr_Cyrl",
    "greek": "ell_Grek",
    "hebrew": "heb_Hebr",
    "swahili": "swh_Latn",
}
OPUS_EN_TARGET = {
    "urdu": "Helsinki-NLP/opus-mt-en-ur",
    "arabic": "Helsinki-NLP/opus-mt-en-ar",
    "hindi": "Helsinki-NLP/opus-mt-en-hi",
    "french": "Helsinki-NLP/opus-mt-en-fr",
    "german": "Helsinki-NLP/opus-mt-en-de",
    "spanish": "Helsinki-NLP/opus-mt-en-es",
    "russian": "Helsinki-NLP/opus-mt-en-ru",
    "chinese": "Helsinki-NLP/opus-mt-en-zh",
}


def _hf_token() -> str:
    return (
        os.getenv("HF_TOKEN")
        or os.getenv("HUGGING_FACE_HUB_TOKEN")
        or os.getenv("HUGGINGFACE_HUB_TOKEN")
        or ""
    ).strip()


def _hf_en_to_target(english: str, target_code: str, target_label: str) -> tuple[str, str] | None:
    """EN→target via Hugging Face Inference (opus-mt / NLLB)."""
    english = (english or "").strip()
    if not english:
        return None
    code = (target_code or "").casefold().strip()
    token = _hf_token()
    if not token:
        return None
    try:
        from huggingface_hub import InferenceClient

        client = InferenceClient(token=token)
        opus = OPUS_EN_TARGET.get(code)
        if opus:
            try:
                out = client.translation(english, model=opus)
                if isinstance(out, list):
                    out = out[0] if out else ""
                if isinstance(out, dict):
                    out = out.get("translation_text") or out.get("generated_text") or ""
                text = str(out).strip()
                if text:
                    return text, opus
            except Exception as err:
                print(f"[free_mt] opus {opus} failed: {err}")

        nllb_tgt = NLLB_TARGET.get(code)
        if nllb_tgt:
            try:
                out = client.translation(
                    english,
                    model="facebook/nllb-200-distilled-600M",
                    src_lang="eng_Latn",
                    tgt_lang=nllb_tgt,
                )
                if isinstance(out, list):
                    out = out[0] if out else ""
                if isinstance(out, dict):
                    out = out.get("translation_text") or out.get("generated_text") or ""
                text = str(out).strip()
                if text:
                    return text, "facebook/nllb-200-distilled-600M"
            except Exception as err:
                print(f"[free_mt] nllb failed: {err}")
    except Exception as err:
        print(f"[free_mt] hf client failed: {err}")
    return None


def groq_configured() -> bool:
    return bool((os.getenv("GROQ_API_KEY") or "").strip())


def groq_model_id() -> str:
    return (os.getenv("GROQ_MODEL") or DEFAULT_GROQ_MODEL).strip() or DEFAULT_GROQ_MODEL


def groq_support_stack() -> list[str]:
    return list(GROQ_SUPPORT)


def mt_support_stack() -> list[str]:
    """Public support labels: adapter/dataset/grammar/lexicon first, then Groq/Google."""
    return list(PRIORITIES) + ([f"groq:{groq_model_id()}", "google_translate"] if groq_configured() else [])


def build_mt_grounding(text: str, route_key: str) -> str:
    """Binding senses from dataset + lexicon + grammar (no adapter)."""
    rk = "en-bal" if route_key in {"en-bal", "any-bal"} else "bal-en"
    parts: list[str] = [
        "UNDERSTANDING PIPELINE (binding order):",
        "1) Search dataset.xlsx for matching Balochi sentences → English parallels",
        "2) Look up words in words_balochi_to_english.xlsx",
        "3) Apply Balochi_grammar_rules.pdf for morphology/syntax sense",
        "4) Produce faithful translation from the Balochi source using (1)–(3)",
        "KEEP-ORIGINAL RULE (mandatory): If a word, phrase, or whole line has no "
        "clear sense from dataset/lexicon/grammar and you cannot translate it "
        "confidently, copy the original Balochi Latin text unchanged into that "
        "place. Do NOT invent, skip, omit, or leave a blank for that part.",
        "Do NOT invent unrelated themes. Do NOT use TranslateGemma adapter drafts.",
    ]
    # Exact / near sentence hits first
    few = retrieve_fewshot(text, rk, k=8)
    if few:
        parts.append("BINDING dataset.xlsx sentence matches:")
        for score, pair in few:
            parts.append(
                f"[score={score:.2f}]\nBalochi: {pair['bal']}\nEnglish: {pair['en']}"
            )
    # Also search each non-empty line as its own query
    line_hits = 0
    for ln in (text or "").split("\n"):
        raw = ln.strip()
        if not raw or line_hits >= 6:
            continue
        exact = exact_lookup(raw, "bal-en" if rk == "bal-en" else "en-bal")
        if exact:
            parts.append(f"EXACT dataset/lexicon line:\nBalochi: {raw}\nEnglish: {exact}")
            line_hits += 1
            continue
        sub = retrieve_fewshot(raw, rk, k=1)
        if sub and sub[0][0] >= 0.2:
            score, pair = sub[0]
            parts.append(
                f"LINE dataset match [score={score:.2f}]:\n"
                f"Balochi: {pair['bal']}\nEnglish: {pair['en']}"
            )
            line_hits += 1
    lex = lexicon_hits(text, rk, limit=40)
    if lex:
        parts.append("BINDING lexicon (words_balochi_to_english.xlsx):")
        for bal, en in lex:
            parts.append(f"- {bal} = {en}")
    grammar = (RESOURCES.get("grammar_excerpt") or "").strip()
    if grammar:
        parts.append(
            "BINDING grammar (Balochi_grammar_rules.pdf) — use for understanding "
            "the Balochi Latin source before writing English:\n"
            + grammar[:1800]
        )
    return "\n".join(parts)


def _looks_like_mt_refusal(text: str) -> bool:
    """True if the model refused instead of translating."""
    t = (text or "").casefold()
    if len(t) > 280:
        return False
    needles = (
        "cannot translate",
        "can't translate",
        "do not know how to translate",
        "don't know how to translate",
        "i'm unable to translate",
        "i am unable to translate",
        "unable to translate",
        "as an ai",
        "ترجمہ کیسے",
        "ترجمہ نہیں",
        "نہیں جانتا",
    )
    return any(n in t for n in needles)


def _looks_like_hallucinated_mt(source: str, translated: str) -> bool:
    """Catch generic invented poems / repeated lines that ignore the source."""
    src_n = non_empty_line_count(source)
    dst_n = non_empty_line_count(translated)
    if src_n >= 2 and dst_n < max(1, int(src_n * 0.7 + 0.001)):
        return True
    lines = [ln.strip() for ln in (translated or "").split("\n") if ln.strip()]
    if len(lines) >= 3:
        top_n = Counter(lines).most_common(1)[0][1]
        if top_n >= 2 and top_n >= len(lines) // 2:
            return True
    # Generic English filler often produced when the model ignores Balochi.
    en = (translated or "").casefold()
    generic = (
        "purpose of life",
        "wandering soul",
        "according to this account",
        "restless and tormented",
        "night spreads its darkness",
    )
    if src_n >= 2 and any(g in en for g in generic):
        # Allow only if source somehow matches — these almost always mean hallucination.
        return True
    return False


def _english_misses_lexicon(
    source: str, english: str, *, strict: bool = False
) -> bool:
    """True when binding lexicon glosses are absent from the English draft."""
    hits = lexicon_hits(source, "bal-en", limit=24)
    if len(hits) < 3:
        return False
    en_fold = fold_accents(english or "").casefold()
    matched = 0
    stop = {
        "the", "and", "for", "with", "from", "that", "this", "into", "onto",
        "over", "under", "have", "been", "were", "will", "shall", "also",
    }
    for _bal, gloss in hits:
        words = [
            w
            for w in re.findall(r"[A-Za-z]{4,}", gloss or "")
            if w.casefold() not in stop
        ]
        if any(w.casefold() in en_fold for w in words):
            matched += 1
    if strict:
        need = max(2, int(len(hits) * 0.35 + 0.999))
        return matched < need
    # Soft: only reject total misses (used for Groq output filter).
    return matched == 0


def _number_source_lines(text: str) -> str:
    """Number non-empty lines so the model must return the same count."""
    out: list[str] = []
    n = 0
    for ln in (text or "").split("\n"):
        if not ln.strip():
            out.append("")
            continue
        n += 1
        out.append(f"{n}. {ln.strip()}")
    return "\n".join(out)


def _strip_line_numbers(text: str) -> str:
    return re.sub(r"(?m)^\s*\d+\s*[.)\-:]\s*", "", text or "")


def _groq_extract_content(data: dict) -> tuple[str, str, dict]:
    """Return (text, finish_reason, usage) from a Groq chat completion."""
    choice = (data.get("choices") or [{}])[0]
    msg = choice.get("message") or {}
    text = (msg.get("content") or "").strip().strip('"').strip("'")
    finish = str(choice.get("finish_reason") or "")
    usage = data.get("usage") or {}
    return text, finish, usage if isinstance(usage, dict) else {}


def groq_chat_translate(
    source_text: str,
    *,
    source_label: str,
    target_code: str,
    target_label: str,
    grounding: str | None = None,
    route_key: str | None = None,
    reject_filters: bool = True,
) -> tuple[str, str] | None:
    """Direct Groq chat translation (Balochi or English → target), optionally grounded."""
    groq_key = (os.getenv("GROQ_API_KEY") or "").strip()
    source_text = (source_text or "").strip()
    if not groq_key or not source_text:
        return None
    code = (target_code or "").casefold().strip()
    label = (target_label or code).strip()
    arabic = code in {
        "urdu",
        "arabic",
        "persian",
        "farsi",
        "pashto",
        "sindhi",
        "hebrew",
    }
    src_lines = non_empty_line_count(source_text)
    rk = route_key or (
        "bal-en"
        if "balochi" in source_label.casefold()
        else ("en-bal" if code.startswith("bal") else "bal-en")
    )
    ground = (grounding or "").strip() or build_mt_grounding(source_text, rk)
    try:
        import httpx

        model = groq_model_id()
        max_out = int(os.getenv("GROQ_MAX_COMPLETION_TOKENS", "8192"))
        script = (
            f"Write {label} in Arabic/Perso-Arabic script only (not Latin/Roman). "
            if arabic
            else f"Use the standard writing system for {label}. "
        )
        numbered = _number_source_lines(source_text)
        prompt = (
            f"{ground}\n\n"
            f"Faithfully translate from {source_label} to {label}.\n"
            f"{script}\n"
            f"There are {src_lines} non-empty source lines. "
            f"Return exactly {src_lines} non-empty output lines "
            "(same blank-line stanza breaks). "
            "Use BINDING lexicon/dataset senses above. "
            "Translate each numbered line's meaning — do NOT invent a new poem, "
            "do NOT summarize, do NOT repeat lines, do NOT transliterate "
            "letter-by-letter. "
            "KEEP-ORIGINAL: if a word/phrase/line cannot be translated from the "
            "bindings, leave that Balochi Latin text as-is in the output. "
            "Never refuse. Reply with the translation ONLY "
            "(you may keep or drop the numbers).\n\n"
            f"{source_label}:\n{numbered}\n\n{label}:"
        )
        body: dict = {
            "model": model,
            "temperature": 0.05,
            "max_completion_tokens": max_out,
            "messages": [
                {
                    "role": "system",
                    "content": (
                        "You are a careful literary translator for Balochi Latin "
                        "(Syáhag). Obey lexicon + dataset.xlsx + grammar rules. "
                        "Translate meaning line-by-line. Never invent themes. "
                        "If you cannot translate a word or line, keep the original "
                        "Balochi Latin text in that place. "
                        "Never refuse. Reply with the translation only."
                    ),
                },
                {"role": "user", "content": prompt},
            ],
        }
        if "gpt-oss" in model.casefold():
            body["reasoning_effort"] = os.getenv("GROQ_REASONING_EFFORT", "low")

        filters_on = reject_filters
        for attempt in range(3):
            r = httpx.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {groq_key}",
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=90.0,
            )
            if r.status_code == 429:
                print(f"[groq] HTTP 429: {r.text[:300]}")
                time.sleep(2.5 * (attempt + 1))
                continue
            if r.status_code >= 300:
                print(f"[groq] HTTP {r.status_code}: {r.text[:300]}")
                return None
            data = r.json()
            text_out, finish, usage = _groq_extract_content(data)
            text_out = _strip_line_numbers(text_out)
            if text_out and _looks_like_mt_refusal(text_out):
                print(f"[groq] refusal-like output: {text_out[:160]!r}")
                text_out = ""
            if filters_on and text_out and _looks_like_hallucinated_mt(source_text, text_out):
                print(
                    f"[groq] hallucinated/short output rejected ({src_lines} src lines)"
                )
                text_out = ""
            if (
                filters_on
                and text_out
                and rk == "bal-en"
                and label.casefold().startswith("english")
                and _english_misses_lexicon(source_text, text_out)
            ):
                print("[groq] rejected: English misses binding lexicon glosses")
                text_out = ""
            if text_out:
                return text_out, f"groq:{model}"
            details = usage.get("completion_tokens_details") or {}
            reasoning_n = details.get("reasoning_tokens")
            print(
                f"[groq] empty/bad content finish={finish!r} "
                f"reasoning_tokens={reasoning_n} usage={usage}"
            )
            if finish == "length" and attempt == 0:
                body["max_completion_tokens"] = min(max_out * 2, 32768)
                continue
            # Drop quality filters so labzank never hard-fails empty.
            if filters_on:
                filters_on = False
                continue
            return None
    except Exception as err:
        print(f"[groq] failed: {err}")
    return None


def groq_translate_chunked(
    source_text: str,
    *,
    source_label: str,
    target_code: str,
    target_label: str,
    text_format: str | None = None,
    route_key: str | None = None,
    adapter_draft: str | None = None,
    strict_quality: bool = True,
) -> tuple[str, str] | None:
    """Translate with dataset/lexicon/grammar grounding; always try a best-effort fallback."""
    source_text = (source_text or "").strip()
    if not source_text:
        return None
    rk = route_key or (
        "bal-en"
        if "balochi" in source_label.casefold()
        else ("en-bal" if (target_code or "").casefold().startswith("bal") else "bal-en")
    )
    ground = build_mt_grounding(source_text, rk)
    # adapter_draft ignored — TranslateGemma is disabled for this pipeline.

    hit = groq_chat_translate(
        source_text,
        source_label=source_label,
        target_code=target_code,
        target_label=target_label,
        grounding=ground,
        route_key=rk,
        reject_filters=strict_quality,
    )
    if hit:
        return preserve_untranslated(source_text, hit[0]), hit[1]

    parts = story_paragraphs(source_text)
    src_lines = non_empty_line_count(source_text)
    should_chunk = len(parts) >= 3 or src_lines >= 8 or (text_format or "") == "poetry"
    if not should_chunk or len(parts) < 2:
        fallback = groq_chat_translate(
            source_text,
            source_label=source_label,
            target_code=target_code,
            target_label=target_label,
            grounding=ground,
            route_key=rk,
            reject_filters=False,
        )
        if fallback:
            return preserve_untranslated(source_text, fallback[0]), fallback[1]
        # Nothing translated — keep original Balochi in place.
        return source_text, f"groq:{groq_model_id()}:keep_original"

    group_size = 1 if (text_format or "") == "poetry" or len(parts) >= 4 else 2
    out_parts: list[str] = []
    model_id = ""
    failed = 0
    n_groups = 0
    for i in range(0, len(parts), group_size):
        n_groups += 1
        chunk = "\n\n".join(parts[i : i + group_size])
        chunk_ground = build_mt_grounding(chunk, rk)
        piece = groq_chat_translate(
            chunk,
            source_label=source_label,
            target_code=target_code,
            target_label=target_label,
            grounding=chunk_ground,
            route_key=rk,
            reject_filters=strict_quality,
        )
        if not piece:
            piece = groq_chat_translate(
                chunk,
                source_label=source_label,
                target_code=target_code,
                target_label=target_label,
                grounding=chunk_ground,
                route_key=rk,
                reject_filters=False,
            )
        if not piece:
            failed += 1
            out_parts.append(chunk)  # keep original Balochi for this chunk
            continue
        text_part, model_id = piece
        out_parts.append(preserve_untranslated(chunk, text_part.strip()))
        time.sleep(0.2)
    joined = "\n\n".join(out_parts).strip()
    if not joined:
        return source_text, f"groq:{groq_model_id()}:keep_original"
    if failed >= n_groups:
        # All chunks kept original — still return original rather than inventing.
        return preserve_untranslated(source_text, joined), model_id or (
            f"groq:{groq_model_id()}:keep_original"
        )
    return (
        preserve_untranslated(source_text, joined),
        model_id or f"groq:{groq_model_id()}",
    )

def google_translate_configured() -> bool:
    # Official Cloud key optional; free gtx client always available as fallback.
    return True


def google_tl_code(target_code: str) -> str | None:
    code = (target_code or "").casefold().strip()
    if code in GOOGLE_TL:
        return GOOGLE_TL[code]
    if len(code) in {2, 3, 5, 7} and code.replace("-", "").isalpha():
        return code
    return None


def _google_translate_cloud(
    text: str, *, source_tl: str, target_tl: str, api_key: str
) -> str | None:
    import httpx

    r = httpx.post(
        "https://translation.googleapis.com/language/translate/v2",
        params={"key": api_key},
        json={
            "q": text,
            "source": source_tl,
            "target": target_tl,
            "format": "text",
        },
        timeout=60.0,
    )
    if r.status_code >= 300:
        print(f"[google_translate] cloud HTTP {r.status_code}: {r.text[:240]}")
        return None
    data = r.json() or {}
    translations = ((data.get("data") or {}).get("translations")) or []
    if not translations:
        return None
    out = (translations[0].get("translatedText") or "").strip()
    return out or None


def _google_translate_gtx(
    text: str, *, source_tl: str, target_tl: str
) -> str | None:
    """Unofficial translate.googleapis.com client=gtx (no key). Chunk long text."""
    import httpx

    def split_parts(src: str, limit: int = 1400) -> list[str]:
        lines = src.split("\n")
        parts: list[str] = []
        buf = ""
        for ln in lines:
            candidate = f"{buf}\n{ln}" if buf else ln
            if len(candidate) > limit and buf:
                parts.append(buf)
                buf = ln
            else:
                buf = candidate
        if buf:
            parts.append(buf)
        return parts or [src]

    out_chunks: list[str] = []
    for part in split_parts(text):
        if not part.strip():
            out_chunks.append(part)
            continue
        r = httpx.get(
            "https://translate.googleapis.com/translate_a/single",
            params={
                "client": "gtx",
                "sl": source_tl,
                "tl": target_tl,
                "dt": "t",
                "q": part,
            },
            timeout=60.0,
        )
        if r.status_code >= 300:
            print(f"[google_translate] gtx HTTP {r.status_code}: {r.text[:240]}")
            return None
        data = r.json()
        segs: list[str] = []
        try:
            for row in data[0] or []:
                if row and isinstance(row[0], str):
                    segs.append(row[0])
        except (TypeError, IndexError, KeyError) as err:
            print(f"[google_translate] gtx parse failed: {err}")
            return None
        piece = "".join(segs).strip()
        if not piece:
            return None
        out_chunks.append(piece)
    return "\n".join(out_chunks).strip() or None


def google_translate_text(
    text: str,
    *,
    target_code: str,
    source_code: str = "english",
) -> tuple[str, str] | None:
    """EN (or other) → target via Google Cloud or free gtx endpoint."""
    text = (text or "").strip()
    if not text:
        return None
    target_tl = google_tl_code(target_code)
    source_tl = google_tl_code(source_code) or "auto"
    if not target_tl:
        return None
    if source_tl == target_tl:
        return text, "google_translate:identity"

    api_key = (
        os.getenv("GOOGLE_TRANSLATE_API_KEY")
        or os.getenv("GOOGLE_API_KEY")
        or ""
    ).strip()
    try:
        if api_key:
            out = _google_translate_cloud(
                text, source_tl=source_tl, target_tl=target_tl, api_key=api_key
            )
            if out:
                return out, "google_translate:cloud"
        out = _google_translate_gtx(text, source_tl=source_tl, target_tl=target_tl)
        if out:
            return out, "google_translate:gtx"
    except Exception as err:
        print(f"[google_translate] failed: {err}")
    return None


def free_en_to_target(english: str, target_code: str, target_label: str) -> tuple[str, str] | None:
    """EN→target: Google Translate first, then Groq, then HF opus/NLLB."""
    english = (english or "").strip()
    if not english:
        return None
    hit = google_translate_text(
        english, target_code=target_code, source_code="english"
    )
    if hit:
        return hit
    hit = groq_chat_translate(
        english,
        source_label="English",
        target_code=target_code,
        target_label=target_label,
    )
    if hit:
        return hit
    return _hf_en_to_target(english, target_code, target_label)


def build_system_prompt() -> str:
    return (
        "PRIMARY LAW: Obey Balochi_grammar_rules.pdf for every Balochi Latin string.\n\n"
        + CORE_ORTHOGRAPHY
        + "\n\n--- Full rules from Balochi_grammar_rules.pdf (binding) ---\n"
        + RESOURCES["grammar_excerpt"]
    )


def style_fewshot_en_bal(k: int = MAX_FEWSHOT) -> list[tuple[float, dict]]:
    out: list[tuple[float, dict]] = []
    for pair in RESOURCES["pairs"]:
        if len(pair["en"]) > 90 or len(pair["bal"]) > 90:
            continue
        out.append((0.45, pair))
        if len(out) >= k:
            break
    return out


def build_user_prompt(
    text: str,
    route: dict,
    fewshot: list[tuple[float, dict]],
    english_pivot: str | None = None,
    text_format: str | None = None,
    support_draft: str | None = None,
    support_draft_rejected: bool = False,
    lex_hits: list[tuple[str, str]] | None = None,
) -> str:
    route_key = route["key"]
    target_label = route["target_label"]
    to_balochi = route_key in {"en-bal", "any-bal"}

    if route_key in {"en-bal", "any-bal"}:
        src_l = route.get("source_label") or (
            "English" if route_key == "en-bal" else route.get("source", "source language")
        )
        tgt_l = "Balochi Latin Syáhag"
        pair_lines = [f"EN: {p['en']}\nBAL: {p['bal']}" for _, p in fewshot]
        if route_key == "any-bal" and not is_english(route.get("source", "")):
            bridge = (
                f"The source is {src_l} (not Balochi). "
                "Understand the meaning, then write natural Balochi Latin Syáhag. "
                "Use English↔Balochi dataset pairs as style/sense guides. "
                "Apply Balochi_grammar_rules.pdf strictly (SOV, endings, orthography)."
            )
        else:
            bridge = ""
        script_rule = (
            "CRITICAL for Balochi output: Latin Syáhag ONLY (never Arabic/Perso-Arabic). "
            "Alphabet A Á B Ch D Dh E É G H I J K L M N O Ó P R S Sh T Th U W Y Z Zh. "
            "Never use f, q, v, x. Follow Balochi grammar (SOV, ent/ant, verb endings)."
        )
    elif route_key == "bal-en":
        src_l, tgt_l = "Balochi Latin Syáhag", "English"
        pair_lines = [f"BAL: {p['bal']}\nEN: {p['en']}" for _, p in fewshot]
        bridge = (
            "FIDELITY RULES (mandatory):\n"
            "- Translate ONLY the Balochi source meaning — nothing else.\n"
            "- Do NOT invent scenes, religions, places, or English idioms "
            "that are not in the source (e.g. no Hajj/Kaaba/pilgrims/'early bird' "
            "unless those ideas are actually present in the Balochi).\n"
            "- CRITICAL: Balochi 'hajjám' means BARBER (not pilgrim). "
            "Do not confuse it with Hajj.\n"
            "- Do NOT replace the text with an unrelated story.\n"
            "- Keep the same paragraph / line structure as the source.\n"
            "- Lexicon glosses are BINDING for those words (hajjám = Barber)."
        )
        script_rule = "Write natural English in Latin script."
    else:
        src_l, tgt_l = "Balochi Latin Syáhag", target_label
        pair_lines = [f"BAL: {p['bal']}\nEN (sense): {p['en']}" for _, p in fewshot]
        bridge = (
            f"The source is Balochi in Latin Syáhag script (not Arabic script). "
            f"Use dataset English glosses only as meaning support, then write natural {target_label}."
        )
        if wants_arabic_script(route):
            script_rule = (
                f"CRITICAL: Write the entire {target_label} translation in Arabic / Perso-Arabic script only. "
                "Do NOT use Latin transliteration (no Roman Urdu / Roman Persian)."
            )
        else:
            script_rule = f"Use the standard writing system for {target_label}."

    priority = (
        "SUPPORT STACK (mandatory order — matches genai_support):\n"
        "1) dataset.xlsx / dataset.csv — search words & sentences (BINDING sense)\n"
        "2) words_balochi_to_english.xlsx — word senses (BINDING; "
        "hajjám = Barber — NEVER pilgrim/Kaaba/Hajj)\n"
        "3) Balochi_grammar_rules.pdf — understand Balochi morphology/syntax\n"
        "4) Produce faithful English from the Balochi source using (1)–(3)\n"
        "5) Then render that English meaning into the selected target language.\n"
        "Do NOT use TranslateGemma / LoRA adapter drafts."
    )

    sections = [
        f"Translate from {src_l} to {tgt_l}.",
        "Keep meaning faithful. Reply with the translation ONLY.",
        priority,
        structure_hint(text, text_format),
        script_rule,
    ]
    if to_balochi:
        sections.append(GRAMMAR_MUST_FOLLOW)
    if bridge:
        sections.append(bridge)
    # support_draft ignored — TranslateGemma adapter is disabled.
    _ = (support_draft, support_draft_rejected)
    if pair_lines:
        sections.append(
            "SUPPORT 1 — parallel examples from dataset.xlsx / dataset.csv "
            "(BINDING sense guides):\n"
            + "\n\n".join(pair_lines)
        )
    if english_pivot:
        sections.append(
            "SUPPORT 1b — exact/near dataset English sense:\n"
            f"{english_pivot}"
        )
    if lex_hits:
        gloss = "\n".join(f"{a} → {b}" for a, b in lex_hits)
        sections.append(
            "SUPPORT 2 — BINDING lexicon glosses from words_balochi_to_english.xlsx:\n"
            f"{gloss}"
        )
    if to_balochi:
        sections.append(
            "SUPPORT 3 — BINDING text from Balochi_grammar_rules.pdf "
            "(follow these rules exactly):\n"
            + RESOURCES["grammar_excerpt"][:12000]
        )
    elif route_key in {"bal-en", "bal-any"}:
        sections.append(
            "SUPPORT 3 — Understand Balochi Latin Syáhag with Balochi_grammar_rules.pdf "
            "(Latin orthography, SOV, endings) before writing the translation."
        )
    sections.append(
        "SUPPORT 4 — Produce the final translation using supports (1)–(3) only.\n"
        f"Source text:\n{text}\n\nTranslation:"
    )
    return "\n\n".join(sections)


def postprocess_balochi(text: str) -> str:
    text = text.strip().strip('"').strip("'")
    text = re.sub(r"[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDFF\uFE70-\uFEFF]+", "", text)
    repl = {"f": "p", "F": "P", "q": "k", "Q": "K", "v": "w", "V": "W", "x": "ks", "X": "Ks"}
    text = "".join(repl.get(ch, ch) for ch in text)
    return re.sub(r"\s+", " ", text).strip()


class TranslationRequest(BaseModel):
    text: str
    source_lang: str
    target_lang: str
    text_format: str | None = Field(
        default=None,
        description="Optional: poetry | story — controls output layout",
    )


class TranslationResponse(BaseModel):
    original_text: str
    translated_text: str
    source_lang: str
    target_lang: str
    method: str = Field(
        description="translategemma | exact_dataset | corpus_near_exact | grounded_model"
    )
    support: dict = Field(default_factory=dict)


@app.get("/health")
def health():
    return {
        "ok": True,
        "model": GEMINI_MODEL,
        "fallbacks": GEMINI_FALLBACKS,
        "thinking_budget": GEMINI_THINKING_BUDGET,
        "gemini_key_set": bool(os.getenv("GEMINI_API_KEY")),
        "groq_key_set": groq_configured(),
        "groq_model": groq_model_id(),
        "google_translate": True,
        "google_translate_cloud_key_set": bool(
            (os.getenv("GOOGLE_TRANSLATE_API_KEY") or os.getenv("GOOGLE_API_KEY") or "").strip()
        ),
        "pairs": len(RESOURCES["pairs"]),
        "lexicon": len(RESOURCES["bal_lex"]),
        "grammar_chars": len(RESOURCES["grammar_excerpt"]),
        "dataset_source": RESOURCES.get("dataset_path"),
        "words_source": RESOURCES.get("words_path"),
        "grammar_source": (
            GRAMMAR_CACHE.name if GRAMMAR_CACHE.exists() else GRAMMAR_PDF.name
        ),
        "supporting_model": {
            "enabled": False,
            "ready": False,
            "note": "disabled — Groq uses dataset + lexicon + grammar only",
        },
        "priority": PRIORITIES,
        "groq_support": groq_support_stack() if groq_configured() else [],
        "genai_support": GENAI_SUPPORT,
        "pipeline": [
            "dataset.xlsx sentence/word search",
            "words_balochi_to_english.xlsx word senses",
            "Balochi_grammar_rules.pdf understanding",
            "groq Balochi→selected language",
        ],
        "priority_adapter": None,
        "translategemma_hub": None,
        "routes": ["balochi→any", "any→balochi", "english→balochi"],
    }


@app.get("/languages")
def languages():
    return {
        "source_default": "balochi",
        "targets": [{"code": code, "name": name} for code, name in TARGET_LANGUAGES],
        "also_supported": (
            "Balochi → any language as target_lang; "
            "any language as source_lang → balochi (grammar rules applied)"
        ),
    }


@app.get("/")
def ui():
    index = STATIC_DIR / "index.html"
    if not index.exists():
        raise HTTPException(status_code=404, detail="UI missing: static/index.html")
    return FileResponse(index)


if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


def corpus_compose_lines(text: str, route_key: str) -> tuple[str, dict] | None:
    """Line-compose from dataset; unmatched lines keep the original text."""
    if route_key not in {"bal-en", "en-bal"}:
        return None
    lines = text.split("\n")
    out: list[str] = []
    matched = 0
    kept_original = 0
    sources: list[str] = []
    for line in lines:
        raw = line.strip()
        if not raw:
            out.append("")
            continue
        hit = exact_lookup(raw, route_key)
        if hit is None:
            out.append(raw)  # keep original Balochi / English
            kept_original += 1
            sources.append(raw)
            continue
        matched += 1
        out.append(hit)
        sources.append(raw)
    if matched == 0:
        return None
    return "\n".join(out), {
        "matched_lines": matched,
        "kept_original_lines": kept_original,
        "sources": sources,
        "note": "unmatched_lines_kept_as_original",
    }


@app.post("/translate", response_model=TranslationResponse)
async def translate_text(request: TranslationRequest):
    text_clean = norm_preserve_structure(request.text)
    if not text_clean:
        raise HTTPException(status_code=400, detail="text is empty")

    route = parse_route(request.source_lang, request.target_lang)
    route_key = route["key"]

    exact_en_or_bal = exact_lookup(text_clean, route_key) if "\n" not in text_clean else None
    if exact_en_or_bal is not None and route_key in {"en-bal", "bal-en"}:
        return TranslationResponse(
            original_text=text_clean,
            translated_text=exact_en_or_bal,
            source_lang=request.source_lang,
            target_lang=request.target_lang,
            method="exact_dataset",
            support={
                "matched": text_clean,
                "dataset_source": RESOURCES.get("dataset_path"),
                "grammar_source": GRAMMAR_CACHE.name if GRAMMAR_CACHE.exists() else GRAMMAR_PDF.name,
                "priority": PRIORITIES,
            },
        )

    if route_key == "any-bal":
        fewshot = retrieve_fewshot(text_clean, "en-bal") or style_fewshot_en_bal()
    elif route_key == "bal-any":
        fewshot = retrieve_fewshot(text_clean, "bal-en")
    else:
        fewshot = retrieve_fewshot(text_clean, route_key)

    text_format = (request.text_format or "").strip().casefold() or None
    if text_format not in {None, "poetry", "story"}:
        text_format = None

    support_draft = None
    support_draft_rejected = False
    support_draft_meta: dict = {}

    # TranslateGemma draft for short en↔bal AND bal→any (English pivot for Urdu etc.).
    # Long labzank essays skip LoRA to avoid ZeroGPU + 90s timeouts.
    tg_route = (
        "bal-en"
        if route_key == "bal-any"
        else (route_key if route_key in {"en-bal", "bal-en"} else None)
    )
    if tg_route and should_run_translategemma(text_clean, text_format):
        try:
            from support_translategemma import draft_translation, translategemma_status

            def _tg_max_tokens(src: str) -> int:
                base = int(os.getenv("TRANSLATEGEMMA_MAX_NEW_TOKENS", "256"))
                words = max(1, len(src.split()))
                lines = max(1, src.count("\n") + 1)
                return min(1024, max(base, words * 2, lines * 32))

            def _tg_draft_primary(src: str, rk: str) -> str | None:
                """Line-by-line LoRA first for multi-line posts; else whole-text."""

                def _from_lines(s: str) -> str | None:
                    out_lines: list[str] = []
                    for line in s.split("\n"):
                        raw = line.strip()
                        if not raw:
                            out_lines.append("")
                            continue
                        part = draft_translation(raw, rk, _tg_max_tokens(raw))
                        if not part or not part.strip():
                            return None
                        out_lines.append(part.strip().strip('"').strip("'"))
                    draft = "\n".join(out_lines).strip()
                    if draft and draft_covers_source(s, draft):
                        return draft
                    return None

                if "\n" in src:
                    lined = _from_lines(src)
                    if lined:
                        return lined
                max_tokens = _tg_max_tokens(src)
                whole = draft_translation(src, rk, max_tokens)
                if whole and whole.strip() and draft_covers_source(src, whole):
                    return whole.strip().strip('"').strip("'")
                if "\n" not in src and whole and whole.strip():
                    return whole.strip().strip('"').strip("'")
                return None

            on_zerogpu = os.getenv("SPACES_ZERO_GPU", "").strip().lower() in {
                "1",
                "true",
                "yes",
            }
            if on_zerogpu:
                support_draft = _tg_draft_primary(text_clean, tg_route)
            else:
                support_draft = await asyncio.to_thread(
                    _tg_draft_primary, text_clean, tg_route
                )
            support_draft_meta = translategemma_status()
            if support_draft:
                translated = support_draft
                if tg_route == "en-bal":
                    translated = postprocess_balochi(translated)
                support_draft = translated
                support_draft_meta = {
                    **support_draft_meta,
                    "draft": support_draft,
                    "adapter_dir": str(
                        Path(
                            os.getenv(
                                "TRANSLATEGEMMA_ADAPTER",
                                str(ROOT / "translategemma_balochi_adapter"),
                            )
                        )
                    ),
                    "priority_role": (
                        "english_pivot_for_bal_any"
                        if route_key == "bal-any"
                        else "draft_for_genai"
                    ),
                }
                if not draft_covers_source(text_clean, translated):
                    support_draft_meta = {
                        **support_draft_meta,
                        "rejected_short_draft": True,
                        "draft_lines": non_empty_line_count(translated),
                        "source_lines": non_empty_line_count(text_clean),
                    }
                    support_draft_rejected = True
                if tg_route == "bal-en":
                    conflicts = lexicon_draft_conflicts(
                        text_clean, translated, "bal-en"
                    )
                    if conflicts:
                        support_draft_meta = {
                            **support_draft_meta,
                            "rejected_lexicon_conflict": conflicts,
                        }
                        support_draft_rejected = True
                    elif _looks_like_hallucinated_mt(text_clean, translated) or (
                        _english_misses_lexicon(text_clean, translated, strict=True)
                        or draft_misses_lexicon_senses(
                            text_clean, translated, "bal-en"
                        )
                    ):
                        support_draft_meta = {
                            **support_draft_meta,
                            "rejected_unrelated_draft": True,
                        }
                        support_draft_rejected = True
                # bal→en / en→bal: accept a clean LoRA draft WITHOUT Gemini.
                if (
                    route_key in {"bal-en", "en-bal"}
                    and not support_draft_rejected
                    and support_draft
                    and draft_covers_source(text_clean, support_draft)
                    and not _looks_like_hallucinated_mt(text_clean, support_draft)
                    and (
                        route_key == "en-bal"
                        or not _english_misses_lexicon(
                            text_clean, support_draft, strict=True
                        )
                    )
                    and (
                        route_key == "bal-en"
                        or not violates_syahag(support_draft)
                    )
                ):
                    return TranslationResponse(
                        original_text=text_clean,
                        translated_text=support_draft,
                        source_lang=request.source_lang,
                        target_lang=request.target_lang,
                        method="translategemma",
                        support={
                            "supporting_model": {
                                **support_draft_meta,
                                "priority_role": "accepted_without_genai",
                            },
                            "dataset_source": RESOURCES.get("dataset_path"),
                            "grammar_source": (
                                GRAMMAR_CACHE.name
                                if GRAMMAR_CACHE.exists()
                                else GRAMMAR_PDF.name
                            ),
                            "priority": PRIORITIES,
                            "genai_support": GENAI_SUPPORT,
                            "note": "adapter_ok_dataset_lexicon_checked",
                            "syahag_ok": (
                                not bool(violates_syahag(support_draft))
                                if route_key == "en-bal"
                                else None
                            ),
                        },
                    )
                # Do NOT soft-accept weak LoRA poems — fall through to
                # dataset / grounded Groq instead of shipping invented lines.
        except Exception as err:
            support_draft_meta = {"ready": False, "error": str(err)}
    elif route_key in {"en-bal", "bal-en", "bal-any"}:
        try:
            from support_translategemma import translategemma_status

            support_draft_meta = {
                **translategemma_status(),
                "skipped_long_text": True,
                "chars": len(text_clean),
                "words": len(text_clean.split()),
                "lines": non_empty_line_count(text_clean),
                "priority_role": "skipped_for_latency",
            }
        except Exception:
            support_draft_meta = {"skipped_long_text": True}

    composed = corpus_compose_lines(text_clean, route_key)
    if composed is not None:
        translated, meta = composed
        if route_key in {"en-bal", "any-bal"}:
            translated = postprocess_balochi(translated)
        return TranslationResponse(
            original_text=text_clean,
            translated_text=translated,
            source_lang=request.source_lang,
            target_lang=request.target_lang,
            method="exact_dataset_lines",
            support={
                **meta,
                "dataset_source": RESOURCES.get("dataset_path"),
                "grammar_source": GRAMMAR_CACHE.name if GRAMMAR_CACHE.exists() else GRAMMAR_PDF.name,
                "supporting_model": support_draft_meta,
                "priority": PRIORITIES,
            },
        )

    if (
        route_key in {"bal-en", "en-bal"}
        and "\n" not in text_clean
        and fewshot
        and fewshot[0][0] >= CORPUS_DIRECT_SCORE
    ):
        score, pair = fewshot[0]
        translated = pair["bal"] if route_key == "en-bal" else pair["en"]
        if route_key == "en-bal":
            translated = postprocess_balochi(translated)
        return TranslationResponse(
            original_text=text_clean,
            translated_text=translated,
            source_lang=request.source_lang,
            target_lang=request.target_lang,
            method="corpus_near_exact",
            support={
                "score": round(score, 3),
                "example": {"en": pair["en"], "bal": pair["bal"]},
                "dataset_source": RESOURCES.get("dataset_path"),
                "grammar_source": GRAMMAR_CACHE.name if GRAMMAR_CACHE.exists() else GRAMMAR_PDF.name,
                "supporting_model": support_draft_meta,
                "priority": PRIORITIES,
            },
        )

    english_pivot = exact_en_or_bal if route_key == "bal-any" else None
    if english_pivot is None and route_key == "bal-any" and fewshot and fewshot[0][0] >= 0.55:
        english_pivot = fewshot[0][1]["en"]
    # Prefer dataset English pivot for bal→Urdu/etc. (no adapter).
    if (
        route_key == "bal-any"
        and fewshot
        and fewshot[0][0] >= 0.45
        and draft_covers_source(text_clean, fewshot[0][1]["en"])
    ):
        english_pivot = fewshot[0][1]["en"]

    lex_route = "en-bal" if route_key in {"en-bal", "any-bal"} else "bal-en"
    lex_hits = lexicon_hits(text_clean, lex_route)

    # bal→Urdu/etc.: Groq reads dataset+lexicon+grammar, then translates
    # Balochi → target itself (no TranslateGemma; Google only as last resort).
    if route_key == "bal-any" and groq_configured():
        tgt_code = route.get("target", "")
        tgt_label = route.get("target_label", "")
        fmt = text_format or "poetry"
        free = groq_translate_chunked(
            text_clean,
            source_label="Balochi Latin Syáhag",
            target_code=tgt_code,
            target_label=tgt_label,
            text_format=fmt,
            route_key="bal-en",
            adapter_draft=None,
            strict_quality=True,
        )
        note = "groq_grounded_balochi_to_target"
        pivot = None
        if not free:
            # Retry without strict filters
            free = groq_translate_chunked(
                text_clean,
                source_label="Balochi Latin Syáhag",
                target_code=tgt_code,
                target_label=tgt_label,
                text_format=fmt,
                route_key="bal-en",
                adapter_draft=None,
                strict_quality=False,
            )
        if not free:
            # Last resort: grounded English then Google / Groq EN→target
            pivot = english_pivot
            if (
                not pivot
                or non_empty_line_count(text_clean) > 6
                or _looks_like_hallucinated_mt(text_clean, pivot or "")
                or _english_misses_lexicon(text_clean, pivot or "")
            ):
                en_hit = groq_translate_chunked(
                    text_clean,
                    source_label="Balochi Latin Syáhag",
                    target_code="english",
                    target_label="English",
                    text_format=fmt,
                    route_key="bal-en",
                    adapter_draft=None,
                )
                if en_hit:
                    pivot = en_hit[0]
            if pivot:
                free = free_en_to_target(pivot, tgt_code, tgt_label)
                if free and str(free[1]).startswith("google_translate"):
                    note = "fallback_google_via_grounded_english"
                elif free:
                    note = "fallback_groq_via_grounded_english"
        if free:
            text, model_id = free
            if fmt != "story":
                text = align_line_structure(text_clean, text)
            method = (
                "google_translate"
                if str(model_id).startswith("google_translate")
                else "groq_grounded"
            )
            return TranslationResponse(
                original_text=text_clean,
                translated_text=text,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method=method,
                support={
                    "model": model_id,
                    "english_pivot": pivot,
                    "note": note,
                    "lexicon_hits": lex_hits[:12],
                    "supporting_model": {
                        "enabled": False,
                        "ready": False,
                        "note": "translategemma_disabled",
                    },
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                    "groq_support": groq_support_stack(),
                    "google_translate": str(model_id).startswith("google_translate"),
                    "routes": ["balochi→any", "any→balochi", "english→balochi"],
                },
            )
        raise HTTPException(
            status_code=502,
            detail=(
                "Translation failed after dataset/lexicon/grammar/Groq "
                "(no TranslateGemma). Retry in a minute."
            ),
        )

    # bal→English: dataset + lexicon + grammar → Groq English. No adapter / Gemini.
    if route_key == "bal-en" and groq_configured():
        lex_hits_be = lexicon_hits(text_clean, "bal-en")
        # 1) Strong dataset near-match / line compose
        if fewshot and fewshot[0][0] >= 0.35:
            score, pair = fewshot[0]
            if draft_covers_source(text_clean, pair["en"]) or non_empty_line_count(text_clean) <= 2:
                return TranslationResponse(
                    original_text=text_clean,
                    translated_text=pair["en"],
                    source_lang=request.source_lang,
                    target_lang=request.target_lang,
                    method="corpus_near_exact",
                    support={
                        "score": round(score, 3),
                        "example": {"en": pair["en"], "bal": pair["bal"]},
                        "lexicon_hits": lex_hits_be[:12],
                        "note": "dataset_sentence_search",
                        "priority": PRIORITIES,
                        "genai_support": GENAI_SUPPORT,
                    },
                )
        composed = corpus_compose_lines(text_clean, "bal-en")
        if composed:
            text, meta = composed
            return TranslationResponse(
                original_text=text_clean,
                translated_text=text,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="exact_dataset_lines",
                support={
                    **meta,
                    "lexicon_hits": lex_hits_be[:12],
                    "note": "dataset_line_compose",
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                },
            )
        # 2) AI English grounded by dataset + lexicon + grammar
        fmt = text_format
        if not fmt:
            lines_n = non_empty_line_count(text_clean)
            avg = len(text_clean) / max(lines_n, 1)
            fmt = "poetry" if lines_n >= 4 and avg < 90 else "story"
        free = groq_translate_chunked(
            text_clean,
            source_label="Balochi Latin Syáhag",
            target_code="english",
            target_label="English",
            text_format=fmt,
            route_key="bal-en",
            adapter_draft=None,
            strict_quality=True,
        )
        if not free:
            free = groq_translate_chunked(
                text_clean,
                source_label="Balochi Latin Syáhag",
                target_code="english",
                target_label="English",
                text_format=fmt,
                route_key="bal-en",
                adapter_draft=None,
                strict_quality=False,
            )
        if free:
            text, model_id = free
            if fmt != "story":
                text = align_line_structure(text_clean, text)
            return TranslationResponse(
                original_text=text_clean,
                translated_text=text,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="groq_grounded",
                support={
                    "model": model_id,
                    "note": "dataset_lexicon_grammar_then_groq_english",
                    "lexicon_hits": lex_hits_be[:12],
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                    "groq_support": groq_support_stack(),
                },
            )
        raise HTTPException(
            status_code=502,
            detail=(
                "Balochi→English failed after dataset/lexicon/grammar/Groq "
                "(TranslateGemma off; Gemini skipped). Retry in a minute."
            ),
        )

    if route_key == "en-bal" and groq_configured():
        free = groq_translate_chunked(
            text_clean,
            source_label="English",
            target_code="balochi",
            target_label="Balochi Latin Syáhag",
            text_format=text_format or "story",
            route_key="en-bal",
            adapter_draft=None,
        )
        if free:
            text, model_id = free
            text = postprocess_balochi(text)
            return TranslationResponse(
                original_text=text_clean,
                translated_text=text,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="groq_grounded",
                support={
                    "model": model_id,
                    "note": "dataset_lexicon_grammar_then_groq_balochi",
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                    "groq_support": groq_support_stack(),
                },
            )
        raise HTTPException(
            status_code=502,
            detail=(
                "English→Balochi failed after dataset/lexicon/grammar/Groq "
                "(TranslateGemma off; Gemini skipped). Retry in a minute."
            ),
        )

    # Short bal↔en without Groq: dataset only (no adapter / Gemini).
    if route_key in {"bal-en", "en-bal"} and not groq_configured():
        if fewshot and fewshot[0][0] >= 0.22:
            score, pair = fewshot[0]
            translated = pair["bal"] if route_key == "en-bal" else pair["en"]
            if route_key == "en-bal":
                translated = postprocess_balochi(translated)
            return TranslationResponse(
                original_text=text_clean,
                translated_text=translated,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="corpus_near_exact",
                support={
                    "score": round(score, 3),
                    "example": {"en": pair["en"], "bal": pair["bal"]},
                    "note": "dataset_only_no_groq_key",
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                },
            )
        raise HTTPException(
            status_code=503,
            detail=(
                "Set GROQ_API_KEY for AI English/target MT after dataset+lexicon+grammar. "
                "TranslateGemma adapter is disabled."
            ),
        )

    # Cooldown: do not call Gemini at all while free-tier is exhausted.
    if not gemini_allowed():
        wait = max(1, int(_gemini_quota_until - time.time()))
        if route_key == "bal-any" and (english_pivot or support_draft):
            pivot = english_pivot or support_draft or ""
            free = free_en_to_target(
                pivot, route.get("target", ""), route.get("target_label", "")
            )
            if free:
                text, model_id = free
                if text_format != "story":
                    text = align_line_structure(text_clean, text)
                return TranslationResponse(
                    original_text=text_clean,
                    translated_text=text,
                    source_lang=request.source_lang,
                    target_lang=request.target_lang,
                    method="free_mt",
                    support={
                        "model": model_id,
                        "note": "gemini_quota_used_free_hf_or_groq",
                        "english_pivot": pivot,
                        "retry_after_sec": wait,
                        "supporting_model": support_draft_meta,
                        "priority": PRIORITIES,
                        "genai_support": GENAI_SUPPORT,
                    },
                )
            return TranslationResponse(
                original_text=text_clean,
                translated_text=pivot,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="english_sense_quota_fallback",
                support={
                    "warning": (
                        f"Gemini cooldown (~{wait}s). Free HF/Groq MT unavailable; "
                        f"showing English sense. Set GROQ_API_KEY or wait / enable billing."
                    ),
                    "retry_after_sec": wait,
                    "supporting_model": support_draft_meta,
                    "priority": PRIORITIES,
                },
            )
        raise HTTPException(
            status_code=429,
            detail=(
                f"Gemini free-tier cooldown — retry in {wait}s, use English, "
                "or enable billing on GEMINI_API_KEY / set GROQ_API_KEY."
            ),
        )

    # Hard stop: with GROQ_API_KEY set, never call Gemini for labzank routes.
    if groq_configured() and route_key in {"bal-en", "en-bal", "bal-any", "any-bal"}:
        raise HTTPException(
            status_code=502,
            detail=(
                "Groq/Google path unavailable; Gemini skipped because "
                "GROQ_API_KEY is set. Retry in a minute."
            ),
        )

    async def _genai_one(
        chunk: str,
        *,
        draft: str | None,
        draft_rejected: bool,
        fmt: str | None,
        pivot: str | None = None,
    ) -> tuple[str, str]:
        # Fast path: EN→target with a tiny prompt when we already have English.
        if route_key == "bal-any" and pivot and chunk == text_clean:
            try:
                pivot_prompt = build_pivot_prompt(
                    pivot,
                    route["target_label"],
                    text_format=fmt,
                    arabic_script=wants_arabic_script(route),
                )
                out, model = await asyncio.to_thread(
                    generate_with_retries, pivot_prompt, light=True
                )
                return out.strip().strip('"').strip("'"), model
            except Exception as err:
                if _is_quota_exhausted(err) or not gemini_allowed():
                    free = free_en_to_target(
                        pivot, route.get("target", ""), route.get("target_label", "")
                    )
                    if free:
                        return free
                raise

        chunk_few = retrieve_fewshot(
            chunk,
            "en-bal"
            if route_key == "any-bal"
            else ("bal-en" if route_key == "bal-any" else route_key),
        )
        if route_key == "any-bal" and not chunk_few:
            chunk_few = fewshot
        chunk_lex = lexicon_hits(chunk, lex_route)
        prompt = build_user_prompt(
            chunk,
            route,
            chunk_few or fewshot,
            english_pivot=pivot if chunk == text_clean else None,
            text_format=fmt,
            support_draft=draft if chunk == text_clean else None,
            support_draft_rejected=draft_rejected if chunk == text_clean else False,
            lex_hits=chunk_lex or lex_hits,
        )
        try:
            out, model = await asyncio.to_thread(generate_with_retries, prompt)
        except Exception as err:
            if (
                route_key == "bal-any"
                and pivot
                and (_is_quota_exhausted(err) or not gemini_allowed())
            ):
                free = free_en_to_target(
                    pivot, route.get("target", ""), route.get("target_label", "")
                )
                if free:
                    return free
            raise
        if route_key in {"en-bal", "any-bal"}:
            out = postprocess_balochi(out)
        else:
            out = out.strip().strip('"').strip("'")
        return out, model

    try:
        used_model = GEMINI_MODEL
        chunked = False
        if should_chunk_genai(text_clean, text_format) and route_key != "bal-any":
            paras = story_paragraphs(text_clean)
            # Merge tiny paragraphs so we don't fire dozens of Gemini calls.
            batches: list[str] = []
            buf = ""
            for p in paras:
                if not buf:
                    buf = p
                elif len(buf) + len(p) + 2 <= GENAI_CHUNK_CHARS:
                    buf = f"{buf}\n\n{p}"
                else:
                    batches.append(buf)
                    buf = p
            if buf:
                batches.append(buf)
            outs: list[str] = []
            for batch in batches:
                part, used_model = await _genai_one(
                    batch,
                    draft=None,
                    draft_rejected=False,
                    fmt=text_format or "story",
                )
                outs.append(part)
            translated = "\n\n".join(outs)
            chunked = True
        else:
            translated, used_model = await _genai_one(
                text_clean,
                draft=support_draft,
                draft_rejected=support_draft_rejected,
                fmt=text_format,
                pivot=english_pivot,
            )

        if text_format != "story" and not chunked:
            translated = align_line_structure(text_clean, translated)

        if route_key in {"en-bal", "any-bal"} and not chunked:
            issues = violates_syahag(translated)
            if issues:
                repair_prompt = (
                    build_user_prompt(
                        text_clean,
                        route,
                        fewshot,
                        english_pivot=english_pivot,
                        text_format=text_format,
                        support_draft=support_draft,
                        support_draft_rejected=support_draft_rejected,
                        lex_hits=lex_hits,
                    )
                    + "\n\nYour previous Balochi draft VIOLATES Balochi_grammar_rules.pdf "
                    f"({', '.join(issues)}).\n"
                    "Rewrite the FULL translation in correct Balochi Latin Syáhag only, "
                    "obeying the PDF rules. Reply with the corrected translation only:\n"
                    f"{translated}\n\nCorrected translation:"
                )
                translated, used_model = await asyncio.to_thread(
                    generate_with_retries, repair_prompt
                )
                translated = postprocess_balochi(translated)
                if text_format != "story":
                    translated = align_line_structure(text_clean, translated)

        if (
            wants_arabic_script(route)
            and arabic_script_ratio(translated) < 0.55
            and not chunked
        ):
            fix_prompt = (
                build_user_prompt(
                    text_clean,
                    route,
                    fewshot,
                    english_pivot=english_pivot,
                    text_format=text_format,
                    support_draft=support_draft,
                    support_draft_rejected=support_draft_rejected,
                    lex_hits=lex_hits,
                )
                + "\n\nYour previous draft used Latin letters. "
                f"Rewrite the full translation in Arabic script for {route['target_label']} only, "
                "preserving the same layout:\n"
                f"{translated}\n\nCorrected translation:"
            )
            translated, used_model = await asyncio.to_thread(generate_with_retries, fix_prompt)
            translated = translated.strip().strip('"').strip("'")
            if text_format != "story":
                translated = align_line_structure(text_clean, translated)

        return TranslationResponse(
            original_text=text_clean,
            translated_text=translated,
            source_lang=request.source_lang,
            target_lang=request.target_lang,
            method="grounded_model",
            support={
                "model": used_model,
                "route": route_key,
                "target_label": route["target_label"],
                "source_label": route.get("source_label"),
                "text_format": text_format or "auto",
                "english_pivot": english_pivot,
                "arabic_script_required": wants_arabic_script(route),
                "source_lines": len(text_clean.split("\n")),
                "output_lines": len(translated.split("\n")),
                "chunked": chunked,
                "fewshot": [{"score": round(s, 3), "en": p["en"], "bal": p["bal"]} for s, p in fewshot],
                "grammar_source": GRAMMAR_CACHE.name if GRAMMAR_CACHE.exists() else GRAMMAR_PDF.name,
                "grammar_chars": len(RESOURCES["grammar_excerpt"]),
                "syahag_ok": not bool(violates_syahag(translated))
                if route_key in {"en-bal", "any-bal"}
                else None,
                "dataset_source": RESOURCES.get("dataset_path"),
                "supporting_model": support_draft_meta,
                "priority": PRIORITIES,
                "genai_support": GENAI_SUPPORT,
                "priority_adapter": TRANSLATEGEMMA_ADAPTER_DIR,
                "translategemma_hub": TRANSLATEGEMMA_HUB,
                "lexicon": len(RESOURCES["bal_lex"]),
            },
        )
    except HTTPException:
        raise
    except Exception as err:
        # bal→en: never die on Gemini quota if we already have a LoRA draft.
        if (
            route_key == "bal-en"
            and support_draft
            and _is_quota_exhausted(err)
            and not support_draft_meta.get("rejected_lexicon_conflict")
        ):
            return TranslationResponse(
                original_text=text_clean,
                translated_text=support_draft,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="translategemma",
                support={
                    "warning": "Gemini quota hit; returning adapter draft.",
                    "gemini_error": str(getattr(err, "message", err)),
                    "supporting_model": {
                        **support_draft_meta,
                        "priority_role": "lora_after_gemini_quota",
                    },
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                },
            )
        # bal→Urdu/etc.: if Gemini free quota is dead but we have English, return it.
        if (
            route_key == "bal-any"
            and english_pivot
            and _is_quota_exhausted(err)
        ):
            return TranslationResponse(
                original_text=text_clean,
                translated_text=english_pivot,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="english_sense_quota_fallback",
                support={
                    "warning": (
                        f"{route['target_label']} needs Gemini; free-tier quota "
                        "exceeded. Showing English sense from the adapter/dataset. "
                        "Wait ~1 min or enable billing on GEMINI_API_KEY."
                    ),
                    "gemini_error": str(getattr(err, "message", err)),
                    "retry_after_sec": _quota_retry_seconds(err),
                    "route": route_key,
                    "target_label": route["target_label"],
                    "english_pivot": english_pivot,
                    "supporting_model": support_draft_meta,
                    "priority": PRIORITIES,
                    "genai_support": GENAI_SUPPORT,
                },
            )
        offline = offline_fallback(
            text_clean,
            route_key,
            fewshot,
            route["target_label"],
        )
        if offline is not None:
            translated, meta = offline
            if route_key in {"en-bal", "any-bal"}:
                translated = postprocess_balochi(translated)
            return TranslationResponse(
                original_text=text_clean,
                translated_text=translated,
                source_lang=request.source_lang,
                target_lang=request.target_lang,
                method="offline_fallback",
                support={
                    "gemini_error": str(getattr(err, "message", err)),
                    "route": route_key,
                    "target_label": route["target_label"],
                    **meta,
                    "priority": PRIORITIES,
                },
            )
        detail = getattr(err, "message", None) or str(err)
        if _is_quota_exhausted(err):
            raise HTTPException(
                status_code=429,
                detail=(
                    f"Gemini free-tier quota exceeded. Retry in "
                    f"{int(_quota_retry_seconds(err))}s, use English, "
                    f"or enable billing on GEMINI_API_KEY. ({detail})"
                ),
            )
        if isinstance(err, APIError):
            raise HTTPException(status_code=502, detail=f"Gemini API Error: {err.message}")
        raise HTTPException(status_code=502, detail=f"Gemini API Error: {err}")


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "8000"))
    uvicorn.run(app, host="0.0.0.0", port=port)