#!/usr/bin/env python3
"""Inference CLI for foksly/wmt26-constrained-submission."""

import argparse
import sys
import time
from pathlib import Path
from typing import Dict, Optional

import torch
import transformers
from packaging.version import Version
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_ID = "foksly/wmt26-constrained-submission"
MIN_TRANSFORMERS_VERSION = Version("4.46.3")


# Public CLI uses short ISO 639-1 codes.
# WMT and ISO 639-2/3 aliases are accepted for convenience.
TARGET_ALIASES: Dict[str, str] = {
    # Russian
    "ru": "ru",
    "rus": "ru",
    "rus_cyrl": "ru",

    # Belarusian
    "be": "be",
    "bel": "be",
    "bel_cyrl": "be",

    # Kazakh
    "kk": "kk",
    "kaz": "kk",
    "kaz_cyrl": "kk",

    # Armenian
    "hy": "hy",
    "hye": "hy",
    "hye_armn": "hy",
}


TRANSLATION_PROMPTS: Dict[str, str] = {
    "ru": "Текст для перевода на русский:",
    "be": "Текст для перевода на белорусский:",
    "kk": "Текст для перевода на казахский:",
    "hy": "Текст для перевода на армянский:",
}


DOMAIN_PROMPTS: Dict[str, str] = {
    "social": (
        "Ты профессиональный переводчик. "
        "Исходный текст — публикация или комментарий из социальных сетей. "
        "Переводи в естественном разговорном стиле. "
        "Не повторяй орфографические ошибки. "
        "Сохраняй ссылки, ники, HTML-разметку и пунктуацию. "
        "Переводи хештеги так, чтобы они естественно звучали "
        "на целевом языке. "
        "В ответе возвращай только перевод."
    ),
    "speech": (
        "Ты профессиональный переводчик. "
        "Исходный текст автоматически расшифрован с устной речи "
        "и может содержать ошибки. "
        "Сохраняй разговорный стиль. "
        "Не включай нелингвистические звуки, но оставляй междометия. "
        "Если слово оборвано, восстанови его, если это возможно, "
        "иначе пропусти. "
        "Каждое предложение выводи с новой строки. "
        "В ответе возвращай только перевод."
    ),
    "news": (
        "Ты профессиональный переводчик. "
        "Исходный текст — новостная статья. "
        "Переводи в официальном журналистском стиле. "
        "Сохраняй HTML-разметку. "
        "В ответе возвращай только перевод."
    ),
    "software": (
        "Ты профессиональный переводчик. "
        "Исходный текст содержит структурированные данные "
        "программного обеспечения. "
        "Переводи только текстовые значения. "
        "Не изменяй ключи, плейсхолдеры, переменные "
        "и служебную разметку. "
        "Верни результат в том же формате, что и исходный текст. "
        "В ответе возвращай только перевод."
    ),
}


def log(message: str, quiet: bool = False) -> None:
    """Write diagnostic messages to stderr."""
    if not quiet:
        print(message, file=sys.stderr, flush=True)


def normalize_target(target: str) -> str:
    """Convert a user-provided language alias to the canonical short code."""
    normalized = target.strip().lower()

    if normalized not in TARGET_ALIASES:
        supported = "ru, be, kk, hy"
        raise ValueError(
            f"Unsupported target language: {target!r}. "
            f"Use one of: {supported}."
        )

    return TARGET_ALIASES[normalized]


def select_dtype(dtype_name: str) -> torch.dtype:
    """Select an inference dtype suitable for the current hardware."""
    explicit_dtypes = {
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
        "float32": torch.float32,
    }

    if dtype_name in explicit_dtypes:
        return explicit_dtypes[dtype_name]

    if torch.cuda.is_available():
        supports_bf16 = getattr(
            torch.cuda,
            "is_bf16_supported",
            lambda: False,
        )

        if supports_bf16():
            return torch.bfloat16

        # V100 and other older CUDA GPUs.
        return torch.float16

    return torch.float32


def read_source_text(text: Optional[str]) -> str:
    """Read source text from --text or stdin."""
    source = text if text is not None else sys.stdin.read()

    if not source.strip():
        raise ValueError(
            "Source text is empty. "
            "Pass it through --text or stdin."
        )

    # Do not strip the source itself: leading whitespace and formatting
    # may be meaningful for software and HTML inputs.
    return source


def read_custom_prompt(
    prompt: Optional[str],
    prompt_file: Optional[str],
) -> Optional[str]:
    """Read a custom instruction from CLI or a UTF-8 text file."""
    if prompt is not None:
        custom_prompt = prompt
    elif prompt_file is not None:
        custom_prompt = Path(prompt_file).read_text(encoding="utf-8")
    else:
        return None

    custom_prompt = custom_prompt.strip()

    if not custom_prompt:
        raise ValueError("Custom prompt is empty.")

    return custom_prompt


def build_prompt(
    source_text: str,
    target: str,
    domain: Optional[str] = None,
    custom_prompt: Optional[str] = None,
) -> str:
    """Construct the exact plain-text prompt passed to the model."""
    parts = []

    if custom_prompt is not None:
        parts.append(custom_prompt)
    elif domain is not None:
        parts.append(DOMAIN_PROMPTS[domain])

    parts.append(TRANSLATION_PROMPTS[target])

    # Single newline between instruction blocks mirrors the WMT setup.
    return "\n".join(parts) + "\n" + source_text


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Translate English text using the "
            "WMT26 constrained submission model."
        )
    )

    parser.add_argument(
        "--target",
        required=True,
        metavar="LANG",
        help=(
            "Target language: ru, be, kk, or hy. "
            "Aliases such as rus, bel, kaz, hye and WMT codes "
            "are also accepted."
        ),
    )

    parser.add_argument(
        "--text",
        help=(
            "English source text. "
            "If omitted, the source is read from stdin."
        ),
    )

    instruction_group = parser.add_mutually_exclusive_group()

    instruction_group.add_argument(
        "--domain",
        choices=sorted(DOMAIN_PROMPTS),
        help=(
            "WMT26 domain prompt: "
            "social, speech, news, or software."
        ),
    )

    instruction_group.add_argument(
        "--prompt",
        help=(
            "Custom instruction. It replaces the domain prompt; "
            "the target-language prompt is still appended automatically."
        ),
    )

    instruction_group.add_argument(
        "--prompt-file",
        help=(
            "Path to a UTF-8 file containing a custom instruction. "
            "It replaces the domain prompt."
        ),
    )

    parser.add_argument(
        "--model",
        default=MODEL_ID,
        help="Hugging Face repo ID or local model directory.",
    )

    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=2048,
        help="Maximum number of generated tokens. Default: 2048.",
    )

    parser.add_argument(
        "--dtype",
        choices=("auto", "float16", "bfloat16", "float32"),
        default="auto",
        help=(
            "Model dtype. Auto uses BF16 when supported, "
            "FP16 on older CUDA GPUs, and FP32 on CPU."
        ),
    )

    parser.add_argument(
        "--show-prompt",
        action="store_true",
        help=(
            "Print the final prompt and exit without loading the model."
        ),
    )

    parser.add_argument(
        "--quiet",
        action="store_true",
        help="Suppress progress messages on stderr.",
    )

    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if Version(transformers.__version__) < MIN_TRANSFORMERS_VERSION:
        raise RuntimeError(
            "transformers>={} is required; found {}.\n"
            "Upgrade it with:\n"
            "python -m pip install -U "
            "'transformers>=4.46.3,<5'".format(
                MIN_TRANSFORMERS_VERSION,
                transformers.__version__,
            )
        )

    if args.max_new_tokens <= 0:
        raise ValueError("--max-new-tokens must be positive.")

    target = normalize_target(args.target)
    source_text = read_source_text(args.text)
    custom_prompt = read_custom_prompt(
        prompt=args.prompt,
        prompt_file=args.prompt_file,
    )

    prompt = build_prompt(
        source_text=source_text,
        target=target,
        domain=args.domain,
        custom_prompt=custom_prompt,
    )

    if args.show_prompt:
        print(prompt)
        return

    dtype = select_dtype(args.dtype)

    log("[1/3] Loading tokenizer...", args.quiet)

    tokenizer = AutoTokenizer.from_pretrained(
        args.model,
        use_fast=False,
    )

    log(
        f"[2/3] Loading model with dtype={dtype}...",
        args.quiet,
    )
    started = time.time()

    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype=dtype,
        device_map="auto",
        low_cpu_mem_usage=True,
    )
    model.eval()

    log(
        f"[2/3] Model loaded in {time.time() - started:.1f}s.",
        args.quiet,
    )

    encoded = tokenizer(
        prompt,
        return_tensors="pt",
    )

    input_device = model.get_input_embeddings().weight.device
    inputs = {
        name: tensor.to(input_device)
        for name, tensor in encoded.items()
    }

    input_length = inputs["input_ids"].shape[1]

    pad_token_id = tokenizer.pad_token_id
    if pad_token_id is None:
        pad_token_id = tokenizer.eos_token_id

    log("[3/3] Generating translation...", args.quiet)
    started = time.time()

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,
            num_beams=1,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=pad_token_id,
        )

    generated_ids = output_ids[0, input_length:]

    translation = tokenizer.decode(
        generated_ids,
        skip_special_tokens=True,
    ).strip()

    log(
        f"[3/3] Done in {time.time() - started:.1f}s.",
        args.quiet,
    )

    if not translation:
        raise RuntimeError(
            "The model generated an empty translation."
        )

    # stdout contains only the translation.
    print(translation)


if __name__ == "__main__":
    main()