import os as _os
HERE = _os.path.dirname(_os.path.abspath(__file__))
ROOT = _os.path.dirname(HERE)
import re, sys, time, types, hashlib
sys.path.insert(0, ROOT)
import torch
from transformers import AutoModel, Qwen3Config
from ogg import model as of, generate as gen

TEXT_VOCAB = 1000

class Out:
    def __init__(self, ids, tensor): self.input_ids = torch.tensor([ids]) if tensor else ids

class StubTok:
    """stands in for OmniVoice's tokenizer: words and <|special|> markers hashed into the text vocabulary"""
    def _ids(self, s):
        parts = re.findall(r"<\|[^|]+\|>|\w+|[^\w\s]", s)
        return [1 + int(hashlib.md5(p.encode()).hexdigest(), 16) % (TEXT_VOCAB - 2) for p in parts]
    def __call__(self, s, return_tensors=None, add_special_tokens=True, **kw):
        return Out(self._ids(s), return_tensors == "pt")
    def convert_tokens_to_ids(self, t): return self._ids(t)[0]

cfg = Qwen3Config(vocab_size=TEXT_VOCAB, hidden_size=64, intermediate_size=128, num_hidden_layers=2, num_attention_heads=4,
                  num_key_value_heads=2, head_dim=16, max_position_embeddings=2048)
cfg._attn_implementation = "sdpa"
torch.manual_seed(0)
llm = AutoModel.from_config(cfg)
model = of.OmniFish(llm).eval()
tok = StubTok()
model.configure_languages(tok)
tok = of.IdTokenizer(tok, model.text_vocab)
print("params M:", sum(p.numel() for p in model.parameters()) / 1e6)

# the state dict round-trips through load_weights (the path the loader takes)
state = {k: v.clone() for k, v in model.state_dict().items()}
fresh = model.load_weights(state, tok)
print("load_weights fresh (False expected, tables present):", fresh)

# a model with the prompt cross-attention keys loads when built with prompt_x
mx = of.OmniFish(AutoModel.from_config(cfg), prompt_x=True)
print("px keys:", sum(".px." in k for k in mx.state_dict()))

prompt = torch.randint(0, 1024, (10, 24)); prompt[0] %= 4096
for label, kw in [("no prompt", dict(prompt_codes=None, prompt_text="")),
                  ("with prompt", dict(prompt_codes=prompt, prompt_text="hello there")),
                  ("left+after ctx", dict(prompt_codes=prompt, prompt_text="hello there", semantic_left=("one", "two"), after=("three",)))]:
    t = time.time()
    last = None
    for clean, codes in gen.blocks_cached(model, tok, text="[sighing] hello world", lang="en", block=16, steps=2, cfg=1.5, max_frames=32, seed=1, **kw):
        last = (clean, codes)
    clean, codes = last
    n = gen.speech_frames(clean)
    print(f"{label}: clean {tuple(clean.shape)} codes {tuple(codes.shape)} speech {n} range ok:",
          int(codes[0, :n].max()) < 4096 if n else None, int(codes[1:, :n].max()) < 1024 if n else None, f"{time.time()-t:.1f}s")

# cfg=0 path (single batch row)
for clean, codes in gen.blocks_cached(model, tok, None, "", "hi", "pl", block=16, steps=2, cfg=0, max_frames=16, seed=2):
    pass
print("cfg=0 ok", tuple(codes.shape))