"""Word-level tokenizer. Unknown words map to id 0 (<pad>), which the
embedding layer ignores, so garbage prompts are harmless."""

import re

import torch

from .palette import COLOR_WORDS
from .templates import CLASSES

CLASS_SYNONYMS = {
    "potion": ["potion", "bottle", "flask", "elixir", "vial"],
    "sword": ["sword", "blade", "dagger"],
    "shield": ["shield", "buckler"],
    "ship": ["ship", "rocket", "spaceship", "starship"],
    "invader": ["invader", "alien", "monster"],
    "slime": ["slime", "blob", "jelly"],
    "skull": ["skull", "bone", "skeleton"],
    "tree": ["tree", "pine", "forest"],
    "key": ["key"],
    "heart": ["heart", "love", "life"],
    "star": ["star", "sparkle"],
    "mushroom": ["mushroom", "fungus", "toadstool", "shroom"],
    "ghost": ["ghost", "spirit", "phantom", "spooky"],
    "bomb": ["bomb", "grenade", "explosive"],
    "coin": ["coin", "token", "money", "medal"],
    "gem": ["gem", "diamond", "jewel", "crystal"],
    "cherry": ["cherry", "cherries", "berry", "fruit"],
    "crown": ["crown"],
    "bat": ["bat", "vampire"],
    "cactus": ["cactus", "desert", "succulent"],
    "robot": ["robot", "droid", "android", "mech"],
}

PAD = "<pad>"

# Accessory words select a second, overlaid part (see parts.py).
ACCESSORY_SYNONYMS = {
    "none": [],
    "crown": ["crown", "king", "queen", "royal", "crowned", "prince", "princess"],
    "halo": ["halo", "angel", "angelic", "holy", "saint"],
    "horns": ["horns", "horned", "devil", "demon", "evil"],
    "helmet": ["helmet", "knight", "armor", "armour", "armored"],
    "antenna": ["antenna", "antennae", "bug", "martian"],
    "wings": ["wings", "winged", "fly", "flying", "bird", "fairy", "butterfly"],
    "torch": ["torch", "torches", "fire", "flame", "burning"],
}
SIZE_WORDS = {"small": 1, "tiny": 1, "mini": 1, "little": 1}
# These split a prompt in two: words after them describe the accessory.
SEPARATORS = ["with", "wearing", "holding", "and", "plus"]


def build_vocab():
    vocab = {PAD: 0}
    groups = [
        set(w for ws in CLASS_SYNONYMS.values() for w in ws),
        set(COLOR_WORDS),
        set(w for ws in ACCESSORY_SYNONYMS.values() for w in ws),
        set(SIZE_WORDS),
        set(SEPARATORS),
    ]
    for g in groups:
        for word in sorted(g):
            if word not in vocab:
                vocab[word] = len(vocab)
    return vocab


VOCAB = build_vocab()

_SPLIT = re.compile(r"[^a-z]+")


def tokenize(prompt, vocab=VOCAB):
    words = [w for w in _SPLIT.split(prompt.lower()) if w]
    ids = [vocab.get(w, 0) for w in words]
    return torch.tensor(ids if ids else [0], dtype=torch.long)


def pad_batch(prompts, vocab=VOCAB):
    toks = [tokenize(p, vocab) for p in prompts]
    L = max(t.numel() for t in toks)
    out = torch.zeros(len(toks), L, dtype=torch.long)
    for i, t in enumerate(toks):
        out[i, : t.numel()] = t
    return out


def split_prompt(prompt):
    """'green slime with gold crown' -> ('green slime', 'gold crown').
    No separator -> whole prompt is the body half, accessory half empty."""
    words = [w for w in _SPLIT.split(prompt.lower()) if w]
    for i, w in enumerate(words):
        if w in SEPARATORS:
            return " ".join(words[:i]), " ".join(words[i + 1:])
    return " ".join(words), ""


def pad_batch_split(prompts, vocab=VOCAB):
    halves = [split_prompt(p) for p in prompts]
    return pad_batch([a for a, _ in halves], vocab), pad_batch([b for _, b in halves], vocab)


def word_class(word):
    for cls, ws in CLASS_SYNONYMS.items():
        if word in ws:
            return CLASSES.index(cls)
    return None