import os as _os
HERE = _os.path.dirname(_os.path.abspath(__file__))
ROOT = _os.path.dirname(HERE)
import sys, os, json, types, math, shutil; sys.path.insert(0, HERE)
import harness, torch
from safetensors.torch import save_file
pack, nodes = harness.load_pack()
shutil.rmtree(harness.OUT, ignore_errors=True); os.makedirs(harness.OUT + "/temp")

# ---- 1. the workflows match the schemas of the nodes (names, order, widget or socket)
import asyncio
async def _nodes():
    return await (await pack.comfy_entrypoint()).get_node_list()
ext = asyncio.run(_nodes())
by_type = {c.__name__: c for c in ext}
node_ids = {}
for c in ext:
    sch = c.define_schema(); node_ids[sch.k["node_id"]] = (c, sch)
print("node ids:", sorted(node_ids))
WIDGET = {"COMBO", "INT", "FLOAT", "STRING", "BOOLEAN"}
for name in ("ogg_text_to_speech", "ogg_voice_clone", "ogg_dialogue", "ogg_voice_design"):
    w = json.load(open(f"{ROOT}/example_workflows/{name}.json"))
    for n in w["nodes"]:
        if n["type"] in ("MarkdownNote", "SaveAudio", "LoadAudio"): continue
        cls, sch = node_ids[n["type"]]
        want = [(i.a[0], i.kind in WIDGET) for i in sch.k["inputs"]]
        have = [(i["name"], "widget" in i) for i in n["inputs"]]
        assert want == have, (n["type"], want, have)
        nw = sum(1 for _, wd in want if wd) + sum(1 for i in sch.k["inputs"] if i.k.get("control_after_generate"))
        if n["type"] != "OggEditor": assert len(n["widgets_values"]) == nw, (n["type"], len(n["widgets_values"]), nw)
    print("workflow OK:", name)

# ---- 2. load_ogg path on a tiny checkpoint (tokenizer files are not available offline: from_config_dir is replaced)
sys.path.insert(0, HERE)
import runpy
ns = runpy.run_path(HERE + "/smoke_core.py")
tiny, tok = ns["model"], ns["tok"]
tmp = harness.OUT + "/tiny_ogg.safetensors"
save_file({k: v.contiguous() for k, v in tiny.state_dict().items()}, tmp)
nodes.ensure_omnivoice_files = lambda: None
seen = {}
def fake_from_config_dir(path, dtype=torch.float32, attn="sdpa", prompt_x=False):
    seen["prompt_x"], seen["path"] = prompt_x, path
    from transformers import AutoModel
    m = nodes.of.OmniFish(AutoModel.from_config(tiny.llm.config), prompt_x=prompt_x).to(dtype)
    t = ns["StubTok"](); m.configure_languages(t); return m, t
nodes.of.OmniFish.from_config_dir = staticmethod(fake_from_config_dir)
ogg = nodes.load_ogg(tmp, torch.float32)
print("load_ogg ok; prompt_x:", seen["prompt_x"], "tokenizer wrapped:", type(ogg.tok).__name__, "| omnivoice dir:", nodes.OMNIVOICE_DIR)
try:
    save_file({"x": torch.zeros(1)}, harness.OUT + "/bad.safetensors"); nodes.load_ogg(harness.OUT + "/bad.safetensors", torch.float32)
except ValueError as e: print("rejects a foreign file:", e)

# ---- 3. a stand-in codec with the real interface: Encode Voice, Generate, Decode, Editor, save/load
class FakeCodecModel:
    def encode(self, wav):
        T = math.ceil(wav.shape[-1] / 2048); torch.manual_seed(1)
        c = torch.randint(0, 1024, (1, 10, T)); c[:, 0] %= 4096; return c
    def decode(self, codes, dtype):
        n = codes.shape[-1] * 2048; t = torch.arange(n) / 44100
        return (0.2 * torch.sin(2 * math.pi * 300 * t) * (torch.arange(n) % 20000 < 15000))[None]
codec = types.SimpleNamespace(model=FakeCodecModel(), patcher=harness.ModelPatcher(None, torch.device("cpu"), torch.device("cpu")))
audio = {"waveform": torch.randn(1, 1, 48000) * 0.1, "sample_rate": 48000}
voice = nodes.OggEncodeVoice.execute(codec, audio, "<|speaker:0|> hello there").args[0]
print("encoded voice codes:", tuple(voice["codes"].shape))
nodes.OggSaveVoice.execute(voice, "ogg/voice")
vf = [f for f in os.listdir(harness.OUT + "/ogg_voices/ogg")][0]
loaded = nodes.OggLoadVoice.execute("ogg/" + vf) if False else None
nodes.folder_paths.get_full_path_or_raise = lambda n, f: os.path.join(harness.OUT, n, f)
loaded = nodes.OggLoadVoice.execute("ogg/" + vf).args[0]
assert torch.equal(loaded["codes"], voice["codes"]) and loaded["transcript"] == voice["transcript"]; print("voice save/load round-trip ok")

r = nodes.OggGenerate.execute(ogg, "[sad] one\ntwo\nthree", "en", 3, 2, 1.5, 0.8, True, 24, "", "", 0.0, 0, 4, loaded).args[0]
print("generate: codes", tuple(r["codes"].shape), "takes", r["takes"], "lines", [(l["start"], l["end"]) for l in r["lines"]])
nodes.OggSaveCodes.execute(r, "ogg/take")
cf = os.listdir(harness.OUT + "/ogg_codes/ogg")[0]
r2 = nodes.OggLoadCodes.execute("ogg/" + cf).args[0]
assert torch.equal(r2["codes"], r["codes"]) and r2["takes"] == r["takes"] and r2["lines"] == r["lines"]; print("codes save/load round-trip ok")

a = nodes.OggDecode.execute(codec, r2, False).args[0]
expect = sum(r["takes"]) * 2048 + 2 * round(0.3 * 44100)
print("decode:", tuple(a["waveform"].shape), a["sample_rate"], "expected samples", expect); assert a["waveform"].shape[-1] == expect
a2 = nodes.OggDecode.execute(codec, r2, True).args[0]; print("decode with shorten_long_pauses:", tuple(a2["waveform"].shape))
e = nodes.OggEditor.execute(codec, r2, None, "")
print("editor audio:", tuple(e.args[0]["waveform"].shape), "ui:", e.ui["ogg_editor"][0])
assert os.path.exists(os.path.join(harness.OUT, "temp", "ogg_editor", e.ui["ogg_editor"][0]["audio"]))
a_ = json.load(open(os.path.join(harness.OUT, "temp", "ogg_editor", e.ui["ogg_editor"][0]["analysis"])))
print("editor analysis: lines", len(a_["lines"]), "speakers", a_["speakers"], "duration", round(a_["duration"], 2))

# ---- 4. the editor's server route logic
from ogg_pack.ogg.editor import DEFAULTS, plan
p = plan(a_, dict(DEFAULTS)); print("plan pieces:", len(p["pieces"]))

# ---- 5. error path of the omnivoice fetch (no network here)
import importlib; importlib.reload(sys.modules["ogg_pack.nodes"]) if False else None
src = open(ROOT + "/nodes.py").read()
ns2 = {}
print("ALL OK")