import os, sys; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import os, math, types, runpy
import harness, torch
pack, nodes = harness.load_pack()
HERE_ = os.path.dirname(os.path.abspath(__file__))

# ---- script parsing
P = nodes.parse_script
names = ["Jean", "Maria", "", "", "", ""]
assert P("Jean: hi\nMaria: [sad] ok: fine\nand more\njean: bye", names) == [(0, "hi"), (1, "[sad] ok: fine"), (1, "and more"), (0, "bye")]
assert P("Maria:\nline one\nline two", names) == [(1, "line one"), (1, "line two")]
assert P("Жан: Привет\nМария: Да", ["Жан", "Мария"] + [""] * 4) == [(0, "Привет"), (1, "Да")]
for bad, why in [("hi there", "no name first"), ("Jean: a\nMaía: b", "misspelled speaker"), ("Jean: a\nHe said: yes", "unknown plain word before colon")]:
    try: P(bad, names); raise SystemExit("no error for: " + why)
    except ValueError as e: print("rejects", why, "->", str(e)[:70])
assert P("Jean: a\nthe time is 5:30, ok", names) == [(0, "a"), (0, "the time is 5:30, ok")]  # digits/punctuation before a colon: a continuation
print("parse_script ok")

# ---- both modes on the tiny model
ns = runpy.run_path(os.path.join(HERE_, "smoke_core.py"))
ogg = nodes.Ogg(ns["model"], ns["tok"], types.SimpleNamespace())
vA = {"transcript": "<|speaker:0|> hello there", "codes": torch.randint(0, 1024, (10, 20))}
vB = {"transcript": "good morning to you", "codes": torch.randint(0, 1024, (10, 30))}
script = "Jean: [excited] one\nMaria: [sighing] two\nJean: three"
kw = dict(name_1="Jean", name_2="Maria", language="ru", seed=5, steps=2, cfg=1.5, temperature=0.8, max_frames=24)

r = nodes.OggDialogue.execute(ogg, script, mode="line by line", voice_1=vA, voice_2=vB, **kw).args[0]
print("line by line:", tuple(r["codes"].shape), r["takes"], [(l["speaker"], l["text"]) for l in r["lines"]])
assert [l["speaker"] for l in r["lines"]] == [0, 1, 0] and sum(r["takes"]) == r["codes"].shape[1]

r = nodes.OggDialogue.execute(ogg, script, mode="one pass", voice_1=vA, voice_2=vB, **kw).args[0]
print("one pass:", tuple(r["codes"].shape), r["takes"], len(r["lines"]))
assert len(r["takes"]) == 1 and r["codes"].shape[0] == 10

# what the model is given in one pass, in both layouts
seen = {}
orig = nodes.gen.blocks_cached
def spy(*a, **k):
    seen.update(prompt_codes=None if a[2] is None else tuple(a[2].shape), prompt_text=a[3], text=a[4], speaker=k.get("speaker"),
                voices=[(v[0], v[1], tuple(v[2].shape)) for v in k.get("voices") or []])
    return orig(*a, **k)
nodes.gen.blocks_cached = spy
nodes.OggDialogue.execute(ogg, "Jean: a\nJean: b\nMaria: c\nJean: d", mode="one pass", voice_1=vA, voice_2=vB, **kw)
print("as trained:", seen)
assert seen["text"] == "<|quad_start|>S1<|quad_end|>a b <|quad_start|>S2<|quad_end|>c <|quad_start|>S1<|quad_end|>d"
assert seen["speaker"] == "dialogue" and seen["voices"] == [] and seen["prompt_codes"] == (10, 50)
assert seen["prompt_text"] == "<|quad_start|>S1<|quad_end|>hello there <|quad_start|>S2<|quad_end|>good morning to you"
nodes.OggDialogue.execute(ogg, "Maria: c\nJean: a", mode="one pass", dialogue_layout="reference script", voice_1=vA, voice_2=vB, **kw)
print("reference script:", seen)
assert seen["text"] == "<|quad_start|>S1<|quad_end|>c <|quad_start|>S2<|quad_end|>a" and seen["prompt_codes"] is None
assert seen["voices"] == [("S1", "good morning to you", (10, 30)), ("S2", "hello there", (10, 20))]  # S1 = whoever speaks first
nodes.gen.blocks_cached = orig

# a speaker with no voice speaks in a trained identity; one pass still needs every voice
seen = {}
orig = nodes.gen.blocks_cached
def spy(*a, **k): seen.setdefault("calls", []).append((a[2] is None, k.get("speaker"))); return orig(*a, **k)
nodes.gen.blocks_cached = spy
r = nodes.OggDialogue.execute(ogg, script, mode="line by line", voice_1=vA, **kw).args[0]
nodes.gen.blocks_cached = orig
print("no voice for Maria:", seen["calls"])
assert [c[0] for c in seen["calls"]] == [False, True, False] and seen["calls"][1][1] == nodes.resolve_speaker("", 0, "Maria") and seen["calls"][0][1] is None
assert r["takes"] == [32, 32, 32]
try: nodes.OggDialogue.execute(ogg, script, mode="one pass", voice_1=vA, **kw); raise SystemExit("no error")
except ValueError as e: print("one pass without a voice ->", str(e)[:90])
# a typed token, a bad token
nodes.gen.blocks_cached = spy; seen.clear()
nodes.OggDialogue.execute(ogg, "Maria: x", mode="line by line", speaker_2="ru_x", **kw)
nodes.gen.blocks_cached = orig; assert seen["calls"] == [(True, "ru_x")]
try: nodes.OggDialogue.execute(ogg, "Maria: x", mode="line by line", speaker_2="pl_3", **kw); raise SystemExit("no error")
except ValueError as e: print("bad token ->", str(e)[:70])
# a defined name that the script does not use needs no voice
r = nodes.OggDialogue.execute(ogg, "Jean: only me", mode="line by line", voice_1=vA, **kw).args[0]
print("unused speaker needs no voice ok:", r["takes"])

# ---- speaker tokens, quality tokens, voice from take
R = nodes.resolve_speaker
assert R("") is None and R("  ") is None
assert R("", 1, "Maria") == R("", 99, " maria ") and R("", 1, "Maria") != R("", 1, "Jean")   # fixed by the name, not the seed
assert R("random", 7) == R("random", 7) and R("random", 7) != R("random", 8) and R("random", 7).startswith("spk_")
assert R("spk_5") == "spk_5" and R("ru_x") == "ru_x"
for bad in ("pl_3", "q5", "spk_99999", "foo"):
    try: R(bad); raise SystemExit("no error for " + bad)
    except ValueError: pass
assert nodes.check_quality("q35") == "q35" and nodes.check_quality("") is None
try: nodes.check_quality("spk_1"); raise SystemExit("no error")
except ValueError: pass
print("speaker and quality tokens ok:", R("random", 7), R("", 0, "Jean"))

# Generate passes the resolved token on
nodes.gen.blocks_cached = spy; seen.clear()
nodes.OggGenerate.execute(ogg, "[sad] hi", "en", 4, 2, 1.5, 0.8, True, 16, "random", "q7", 0.0, 0, 4, None)
nodes.gen.blocks_cached = orig; assert seen["calls"] == [(True, R("random", 4))], seen
print("generate with random speaker ok:", seen["calls"])

take = {"codes": torch.arange(10 * 60).reshape(10, 60) % 1000, "lines": [{"speaker": 0, "text": "[sad]  hello [pause] there", "chunk": 0, "start": 0, "end": 25},
        {"speaker": 0, "text": "second", "chunk": 1, "start": 25, "end": 60}], "takes": [25, 35]}
v = nodes.OggVoiceFromTake.execute(take, 0, True).args[0]
assert v["transcript"] == "hello there" and tuple(v["codes"].shape) == (10, 25) and torch.equal(v["codes"], take["codes"][:, :25])
v = nodes.OggVoiceFromTake.execute(take, 1, False).args[0]; assert v["transcript"] == "second" and v["codes"].shape[1] == 35
assert nodes.OggVoiceFromTake.execute(take, 0, False).args[0]["transcript"] == "[sad]  hello [pause] there"
try: nodes.OggVoiceFromTake.execute(take, 2, True); raise SystemExit("no error")
except ValueError as e: print("take out of range ->", e)
print("voice from take ok")
print("DIALOGUE OK")