# syntax=docker/dockerfile:1

FROM python:3.12-slim-bookworm

ARG MODEL_REPO="Nextcloud-AI/madlad400-3b-mt-ct2-int8"

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    HF_HOME=/tmp/huggingface \
    MODEL_PATH=/models/madlad400-3b-mt-ct2-int8 \
    MODEL_COMPUTE_TYPE=int8 \
    MODEL_CONCURRENCY=1 \
    MAX_INPUT_TOKENS=512 \
    MAX_OUTPUT_TOKENS=512 \
    LT_DISABLE_FILES_TRANSLATION=true \
    LT_REQ_LIMIT=60 \
    LT_BATCH_LIMIT=5 \
    LT_CHAR_LIMIT=1000 \
    LT_THREADS=2 \
    LT_FRONTEND_TIMEOUT=500

RUN apt-get update \
    && apt-get install -y --no-install-recommends libgomp1 \
    && rm -rf /var/lib/apt/lists/* \
    && pip install \
        "beautifulsoup4>=4.13,<5" \
        "ctranslate2>=4.6,<5" \
        "flask>=3.1,<4" \
        "flask-cors>=5,<7" \
        "flask-limiter>=3.12,<5" \
        "gunicorn>=23,<24" \
        "huggingface-hub>=0.34,<2" \
        "lingua-language-detector>=2.1,<3" \
        "pycountry>=24.6,<27" \
        "sentencepiece>=0.2,<1"

RUN mkdir -p /models/madlad400-3b-mt-ct2-int8 /app \
    && python - <<PY
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="${MODEL_REPO}",
    local_dir="/models/madlad400-3b-mt-ct2-int8",
    allow_patterns=[
        "model.bin",
        "config.json",
        "shared_vocabulary.json",
        "spiece.model",
        "sentencepiece.model",
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
        "added_tokens.json",
        "generation_config.json",
    ],
)
PY

RUN cat > /app/app.py <<'PY'
from __future__ import annotations

import html
import json
import os
import re
import threading
import time
import uuid
from pathlib import Path
from typing import Any

import ctranslate2
import pycountry
import sentencepiece as spm
from bs4 import BeautifulSoup, Comment
from flask import Flask, Response, g, jsonify, request
from flask_cors import CORS
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from lingua import LanguageDetectorBuilder


MODEL_PATH = Path(os.getenv("MODEL_PATH", "/models/madlad400-3b-mt-ct2-int8"))
THREADS = max(1, int(os.getenv("LT_THREADS", "2")))
REQ_LIMIT = int(os.getenv("LT_REQ_LIMIT", "60"))
BATCH_LIMIT = int(os.getenv("LT_BATCH_LIMIT", "5"))
CHAR_LIMIT = int(os.getenv("LT_CHAR_LIMIT", "1000"))
FRONTEND_TIMEOUT = int(os.getenv("LT_FRONTEND_TIMEOUT", "500"))
MAX_INPUT_TOKENS = max(64, int(os.getenv("MAX_INPUT_TOKENS", "512")))
MAX_OUTPUT_TOKENS = max(64, int(os.getenv("MAX_OUTPUT_TOKENS", "512")))
MODEL_CONCURRENCY = max(1, int(os.getenv("MODEL_CONCURRENCY", "1")))
FILES_DISABLED = os.getenv("LT_DISABLE_FILES_TRANSLATION", "true").lower() not in {
    "0",
    "false",
    "no",
}

TOKENIZER_PATH = next(
    (
        path
        for path in (
            MODEL_PATH / "spiece.model",
            MODEL_PATH / "sentencepiece.model",
        )
        if path.exists()
    ),
    None,
)
if TOKENIZER_PATH is None:
    raise RuntimeError(f"SentencePiece model is missing from {MODEL_PATH}")

sentencepiece = spm.SentencePieceProcessor(model_file=str(TOKENIZER_PATH))
translator = ctranslate2.Translator(
    str(MODEL_PATH),
    device="cpu",
    compute_type=os.getenv("MODEL_COMPUTE_TYPE", "int8"),
    inter_threads=1,
    intra_threads=THREADS,
)
model_slots = threading.BoundedSemaphore(MODEL_CONCURRENCY)

# Lingua is Apache-2.0 and performs well on short social messages. MADLAD does
# not need a source-language tag, so detection is used only for LibreTranslate
# API compatibility and same-language short-circuiting.
language_detector = (
    LanguageDetectorBuilder.from_all_languages()
    .with_preloaded_language_models()
    .build()
)

app = Flask(__name__)
CORS(app)
limiter = Limiter(
    key_func=get_remote_address,
    app=app,
    default_limits=[] if REQ_LIMIT < 0 else [f"{REQ_LIMIT} per minute"],
    storage_uri="memory://",
)


_TRANSLATION_LOG_PATHS = {"/detect", "/translate"}


def _trace_header(name: str) -> str:
    return (request.headers.get(name) or "").strip()[:200]


def _set_translation_log_details(**details: Any) -> None:
    if request.path not in _TRANSLATION_LOG_PATHS:
        return
    current = getattr(g, "translation_log_details", {})
    current.update({key: value for key, value in details.items() if value is not None})
    g.translation_log_details = current


@app.before_request
def begin_translation_request_log() -> None:
    if request.path not in _TRANSLATION_LOG_PATHS:
        return
    g.translation_started_at = time.perf_counter()
    g.translation_request_id = (
        _trace_header("X-Translation-Request-Id") or uuid.uuid4().hex[:16]
    )
    g.translation_log_details = {}


@app.after_request
def finish_translation_request_log(response: Response) -> Response:
    if request.path not in _TRANSLATION_LOG_PATHS:
        return response

    started_at = getattr(g, "translation_started_at", time.perf_counter())
    request_id = getattr(g, "translation_request_id", uuid.uuid4().hex[:16])
    elapsed_ms = round((time.perf_counter() - started_at) * 1000, 1)
    details = dict(getattr(g, "translation_log_details", {}))
    log_entry: dict[str, Any] = {
        "event": "translation_api_request",
        "request_id": request_id,
        "mode": _trace_header("X-Translation-Mode") or "unknown",
        "job_id": _trace_header("X-Translation-Job-Id") or None,
        "content_table": _trace_header("X-Translation-Content-Table") or None,
        "content_id": _trace_header("X-Translation-Content-Id") or None,
        "operation": request.path.removeprefix("/"),
        "http_status": response.status_code,
        "outcome": "success" if response.status_code < 400 else "error",
        "duration_ms": elapsed_ms,
        "request_bytes": request.content_length or 0,
        "response_bytes": response.calculate_content_length() or 0,
        **details,
    }
    print(
        "TRANSLATION_API "
        + json.dumps(log_entry, ensure_ascii=False, separators=(",", ":")),
        flush=True,
    )
    response.headers["X-Translation-Request-Id"] = request_id
    return response


CODE_ALIASES = {
    "iw": "he",
    "in": "id",
    "ji": "yi",
    "fil": "tl",
    "nb": "no",
    "zh-cn": "zh",
    "zh-hans": "zh",
    "zh-sg": "zh",
    "zh-tw": "zh",
    "zh-hant": "zh",
    "zh-hk": "zh",
}


def _extract_model_codes() -> list[str]:
    codes: set[str] = set()
    for index in range(sentencepiece.get_piece_size()):
        piece = sentencepiece.id_to_piece(index)
        match = re.fullmatch(r"<2([^<>\s]+)>", piece)
        if match:
            codes.add(match.group(1))

    # This should not be needed for the official MADLAD tokenizer, but keeping
    # a fallback prevents the API from becoming unusable after a tokenizer
    # packaging change.
    if not codes:
        codes.update(
            "af am ar az be bg bn bs ca cs cy da de el en eo es et eu fa fi fr "
            "ga gl gu he hi hr hu hy id is it ja ka kk km kn ko lo lt lv mk ml "
            "mn mr ms mt my ne nl no pa pl pt ro ru si sk sl sq sr sv sw ta te "
            "th tl tr uk ur uz vi zh zu".split()
        )

    load_only = os.getenv("LT_LOAD_ONLY", "").strip()
    if load_only:
        requested = {item.strip() for item in load_only.split(",") if item.strip()}
        codes.intersection_update(requested)

    return sorted(codes)


SUPPORTED_CODES = _extract_model_codes()
SUPPORTED_CODE_LOOKUP = {code.lower(): code for code in SUPPORTED_CODES}


def _language_name(code: str) -> str:
    overrides = {
        "zh": "Chinese",
        "he": "Hebrew",
        "yi": "Yiddish",
        "tl": "Tagalog",
        "no": "Norwegian",
    }
    lowered = code.lower()
    if lowered in overrides:
        return overrides[lowered]

    base = re.split(r"[-_]", lowered, maxsplit=1)[0]
    try:
        language = (
            pycountry.languages.get(alpha_2=base)
            if len(base) == 2
            else pycountry.languages.get(alpha_3=base)
        )
        if language is not None:
            return getattr(language, "common_name", language.name)
    except (KeyError, LookupError):
        pass
    return code


def _resolve_target(code: str) -> str | None:
    raw = str(code).strip()
    if not raw:
        return None

    candidates = [raw, raw.lower(), raw.replace("_", "-").lower()]
    alias = CODE_ALIASES.get(candidates[-1])
    if alias:
        candidates.append(alias)

    base = re.split(r"[-_]", candidates[-1], maxsplit=1)[0]
    candidates.append(base)

    try:
        if len(base) == 2:
            language = pycountry.languages.get(alpha_2=base)
            if language and hasattr(language, "alpha_3"):
                candidates.append(language.alpha_3.lower())
        elif len(base) == 3:
            language = pycountry.languages.get(alpha_3=base)
            if language and hasattr(language, "alpha_2"):
                candidates.append(language.alpha_2.lower())
    except (KeyError, LookupError):
        pass

    for candidate in candidates:
        resolved = SUPPORTED_CODE_LOOKUP.get(candidate.lower())
        if resolved:
            return resolved
    return None


def _iso_code(language: Any) -> str:
    iso1 = getattr(language, "iso_code_639_1", None)
    if iso1 is not None:
        return iso1.name.lower()
    iso3 = getattr(language, "iso_code_639_3", None)
    if iso3 is not None:
        return iso3.name.lower()
    return "en"


def _detect(text: str, limit: int = 3) -> list[dict[str, Any]]:
    cleaned = re.sub(r"\s+", " ", BeautifulSoup(text, "html.parser").get_text(" ")).strip()
    if not cleaned:
        return [{"confidence": 0.0, "language": "en"}]

    values = language_detector.compute_language_confidence_values(cleaned)
    detections: list[dict[str, Any]] = []
    for value in values[: max(1, limit)]:
        detections.append(
            {
                "confidence": round(float(value.value) * 100.0, 2),
                "language": _iso_code(value.language),
            }
        )
    return detections or [{"confidence": 0.0, "language": "en"}]


def _request_payload() -> dict[str, Any]:
    payload = request.get_json(silent=True)
    if isinstance(payload, dict):
        return payload

    if request.form:
        form_payload: dict[str, Any] = request.form.to_dict(flat=True)
        q_values = request.form.getlist("q")
        if len(q_values) > 1:
            form_payload["q"] = q_values
        return form_payload

    return {}


def _error(message: str, status: int = 400):
    return jsonify({"error": message}), status


def _validate_texts(raw_q: Any) -> tuple[list[str] | None, bool, Any]:
    is_batch = isinstance(raw_q, list)
    if is_batch:
        if not raw_q:
            return None, True, _error("Invalid request: q must not be empty")
        if BATCH_LIMIT >= 0 and len(raw_q) > BATCH_LIMIT:
            return None, True, _error(f"Invalid request: batch limit is {BATCH_LIMIT}")
        if not all(isinstance(item, str) for item in raw_q):
            return None, True, _error("Invalid request: every q item must be a string")
        texts = raw_q
    elif isinstance(raw_q, str):
        texts = [raw_q]
    else:
        return None, False, _error("Invalid request: q is required")

    total_characters = sum(len(item) for item in texts)
    if CHAR_LIMIT >= 0 and total_characters > CHAR_LIMIT:
        return None, is_batch, _error(
            f"Invalid request: character limit is {CHAR_LIMIT}"
        )
    return texts, is_batch, None


def _tokenize_for_target(text: str, target: str) -> list[str]:
    return sentencepiece.encode(f"<2{target}> {text}", out_type=str)


def _decode(tokens: list[str]) -> str:
    return sentencepiece.decode(tokens).strip()


def _split_oversized_text(text: str, target: str) -> list[str]:
    if len(_tokenize_for_target(text, target)) <= MAX_INPUT_TOKENS:
        return [text]

    parts = re.split(r"(?<=[.!?。！？])\s+|\n+", text)
    chunks: list[str] = []
    current = ""
    for part in parts:
        part = part.strip()
        if not part:
            continue
        candidate = f"{current} {part}".strip()
        if current and len(_tokenize_for_target(candidate, target)) > MAX_INPUT_TOKENS:
            chunks.append(current)
            current = part
        else:
            current = candidate

        if len(_tokenize_for_target(current, target)) > MAX_INPUT_TOKENS:
            raw_tokens = sentencepiece.encode(current, out_type=str)
            current = ""
            for start in range(0, len(raw_tokens), MAX_INPUT_TOKENS - 8):
                chunks.append(_decode(raw_tokens[start : start + MAX_INPUT_TOKENS - 8]))

    if current:
        chunks.append(current)
    return chunks or [text]


def _run_model(
    texts: list[str], target: str, alternatives: int = 0
) -> tuple[list[str], list[list[str]]]:
    if not texts:
        return [], []

    hypotheses_requested = max(1, alternatives + 1)
    beam_size = max(1, hypotheses_requested)

    flattened: list[str] = []
    ownership: list[int] = []
    for owner, text in enumerate(texts):
        chunks = _split_oversized_text(text, target)
        flattened.extend(chunks)
        ownership.extend([owner] * len(chunks))

    token_batches = [_tokenize_for_target(text, target) for text in flattened]
    with model_slots:
        results = translator.translate_batch(
            token_batches,
            beam_size=beam_size,
            num_hypotheses=hypotheses_requested,
            max_decoding_length=MAX_OUTPUT_TOKENS,
            batch_type="tokens",
            max_batch_size=1024,
            repetition_penalty=1.1,
        )

    primary_chunks: list[list[str]] = [[] for _ in texts]
    alternative_chunks: list[list[list[str]]] = [
        [[] for _ in range(alternatives)] for _ in texts
    ]

    for owner, result in zip(ownership, results, strict=True):
        primary_chunks[owner].append(_decode(result.hypotheses[0]))
        for alternative_index in range(alternatives):
            hypothesis_index = alternative_index + 1
            if hypothesis_index < len(result.hypotheses):
                translated = _decode(result.hypotheses[hypothesis_index])
            else:
                translated = _decode(result.hypotheses[0])
            alternative_chunks[owner][alternative_index].append(translated)

    primary = [" ".join(chunks).strip() for chunks in primary_chunks]
    alternative_results = [
        [" ".join(chunks).strip() for chunks in per_text]
        for per_text in alternative_chunks
    ]
    return primary, alternative_results


def _translate_html(text: str, target: str) -> str:
    soup = BeautifulSoup(text, "html.parser")
    nodes = [
        node
        for node in soup.find_all(string=True)
        if not isinstance(node, Comment)
        and node.parent is not None
        and node.parent.name not in {"script", "style", "code", "pre"}
        and str(node).strip()
    ]
    if not nodes:
        return text

    translated, _ = _run_model([str(node) for node in nodes], target, alternatives=0)
    for node, replacement in zip(nodes, translated, strict=True):
        leading = re.match(r"^\s*", str(node)).group(0)
        trailing = re.search(r"\s*$", str(node)).group(0)
        node.replace_with(f"{leading}{replacement}{trailing}")
    return str(soup)


LANGUAGE_TARGETS = SUPPORTED_CODES
LANGUAGES_RESPONSE = [
    {
        "code": code,
        "name": _language_name(code),
        "targets": [target for target in LANGUAGE_TARGETS if target != code],
    }
    for code in SUPPORTED_CODES
]
LANGUAGES_JSON = json.dumps(LANGUAGES_RESPONSE, ensure_ascii=False)


@app.errorhandler(429)
def rate_limit_error(_error_value: Any):
    return _error("Slow down", 429)


@app.errorhandler(500)
def internal_error(_error_value: Any):
    return _error("Internal server error", 500)


@app.get("/")
def index():
    return Response(
        """<!doctype html><html><head><meta charset="utf-8"><title>LibreTranslate-compatible MADLAD-400 API</title></head><body><h1>LibreTranslate-compatible translation API</h1><p>MADLAD-400 3B INT8 is loaded.</p><p>Endpoints: <code>/translate</code>, <code>/detect</code>, <code>/languages</code>, <code>/health</code>.</p></body></html>""",
        mimetype="text/html",
    )


@app.get("/health")
def health():
    return jsonify({"status": "ok"})


@app.get("/languages")
def languages():
    return Response(LANGUAGES_JSON, mimetype="application/json")


@app.get("/frontend/settings")
def frontend_settings():
    default_source = "auto"
    default_target = _resolve_target(os.getenv("LT_FRONTEND_LANGUAGE_TARGET", "en"))
    if default_target is None:
        default_target = "en" if "en" in SUPPORTED_CODE_LOOKUP else SUPPORTED_CODES[0]

    return jsonify(
        {
            "apiKeys": False,
            "charLimit": CHAR_LIMIT,
            "frontendTimeout": FRONTEND_TIMEOUT,
            "keyRequired": False,
            "language": {
                "source": {"code": default_source, "name": "Detect language"},
                "target": {
                    "code": default_target,
                    "name": _language_name(default_target),
                },
            },
            "suggestions": False,
            "supportedFilesFormat": [],
        }
    )


@app.post("/detect")
def detect():
    payload = _request_payload()
    raw_q = payload.get("q")
    if not isinstance(raw_q, str) or not raw_q.strip():
        return _error("Invalid request: q is required")
    if CHAR_LIMIT >= 0 and len(raw_q) > CHAR_LIMIT:
        _set_translation_log_details(characters=len(raw_q))
        return _error(f"Invalid request: character limit is {CHAR_LIMIT}")

    detections = _detect(raw_q)
    primary = detections[0] if detections else {}
    _set_translation_log_details(
        characters=len(raw_q),
        detected_language=primary.get("language"),
        confidence=primary.get("confidence"),
    )
    return jsonify(detections)


@app.post("/translate")
def translate():
    payload = _request_payload()
    texts, is_batch, validation_error = _validate_texts(payload.get("q"))
    if validation_error is not None:
        return validation_error
    assert texts is not None

    source = str(payload.get("source", "")).strip().lower()
    requested_target = str(payload.get("target", "")).strip()
    text_format = str(payload.get("format", "text")).strip().lower()
    _set_translation_log_details(
        items=len(texts),
        characters=sum(len(text) for text in texts),
        source_requested=source or None,
        target_requested=requested_target or None,
        format=text_format or None,
    )

    if not source:
        return _error("Invalid request: source is required")
    if not requested_target:
        return _error("Invalid request: target is required")
    if text_format not in {"text", "html"}:
        return _error("Invalid request: format must be text or html")

    target = _resolve_target(requested_target)
    if target is None:
        return _error(f"Invalid target language: {html.escape(requested_target)}")

    if source != "auto" and _resolve_target(source) is None:
        return _error(f"Invalid source language: {html.escape(source)}")

    try:
        alternatives = int(payload.get("alternatives", 0) or 0)
    except (TypeError, ValueError):
        return _error("Invalid request: alternatives must be an integer")
    alternatives = max(0, min(alternatives, 10))

    detections = [_detect(text, limit=1)[0] for text in texts] if source == "auto" else []
    resolved_source = _resolve_target(source) if source != "auto" else None

    output: list[str] = [""] * len(texts)
    output_alternatives: list[list[str]] = [[] for _ in texts]
    model_texts: list[str] = []
    model_indexes: list[int] = []

    for index, text in enumerate(texts):
        detected_source = (
            _resolve_target(detections[index]["language"])
            if source == "auto"
            else resolved_source
        )
        if not text.strip() or detected_source == target:
            output[index] = text
            output_alternatives[index] = [text] * alternatives
        elif text_format == "html":
            output[index] = _translate_html(text, target)
            output_alternatives[index] = [output[index]] * alternatives
        else:
            model_indexes.append(index)
            model_texts.append(text)

    detected_languages = sorted(
        {
            str(item.get("language", "")).strip().lower()
            for item in detections
            if str(item.get("language", "")).strip()
        }
    )
    _set_translation_log_details(
        target=target,
        source_resolved=resolved_source,
        detected_languages=detected_languages or None,
        model_items=len(model_texts),
        short_circuited_items=len(texts) - len(model_texts),
        alternatives=alternatives,
    )

    if model_texts:
        translated, translated_alternatives = _run_model(
            model_texts, target, alternatives=alternatives
        )
        for local_index, original_index in enumerate(model_indexes):
            output[original_index] = translated[local_index]
            output_alternatives[original_index] = translated_alternatives[local_index]

    response: dict[str, Any] = {
        "translatedText": output if is_batch else output[0],
    }
    if source == "auto":
        response["detectedLanguage"] = detections if is_batch else detections[0]
    if alternatives > 0:
        response["alternatives"] = (
            output_alternatives if is_batch else output_alternatives[0]
        )
    return jsonify(response)


@app.post("/translate_file")
def translate_file():
    if FILES_DISABLED:
        return _error("File translation is disabled", 400)
    return _error("File translation is not implemented", 501)


@app.post("/suggest")
def suggest():
    return _error("Suggestions are disabled", 403)
PY

RUN useradd --create-home --uid 1000 translator \
    && chown -R translator:translator /app /models \
    && rm -rf /tmp/huggingface

USER translator
WORKDIR /app

EXPOSE 7860

HEALTHCHECK --interval=30s --timeout=10s --start-period=180s --retries=3 \
    CMD python -c "import json,urllib.request; assert json.load(urllib.request.urlopen('http://127.0.0.1:7860/health', timeout=5))['status']=='ok'"

CMD ["gunicorn", "--bind", "0.0.0.0:7860", "--workers", "1", "--worker-class", "gthread", "--threads", "4", "--timeout", "300", "--graceful-timeout", "30", "--capture-output", "--access-logfile", "/dev/null", "--error-logfile", "-", "app:app"]