from src.thinking_parser import parse_thinking_trace


def test_parse_xml_thinking_tag():
    raw = "<thinking>\nLet me analyze this problem step by step.\n</thinking>\nHere is the answer."
    segments, response = parse_thinking_trace(raw)
    assert len(segments) == 1
    assert segments[0].content == "Let me analyze this problem step by step."
    assert segments[0].marker == "thinking"
    assert response == "Here is the answer."


def test_parse_xml_thought_tag():
    raw = "<thought>This is my reasoning process</thought>\nFinal response here."
    segments, response = parse_thinking_trace(raw)
    assert len(segments) == 1
    assert segments[0].content == "This is my reasoning process"
    assert segments[0].marker == "thought"
    assert response == "Final response here."


def test_parse_text_thinking_prefix():
    raw = "Thinking:\nThis is a text-based thinking trace.\n\nNow for the actual response."
    segments, response = parse_thinking_trace(raw)
    assert len(segments) == 1
    assert segments[0].content == "This is a text-based thinking trace."
    assert segments[0].marker == "Thinking:"
    assert response == "Now for the actual response."


def test_parse_no_thinking():
    raw = "This is a normal response without any thinking markers."
    segments, response = parse_thinking_trace(raw)
    assert len(segments) == 0
    assert response == raw


def test_parse_empty_response():
    segments, response = parse_thinking_trace("")
    assert len(segments) == 0
    assert response == ""


def test_parse_multiple_markers():
    raw = "<thinking>First thinking</thinking>\n<thought>Second thought</thought>\nResponse"
    segments, response = parse_thinking_trace(raw)
    assert len(segments) == 2
    assert segments[0].content == "First thinking"
    assert segments[1].content == "Second thought"


def test_parse_thinking_with_empty_response():
    raw = "<thinking>Just thinking, no response</thinking>"
    segments, response = parse_thinking_trace(raw)
    assert len(segments) == 1
    assert segments[0].content == "Just thinking, no response"
    assert response == ""


if __name__ == "__main__":
    test_parse_xml_thinking_tag()
    test_parse_xml_thought_tag()
    test_parse_text_thinking_prefix()
    test_parse_no_thinking()
    test_parse_empty_response()
    test_parse_multiple_markers()
    test_parse_thinking_with_empty_response()
    print("All thinking trace tests passed!")