# SPDX-License-Identifier: Apache-2.0
"""
Tests for HarmonyStreamingParser.

Tests the streaming parser for gpt-oss models using Harmony format.
Uses real openai_harmony encoding for accurate testing.
"""

import pytest

from omlx.adapter.harmony import (
    HarmonyStreamingParser,
    load_harmony_gpt_oss_encoding,
    preprocess_harmony_messages,
)


class RealTokenizer:
    """Tokenizer wrapper using real openai_harmony encoding."""

    def __init__(self, enc):
        self._enc = enc

    def convert_tokens_to_ids(self, token: str) -> int:
        """Convert token string to ID using harmony encoding."""
        try:
            ids = self._enc.encode(token, allowed_special="all")
            return ids[0] if ids else -1
        except Exception:
            return -1


@pytest.fixture(scope="session")
def harmony_encoding():
    """Get the real openai_harmony encoding."""
    return load_harmony_gpt_oss_encoding()


@pytest.fixture
def tokenizer(harmony_encoding):
    """Create a tokenizer that uses real openai_harmony encoding."""
    return RealTokenizer(harmony_encoding)


@pytest.fixture
def parser(tokenizer):
    """Create a HarmonyStreamingParser with real tokenizer."""
    return HarmonyStreamingParser(tokenizer)


class TestHarmonyStreamingParser:
    """Tests for HarmonyStreamingParser."""

    def test_get_stop_token_ids(self, parser):
        """Test stop token ID extraction."""
        stop_ids = parser.get_stop_token_ids()
        # Real openai_harmony stop tokens for assistant actions
        assert 200002 in stop_ids  # <|return|>
        assert 200012 in stop_ids  # <|call|>
        assert len(stop_ids) == 2

    def test_final_channel_returns_token_ids(self, parser, harmony_encoding):
        """Test that final channel returns token IDs for both stream and visible."""
        # Encode a real harmony message sequence
        # Note: Model generates from <|channel|> since <|start|>assistant is in prompt
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>Hello world<|end|>",
            allowed_special="all",
        )
        # tokens: [200005, 17196, 200008, 13225, 2375, 200007]
        #         channel  final message Hello   world  end

        results = []
        for token in tokens:
            result = parser.process_token(token)
            results.append(result)

        # Find content tokens (after <|message|>)
        message_idx = tokens.index(200008)  # <|message|>
        hello_token = tokens[message_idx + 1]  # Hello
        world_token = tokens[message_idx + 2]  # world

        # Content tokens should return both stream and visible
        hello_result = results[message_idx + 1]
        assert hello_result[1] == hello_token  # stream_token
        assert hello_result[2] == hello_token  # visible_token
        assert hello_result[3] is False  # not stop

        world_result = results[message_idx + 2]
        assert world_result[1] == world_token  # stream_token
        assert world_result[2] == world_token  # visible_token

    def test_analysis_channel_wrapped_in_think(self, parser, harmony_encoding):
        """Test that analysis channel outputs <think> control and stream tokens only."""
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>thinking content<|end|>",
            allowed_special="all",
        )

        results = []
        for token in tokens:
            result = parser.process_token(token)
            results.append(result)

        # Collect all control text and stream/visible tokens
        control_parts = [r[0] for r in results if r[0]]
        stream_tokens = [r[1] for r in results if r[1] is not None]
        visible_tokens = [r[2] for r in results if r[2] is not None]

        # Should have <think> opening and </think> closing
        full_control = "".join(control_parts)
        assert "<think>" in full_control
        assert "</think>" in full_control

        # Analysis content should stream but not be visible
        assert len(stream_tokens) > 0  # Has stream tokens
        assert len(visible_tokens) == 0  # No visible tokens

    def test_return_token_is_stop(self, parser, harmony_encoding):
        """Test that <|return|> signals stop."""
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>Hello<|return|>",
            allowed_special="all",
        )

        results = []
        for token in tokens:
            result = parser.process_token(token)
            results.append(result)

        # Last token (<|return|>) should signal stop
        assert results[-1][3] is True

    def test_finalize_closes_think_tag(self, parser, harmony_encoding):
        """Test that finalize() closes open think tags."""
        # Enter analysis channel without <|end|>
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>thinking",
            allowed_special="all",
        )
        for token in tokens:
            parser.process_token(token)

        # Parser should be in think tag
        assert parser._in_think_tag is True

        # Finalize should close it
        final = parser.finalize()
        assert final == "</think>\n"
        assert parser._in_think_tag is False

    def test_finalize_no_output_when_no_open_tag(self, parser, harmony_encoding):
        """Test that finalize() returns empty when no open tags."""
        # Process a complete final channel message
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>Hello<|return|>",
            allowed_special="all",
        )
        for token in tokens:
            parser.process_token(token)

        final = parser.finalize()
        assert final == ""

    def test_reset_clears_state(self, parser, harmony_encoding):
        """Test that reset() clears all parser state."""
        # Process some tokens to set state
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>thinking",
            allowed_special="all",
        )
        for token in tokens:
            parser.process_token(token)

        # Verify state is set
        assert parser._in_think_tag is True
        assert parser.current_channel == "analysis"

        # Reset
        parser.reset()

        # Verify state is cleared
        assert parser._in_think_tag is False
        assert parser._prev_channel is None

    def test_multiple_messages_in_sequence(self, parser, harmony_encoding):
        """Test handling multiple messages in sequence."""
        # analysis -> final sequence
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>thinking<|end|>"
            "<|start|>assistant<|channel|>final<|message|>Hello<|return|>",
            allowed_special="all",
        )

        control_parts = []
        stream_tokens = []
        visible_tokens = []

        for token in tokens:
            control, stream_token, visible_token, _ = parser.process_token(token)
            if control:
                control_parts.append(control)
            if stream_token is not None:
                stream_tokens.append(stream_token)
            if visible_token is not None:
                visible_tokens.append(visible_token)

        # Should have think tags
        full_control = "".join(control_parts)
        assert "<think>" in full_control
        assert "</think>" in full_control

        # Stream should have both analysis and final tokens
        assert len(stream_tokens) > 0

        # Visible should only have final channel tokens
        # Analysis tokens should not be in visible
        assert len(visible_tokens) > 0

    def test_channel_transition_closes_think_tag(self, parser, harmony_encoding):
        """Test that transitioning from analysis to final closes think tag."""
        # First: enter analysis channel
        tokens1 = harmony_encoding.encode(
            "<|channel|>analysis<|message|>thinking",
            allowed_special="all",
        )
        for token in tokens1:
            parser.process_token(token)

        assert parser._in_think_tag is True

        # Second: transition to final channel (via new message)
        tokens2 = harmony_encoding.encode(
            "<|end|><|start|>assistant<|channel|>final<|message|>",
            allowed_special="all",
        )

        control_parts = []
        for token in tokens2:
            control, _, _, _ = parser.process_token(token)
            if control:
                control_parts.append(control)

        # Should have closing think tag
        full_control = "".join(control_parts)
        assert "</think>" in full_control
        assert parser._in_think_tag is False


class TestHarmonyToolCalling:
    """Tests for Harmony tool call extraction."""

    def test_tool_call_extraction(self, tokenizer, harmony_encoding):
        """Test tool call extraction from commentary channel."""
        parser = HarmonyStreamingParser(tokenizer)

        # Simulate tool call message format
        # commentary channel with recipient functions.get_weather
        tokens = harmony_encoding.encode(
            '<|channel|>commentary to=functions.get_weather<|message|>{"location": "Seoul"}<|call|>',
            allowed_special="all",
        )

        for token in tokens:
            _, _, _, is_stop = parser.process_token(token)

        # Last token should be stop
        assert is_stop is True

        tool_calls = parser.get_tool_calls()
        assert len(tool_calls) == 1
        assert tool_calls[0]["name"] == "get_weather"
        assert "Seoul" in tool_calls[0]["arguments"]

    def test_no_tool_calls_for_final_channel(self, parser, harmony_encoding):
        """Test that final channel doesn't produce tool calls."""
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>Hello<|return|>",
            allowed_special="all",
        )
        for token in tokens:
            parser.process_token(token)

        tool_calls = parser.get_tool_calls()
        assert len(tool_calls) == 0


class TestHarmonyEdgeCases:
    """Edge case tests for HarmonyStreamingParser."""

    def test_empty_content(self, parser, harmony_encoding):
        """Test handling of empty message content."""
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|><|end|>",
            allowed_special="all",
        )

        stream_tokens = []
        for token in tokens:
            _, stream_token, _, _ = parser.process_token(token)
            if stream_token is not None:
                stream_tokens.append(stream_token)

        # Should have no content tokens
        assert stream_tokens == []

    def test_unknown_channel_without_recipient_falls_back_to_final(
        self, parser, harmony_encoding
    ):
        """Test handling of malformed channel names."""
        # Note: openai_harmony may handle unknown channels differently
        tokens = harmony_encoding.encode(
            "<|channel|>someunknownchannel<|message|>content<|end|>",
            allowed_special="all",
        )

        stream_tokens = []
        visible_tokens = []
        for token in tokens:
            _, stream_token, visible_token, _ = parser.process_token(token)
            if stream_token is not None:
                stream_tokens.append(stream_token)
            if visible_token is not None:
                visible_tokens.append(visible_token)

        assert "content" in harmony_encoding.decode(stream_tokens)
        assert "content" in harmony_encoding.decode(visible_tokens)

    def test_consecutive_analysis_channels(self, parser, harmony_encoding):
        """Test multiple consecutive analysis channels."""
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>first<|end|>"
            "<|start|>assistant<|channel|>analysis<|message|>second<|end|>",
            allowed_special="all",
        )

        control_parts = []
        for token in tokens:
            control, _, _, _ = parser.process_token(token)
            if control:
                control_parts.append(control)

        # Should have proper think tag wrapping for both
        assert control_parts.count("<think>\n") == 2
        assert control_parts.count("</think>\n") == 2


class TestHarmonyIntegration:
    """Integration-style tests for realistic Harmony sequences."""

    def test_typical_response_sequence(self, parser, harmony_encoding):
        """Test a typical Harmony response: analysis -> final."""
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>Let me think...<|end|>"
            "<|start|>assistant<|channel|>final<|message|>Hello world<|return|>",
            allowed_special="all",
        )

        control_output = []
        stream_tokens = []
        visible_tokens = []

        for token in tokens:
            control, stream_token, visible_token, _ = parser.process_token(token)
            if control:
                control_output.append(control)
            if stream_token is not None:
                stream_tokens.append(stream_token)
            if visible_token is not None:
                visible_tokens.append(visible_token)

        # Control should have think tags
        full_control = "".join(control_output)
        assert "<think>" in full_control
        assert "</think>" in full_control

        # Stream should have analysis + final tokens
        assert len(stream_tokens) > 0

        # Visible should only have final tokens (analysis excluded)
        assert len(visible_tokens) > 0
        assert len(visible_tokens) < len(stream_tokens)

    def test_accumulate_visible_tokens(self, parser, harmony_encoding):
        """Test that visible tokens accumulate correctly for storage."""
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>Hello world<|return|>",
            allowed_special="all",
        )

        visible_tokens = []
        for token in tokens:
            _, _, visible_token, _ = parser.process_token(token)
            if visible_token is not None:
                visible_tokens.append(visible_token)

        # Decode accumulated tokens
        output_text = harmony_encoding.decode(visible_tokens)
        assert "Hello" in output_text
        assert "world" in output_text


class TestHarmonyPreprocess:
    """Tests for preprocess_harmony_messages."""

    def test_strips_think_tags_from_assistant(self):
        """Test that <think> tags are stripped from assistant messages."""
        messages = [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "<think>\nLet me think...\n</think>\nHi there!"},
            {"role": "user", "content": "How are you?"},
        ]

        result = preprocess_harmony_messages(messages)

        assert result[0]["content"] == "Hello"  # user unchanged
        assert result[1]["content"] == "Hi there!"  # think stripped
        assert result[2]["content"] == "How are you?"  # user unchanged

    def test_preserves_user_messages(self):
        """Test that user messages are not modified."""
        messages = [
            {"role": "user", "content": "Tell me about <think> tags"},
        ]

        result = preprocess_harmony_messages(messages)
        assert result[0]["content"] == "Tell me about <think> tags"

    def test_handles_empty_content(self):
        """Test handling of empty or None content."""
        messages = [
            {"role": "assistant", "content": ""},
            {"role": "assistant", "content": None},
            {"role": "assistant"},
        ]

        result = preprocess_harmony_messages(messages)
        assert len(result) == 3

    def test_handles_no_think_tags(self):
        """Test that messages without think tags are unchanged."""
        messages = [
            {"role": "assistant", "content": "Just a normal response"},
        ]

        result = preprocess_harmony_messages(messages)
        assert result[0]["content"] == "Just a normal response"

    def test_strips_multiple_think_blocks(self):
        """Test stripping multiple <think> blocks."""
        messages = [
            {"role": "assistant", "content": "<think>first</think>Hello<think>second</think>World"},
        ]

        result = preprocess_harmony_messages(messages)
        assert result[0]["content"] == "HelloWorld"

    def test_preserves_tool_messages_unchanged(self):
        """Test that tool role messages are preserved unchanged for chat_template."""
        messages = [
            {"role": "user", "content": "What's the weather?"},
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [
                    {
                        "id": "call_123",
                        "function": {"name": "get_weather", "arguments": '{"location": "Seoul"}'},
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_123", "content": '{"temp": 20}'},
        ]

        result = preprocess_harmony_messages(messages)

        # Tool message should be preserved unchanged
        # Chat template expects role="tool" and uses last_tool_call.name from assistant
        assert result[2]["role"] == "tool"
        assert result[2]["tool_call_id"] == "call_123"
        assert result[2]["content"] == '{"temp": 20}'


class TestHarmonyUTF8:
    """Tests for UTF-8 handling with Harmony parser.

    These tests verify that the parser returns token IDs (not decoded text)
    so the caller can use streaming detokenizer for proper UTF-8 handling.
    """

    def test_returns_token_ids_not_text(self, parser, harmony_encoding):
        """Verify that parser returns token IDs, not decoded text."""
        # Process a token in final channel
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>Hello",
            allowed_special="all",
        )

        result = None
        for token in tokens:
            result = parser.process_token(token)

        # Last token should return token ID, not string
        control, stream_token, visible_token, is_stop = result
        assert isinstance(stream_token, int)
        assert isinstance(visible_token, int)

    def test_control_text_is_string(self, parser, harmony_encoding):
        """Verify that control text is always a string."""
        # Enter analysis channel to get <think> control
        tokens = harmony_encoding.encode(
            "<|channel|>analysis<|message|>",
            allowed_special="all",
        )

        control_found = None
        for token in tokens:
            result = parser.process_token(token)
            if result[0]:  # control_text
                control_found = result[0]

        assert isinstance(control_found, str)
        assert control_found == "<think>\n"

    def test_korean_text_handling(self, parser, harmony_encoding):
        """Test that Korean text is handled correctly via token IDs."""
        # Korean text should be tokenized and returned as token IDs
        tokens = harmony_encoding.encode(
            "<|channel|>final<|message|>안녕하세요<|return|>",
            allowed_special="all",
        )

        visible_tokens = []
        for token in tokens:
            _, _, visible_token, _ = parser.process_token(token)
            if visible_token is not None:
                visible_tokens.append(visible_token)

        # Should have token IDs (integers)
        assert all(isinstance(t, int) for t in visible_tokens)

        # Decoding should give Korean text
        decoded = harmony_encoding.decode(visible_tokens)
        assert "안녕" in decoded


class TestExtractHarmonyMessages:
    """Tests for extract_harmony_messages function."""

    def test_preserves_tool_role_and_tool_call_id(self):
        """Test that tool messages preserve role and tool_call_id."""
        from omlx.api.utils import extract_harmony_messages

        # Create mock messages
        class MockMessage:
            def __init__(self, role, content, tool_call_id=None, tool_calls=None):
                self.role = role
                self.content = content
                self.tool_call_id = tool_call_id
                self.tool_calls = tool_calls

        messages = [
            MockMessage("user", "What's the weather?"),
            MockMessage(
                "assistant",
                "",
                tool_calls=[
                    {
                        "id": "call_123",
                        "function": {"name": "get_weather", "arguments": '{"location": "Seoul"}'},
                    }
                ],
            ),
            MockMessage("tool", '{"temp": 20}', tool_call_id="call_123"),
        ]

        result = extract_harmony_messages(messages)

        # Tool message should preserve role and tool_call_id
        assert result[2]["role"] == "tool"
        assert result[2]["tool_call_id"] == "call_123"
        # Content is parsed to dict to avoid double-encoding by chat_template |tojson
        assert result[2]["content"] == {"temp": 20}

    def test_preserves_assistant_tool_calls(self):
        """Test that assistant messages preserve tool_calls field."""
        from omlx.api.utils import extract_harmony_messages

        class MockMessage:
            def __init__(self, role, content, tool_calls=None):
                self.role = role
                self.content = content
                self.tool_calls = tool_calls

        messages = [
            MockMessage(
                "assistant",
                "Let me check the weather.",
                tool_calls=[
                    {
                        "id": "call_456",
                        "function": {"name": "get_weather", "arguments": '{"location": "Tokyo"}'},
                    }
                ],
            ),
        ]

        result = extract_harmony_messages(messages)

        # Assistant message should preserve tool_calls
        assert result[0]["role"] == "assistant"
        assert result[0]["content"] == "Let me check the weather."
        assert "tool_calls" in result[0]
        assert len(result[0]["tool_calls"]) == 1
        assert result[0]["tool_calls"][0]["id"] == "call_456"
        assert result[0]["tool_calls"][0]["function"]["name"] == "get_weather"
        # Arguments are parsed to dict to avoid double-encoding by chat_template |tojson
        assert result[0]["tool_calls"][0]["function"]["arguments"] == {"location": "Tokyo"}

    def test_handles_content_array(self):
        """Test that content arrays are properly extracted."""
        from omlx.api.utils import extract_harmony_messages

        class MockMessage:
            def __init__(self, role, content):
                self.role = role
                self.content = content

        messages = [
            MockMessage("user", [
                {"type": "text", "text": "Hello"},
                {"type": "image_url", "image_url": {"url": "http://example.com/img.png"}},
                {"type": "text", "text": "World"},
            ]),
        ]

        result = extract_harmony_messages(messages)

        # Should extract only text parts
        assert result[0]["role"] == "user"
        assert result[0]["content"] == "Hello\nWorld"

    def test_handles_none_content(self):
        """Test that None content is handled correctly."""
        from omlx.api.utils import extract_harmony_messages

        class MockMessage:
            def __init__(self, role, content, tool_calls=None):
                self.role = role
                self.content = content
                self.tool_calls = tool_calls

        messages = [
            MockMessage("assistant", None, tool_calls=[
                {"id": "call_789", "function": {"name": "search", "arguments": '{}'}}
            ]),
        ]

        result = extract_harmony_messages(messages)

        assert result[0]["role"] == "assistant"
        assert result[0]["content"] == ""
        assert "tool_calls" in result[0]

    def test_differs_from_extract_text_content(self):
        """Test that extract_harmony_messages differs from extract_text_content for tool messages."""
        from omlx.api.utils import extract_harmony_messages, extract_text_content

        class MockMessage:
            def __init__(self, role, content, tool_call_id=None, tool_calls=None):
                self.role = role
                self.content = content
                self.tool_call_id = tool_call_id
                self.tool_calls = tool_calls

        messages = [
            MockMessage("tool", '{"result": "success"}', tool_call_id="call_abc"),
        ]

        harmony_result = extract_harmony_messages(messages)
        text_result = extract_text_content(messages)

        # Harmony should preserve tool role
        assert harmony_result[0]["role"] == "tool"
        assert harmony_result[0]["tool_call_id"] == "call_abc"

        # Text content converts to user role with formatted text
        assert text_result[0]["role"] == "user"
        assert "[Tool Result" in text_result[0]["content"]

    def test_parses_json_tool_content_to_dict(self):
        """Test that JSON string content is parsed to dict to avoid double-encoding."""
        from omlx.api.utils import extract_harmony_messages

        class MockMessage:
            def __init__(self, role, content, tool_call_id=None):
                self.role = role
                self.content = content
                self.tool_call_id = tool_call_id

        messages = [
            MockMessage("tool", '{"temp": 20, "condition": "sunny"}', tool_call_id="call_123"),
        ]

        result = extract_harmony_messages(messages)

        # Content should be parsed to dict (for chat_template |tojson filter)
        assert result[0]["role"] == "tool"
        assert isinstance(result[0]["content"], dict)
        assert result[0]["content"]["temp"] == 20
        assert result[0]["content"]["condition"] == "sunny"

    def test_keeps_non_json_tool_content_as_string(self):
        """Test that non-JSON content remains as string."""
        from omlx.api.utils import extract_harmony_messages

        class MockMessage:
            def __init__(self, role, content, tool_call_id=None):
                self.role = role
                self.content = content
                self.tool_call_id = tool_call_id

        messages = [
            MockMessage("tool", "The weather is sunny and 20 degrees.", tool_call_id="call_123"),
        ]

        result = extract_harmony_messages(messages)

        # Content should remain as string
        assert result[0]["role"] == "tool"
        assert isinstance(result[0]["content"], str)
        assert result[0]["content"] == "The weather is sunny and 20 degrees."

    def test_parses_tool_call_arguments_to_dict(self):
        """Test that tool_call arguments are parsed to dict to avoid double-encoding."""
        from omlx.api.utils import extract_harmony_messages

        class MockMessage:
            def __init__(self, role, content, tool_calls=None):
                self.role = role
                self.content = content
                self.tool_calls = tool_calls

        messages = [
            MockMessage(
                "assistant",
                "",
                tool_calls=[
                    {
                        "id": "call_456",
                        "function": {"name": "get_weather", "arguments": '{"location": "Seoul", "unit": "celsius"}'},
                    }
                ],
            ),
        ]

        result = extract_harmony_messages(messages)

        # Arguments should be parsed to dict (for chat_template |tojson filter)
        assert result[0]["role"] == "assistant"
        assert "tool_calls" in result[0]
        arguments = result[0]["tool_calls"][0]["function"]["arguments"]
        assert isinstance(arguments, dict)
        assert arguments["location"] == "Seoul"
        assert arguments["unit"] == "celsius"