"""Shared prompt construction and model loading for Qwen3.5 recognition and translation."""
import os
import re

import torch
from PIL import Image

LANG_NAME = {"en": "English", "sv": "Swedish", "de": "German"}

OCR_PROMPT = (
    "Transcribe the runic inscription in this photograph. Output only the transliteration "
    "in lowercase Latin letters following the Rundata convention (þ, ð, æ, ø, ą, and R for "
    "the yr rune), with words separated by single spaces."
)


def mt_prompt(src, lang, rl=None):
    """rl: runic-language tag (lang_tags.py); with it the prompt carries <rl> <2lang> tags before the input."""
    tags = f"Tags: <{rl}> <2{lang}>\n" if rl else ""
    return (
        f"Translate this runic inscription, given as a Rundata-style transliteration, into "
        f"{LANG_NAME[lang]}. Output only the translation.\n\n{tags}Transliteration: {src}"
    )


MAX_IMAGE_PIXELS = int(os.environ.get("RUNIC_MAX_PIXELS", 640_000))


def load_image(path):
    im = Image.open(path).convert("RGB")
    w, h = im.size
    if w * h > MAX_IMAGE_PIXELS:
        s = (MAX_IMAGE_PIXELS / (w * h)) ** 0.5
        im = im.resize((max(32, int(w * s)), max(32, int(h * s))), Image.BICUBIC)
    return im


def user_messages(task, text=None, lang=None, rl=None):
    if task == "ocr":
        return [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": OCR_PROMPT}]}]
    return [{"role": "user", "content": [{"type": "text", "text": mt_prompt(text, lang, rl)}]}]


def _prompt_text(processor, task, text=None, lang=None, rl=None):
    msgs = user_messages(task, text, lang, rl)
    if getattr(processor, "text_only", False):
        # text-only chat templates expect the content as a plain string; some (Qwen3) silently render a list as ""
        msgs = [dict(m, content=m["content"][0]["text"]) for m in msgs]
    try:
        return processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    except TypeError:
        if task == "ocr":
            raise
        # chat templates of some text-only models (Mistral) expect the content as a plain string
        msgs = [dict(m, content=m["content"][0]["text"]) for m in msgs]
        return processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)


def prompt_text(processor, task, text=None, lang=None, rl=None):
    out = _prompt_text(processor, task, text, lang, rl)
    if task == "mt" and text and text not in out:  # a chat template that drops the user content
        raise ValueError(f"source text missing from the rendered prompt: {out[:300]!r}")
    return out


def _check_tokenizer(model_path, tok):
    """Some checkpoints (Ministral-8B-Instruct-2410 under transformers 5.x) load as a slow LlamaTokenizer that drops
    spaces when encoding; fall back to the fast tokenizer built from tokenizer.json, keeping the chat template and
    special tokens."""
    probe = "had this stone raised"
    if tok.decode(tok(probe, add_special_tokens=False)["input_ids"]).strip() == probe:
        return tok
    from transformers import PreTrainedTokenizerFast
    f = os.path.join(model_path, "tokenizer.json")
    fast = PreTrainedTokenizerFast(tokenizer_file=f, bos_token=tok.bos_token, eos_token=tok.eos_token,
                                   unk_token=tok.unk_token, pad_token=tok.pad_token)
    fast.chat_template = tok.chat_template
    if fast.decode(fast(probe, add_special_tokens=False)["input_ids"]).strip() != probe:
        raise ValueError(f"tokenizer of {model_path} does not round-trip text")
    print("tokenizer: slow tokenizer drops spaces, using tokenizer.json", flush=True)
    return fast


def load_model(model_path, adapter=None, trainable=False):
    """Vision-language checkpoints load as image-text-to-text models with their processor; text-only checkpoints
    (alternative translators) load as causal LMs, and their tokenizer stands in for the processor."""
    from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor
    from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES

    cfg = AutoConfig.from_pretrained(model_path)
    processor = AutoProcessor.from_pretrained(model_path)
    if cfg.model_type in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES:
        model = AutoModelForImageTextToText.from_pretrained(
            model_path, dtype=torch.bfloat16, attn_implementation="sdpa"
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_path, dtype=torch.bfloat16, attn_implementation="eager" if "gemma" in cfg.model_type else "sdpa"
        )
    if not hasattr(processor, "tokenizer"):
        processor = _check_tokenizer(model_path, processor)
        processor.tokenizer = processor
        processor.text_only = True
        if processor.pad_token is None:  # Mistral tokenizers have no pad token; padded positions are masked anyway
            processor.pad_token = processor.unk_token or processor.eos_token
    if adapter:
        from peft import PeftModel

        model = PeftModel.from_pretrained(model, adapter, is_trainable=trainable)
    return processor, model


def _has_token(tokenizer, t):
    i = tokenizer.convert_tokens_to_ids(t)
    return i is not None and i >= 0 and i != tokenizer.unk_token_id


def end_text(tokenizer):
    """Text that closes an assistant turn in the training target."""
    if _has_token(tokenizer, "<|im_end|>"):
        return "<|im_end|>\n"
    if _has_token(tokenizer, "<end_of_turn>"):
        return "<end_of_turn>\n"
    return tokenizer.eos_token


def encode_kwargs(processor):
    """Text-only chat templates (Gemma) already contain <bos>; do not add a second one."""
    return {"add_special_tokens": False} if getattr(processor, "text_only", False) else {}


def stop_ids(tokenizer):
    ids = {tokenizer.convert_tokens_to_ids("<|im_end|>")}
    if _has_token(tokenizer, "<end_of_turn>"):
        ids.add(tokenizer.convert_tokens_to_ids("<end_of_turn>"))
    if tokenizer.eos_token_id is not None:
        ids.add(tokenizer.eos_token_id)
    return sorted(i for i in ids if i is not None and i >= 0)


_THINK = re.compile(r"<think>.*?</think>", re.S)


def clean_generation(text):
    text = _THINK.sub("", text)
    text = text.replace("<think>", "").replace("</think>", "")
    for tok in ("<|im_end|>", "<|im_start|>", "<|endoftext|>", "<end_of_turn>", "<eos>", "</s>", "<unk>"):
        text = text.split(tok)[0]
    text = re.sub(r"^\s*(assistant|user|model)\s*\n", "", text)
    return text.strip()