from llmux.harmony import extract_final_text, HarmonyStreamFilter


def test_extract_llamacpp_format():
    text = '<|channel|>analysis<|message|>User greeting. Simple.<|end|><|start|>assistant<|channel|>final<|message|>Hello! How can I help you today?'
    assert extract_final_text(text) == "Hello! How can I help you today?"


def test_extract_llamacpp_with_end_tag():
    text = '<|channel|>analysis<|message|>thinking...<|end|><|start|>assistant<|channel|>final<|message|>The answer is 42.<|end|>'
    assert extract_final_text(text) == "The answer is 42."


def test_extract_transformers_format():
    text = 'analysisUser greeting. Just respond friendly.assistantfinalHello! I am doing great.'
    assert extract_final_text(text) == "Hello! I am doing great."


def test_extract_non_harmony_passthrough():
    text = "Hello! I'm doing well, thanks for asking."
    assert extract_final_text(text) == text


def test_stream_filter_llamacpp():
    f = HarmonyStreamFilter()
    chunks = [
        "<|channel|>", "analysis", "<|message|>", "User ", "greeting.",
        "<|end|>", "<|start|>", "assistant", "<|channel|>", "final",
        "<|message|>", "Hello!", " How ", "are you?"
    ]
    output = ""
    for c in chunks:
        output += f.feed(c)
    output += f.flush()
    assert output == "Hello! How are you?"


def test_stream_filter_transformers():
    f = HarmonyStreamFilter()
    chunks = ["analysis", "User ", "greeting.", "assistant", "final", "Hello!", " Great day!"]
    output = ""
    for c in chunks:
        output += f.feed(c)
    output += f.flush()
    assert output == "Hello! Great day!"


def test_stream_filter_non_harmony():
    f = HarmonyStreamFilter()
    chunks = ["Hello", " world", "!"]
    output = ""
    for c in chunks:
        output += f.feed(c)
    output += f.flush()
    assert output == "Hello world!"