"""Motor de TokenScope (lado Space).

Mode 1: pide la distribucion del siguiente token a un endpoint de Modal (Qwen).
Mode 2: pide la similitud semantica entre dos palabras a otro endpoint (Nemotron
        Embed VL). Los modelos corren en Modal; el Space queda sin GPU.

STUB local para probar sin red ni GPU:  TOKENSCOPE_STUB=1
"""
import json
import os
import re
import unicodedata
import urllib.request

STUB = os.environ.get("TOKENSCOPE_STUB") == "1"

LOGITS_URL = os.environ.get(
    "TOKENSCOPE_LOGITS_URL",
    "https://sanightblack--tokenscope-logits-logits-topk.modal.run",
)
EMBED_URL = os.environ.get(
    "TOKENSCOPE_EMBED_URL",
    "https://sanightblack--tokenscope-embed-embed-similarity.modal.run",
)

SEARCH_K = 200
SHOW_K = 10


def normalize(word: str) -> str:
    word = word.strip().lower()
    word = "".join(
        c for c in unicodedata.normalize("NFD", word)
        if unicodedata.category(c) != "Mn"
    )
    return re.sub(r"[^\w]", "", word)


def first_word(text: str) -> str:
    parts = text.strip().split()
    return parts[0] if parts else ""


def _post(url: str, payload: dict, timeout: int = 180) -> dict:
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as res:
        return json.loads(res.read().decode("utf-8"))


# --------------------------------------------------------------------------- #
# Mode 1 — distribucion del siguiente token                                   #
# --------------------------------------------------------------------------- #
def _dist_stub(context: str, k: int):
    import hashlib
    import math

    base = ["the", "a", "of", "and", "to", "in", "is", "with", "for", "on",
            "it", "you", "that", "be", "as", "more", "home", "water", "sun", "eyes"]
    seed = int(hashlib.md5(context.encode()).hexdigest(), 16)
    rng = [(seed >> (i * 7)) % 1000 for i in range(len(base))]
    weighted = sorted(zip(base, rng), key=lambda x: -x[1])[:k]
    total = sum(math.exp(w / 200) for _, w in weighted)
    return [(" " + tok, math.exp(w / 200) / total) for tok, w in weighted]


def next_token_dist(context: str, k: int = SHOW_K):
    if STUB:
        return _dist_stub(context, k)
    return [(tok, float(p)) for tok, p in _post(LOGITS_URL, {"context": context, "k": k})["dist"]]


def analyze(context: str, player_word: str):
    dist = next_token_dist(context, SEARCH_K)
    target = normalize(first_word(player_word))
    rank, prob = None, 0.0
    if target:
        for r, (tok, p) in enumerate(dist, start=1):
            tn = normalize(tok)
            if not tn:
                continue
            if tn == target or target.startswith(tn):
                rank, prob = r, p
                break
    return dist[:SHOW_K], rank, prob


def score_for_rank(rank):
    if rank is None:
        return 0
    if rank == 1:
        return 100
    if rank <= 5:
        return 50
    if rank <= 10:
        return 25
    if rank <= 50:
        return 10
    return 5


# --------------------------------------------------------------------------- #
# Mode 2 — similitud semantica (embeddings)                                   #
# --------------------------------------------------------------------------- #
def _sim_stub(guess: str, target: str) -> float:
    a, b = normalize(guess), normalize(target)
    if a == b:
        return 1.0
    common = len(set(a) & set(b))
    return min(0.92, 0.15 + 0.09 * common)


def embed_similarity(guess: str, target: str) -> float:
    if STUB:
        return _sim_stub(guess, target)
    return float(_post(EMBED_URL, {"guess": guess, "target": target})["similarity"])