import re
from wordsegment import load, segment
import tiktoken

load()  # load once
enc = tiktoken.get_encoding("cl100k_base")

def _tiktoken_split(word: str):
    token_ids = enc.encode(word)
    pieces = [enc.decode([tid]) for tid in token_ids]
    merged, buf = [], ""

    for p in pieces:
        if p.isalpha():
            buf += p
        else:
            if buf:
                merged.append(buf)
                buf = ""
    if buf:
        merged.append(buf)

    return merged or [word]

def tokenize_domain(domain: str):
    domain = domain.lower()
    root = re.sub(r"\.[a-z0-9]{2,6}$", "", domain)

    primary = segment(root)

    if len(primary) <= 1:
        return _tiktoken_split(root)

    return primary
