"""Tests for ieum.stream_handler -- base stream handler and token normalization.

Covers:
- normalize_token_usage for multiple LLM provider formats
- _is_summarization_chunk detection
- BaseStreamHandler._process_tool_call_block accumulation
- _extract_tokens method with various message shapes
"""

import pytest

from ieum.stream_handler import (
    BaseStreamHandler,
    _is_summarization_chunk,
    normalize_token_usage,
)


# ---------------------------------------------------------------------------
# normalize_token_usage
# ---------------------------------------------------------------------------


class TestNormalizeTokenUsage:
    """Test token usage normalization across provider formats."""

    def test_none_input(self):
        result = normalize_token_usage(None)
        assert result == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}

    def test_empty_dict(self):
        result = normalize_token_usage({})
        assert result == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}

    def test_openai_format(self):
        usage = {
            "prompt_tokens": 100,
            "completion_tokens": 50,
            "total_tokens": 150,
        }
        result = normalize_token_usage(usage)
        assert result["input_tokens"] == 100
        assert result["output_tokens"] == 50
        assert result["total_tokens"] == 150

    def test_anthropic_format(self):
        usage = {
            "input_tokens": 200,
            "output_tokens": 75,
        }
        result = normalize_token_usage(usage)
        assert result["input_tokens"] == 200
        assert result["output_tokens"] == 75
        assert result["total_tokens"] == 275

    def test_ollama_format(self):
        usage = {
            "prompt_eval_count": 300,
            "eval_count": 100,
        }
        result = normalize_token_usage(usage)
        assert result["input_tokens"] == 300
        assert result["output_tokens"] == 100
        assert result["total_tokens"] == 400

    def test_total_tokens_override(self):
        """If total_tokens is explicitly set, use it."""
        usage = {
            "input_tokens": 100,
            "output_tokens": 50,
            "total_tokens": 200,  # Explicit total (maybe includes cache)
        }
        result = normalize_token_usage(usage)
        assert result["total_tokens"] == 200

    def test_partial_data(self):
        usage = {"prompt_tokens": 100}
        result = normalize_token_usage(usage)
        assert result["input_tokens"] == 100
        assert result["output_tokens"] == 0
        assert result["total_tokens"] == 100


# ---------------------------------------------------------------------------
# _is_summarization_chunk
# ---------------------------------------------------------------------------


class TestIsSummarizationChunk:
    """Test filtering of summarization middleware chunks."""

    def test_none_metadata(self):
        assert _is_summarization_chunk(None) is False

    def test_empty_metadata(self):
        assert _is_summarization_chunk({}) is False

    def test_summarization_chunk(self):
        assert _is_summarization_chunk({"lc_source": "summarization"}) is True

    def test_non_summarization_chunk(self):
        assert _is_summarization_chunk({"lc_source": "agent"}) is False

    def test_missing_key(self):
        assert _is_summarization_chunk({"other_key": "value"}) is False


# ---------------------------------------------------------------------------
# _process_tool_call_block
# ---------------------------------------------------------------------------


class TestProcessToolCallBlock:
    """Test tool call chunk accumulation logic."""

    def test_single_complete_block(self):
        buffers = {}
        block = {
            "type": "tool_call",
            "name": "shell",
            "id": "tc-1",
            "args": {"command": "ls"},
            "index": 0,
        }
        BaseStreamHandler._process_tool_call_block(block, buffers)
        assert 0 in buffers
        assert buffers[0]["name"] == "shell"
        assert buffers[0]["id"] == "tc-1"
        assert buffers[0]["args"] == {"command": "ls"}

    def test_streaming_chunks_accumulate(self):
        buffers = {}
        # First chunk: name and id
        block1 = {
            "type": "tool_call_chunk",
            "name": "shell",
            "id": "tc-1",
            "args": '{"co',
            "index": 0,
        }
        BaseStreamHandler._process_tool_call_block(block1, buffers)
        # Second chunk: more args
        block2 = {
            "type": "tool_call_chunk",
            "name": None,
            "id": None,
            "args": 'mmand":',
            "index": 0,
        }
        BaseStreamHandler._process_tool_call_block(block2, buffers)
        # Third chunk: finish args
        block3 = {
            "type": "tool_call_chunk",
            "name": None,
            "id": None,
            "args": '"ls"}',
            "index": 0,
        }
        BaseStreamHandler._process_tool_call_block(block3, buffers)

        assert buffers[0]["name"] == "shell"
        assert buffers[0]["id"] == "tc-1"
        assert buffers[0]["args"] == '{"command":"ls"}'

    def test_dict_args_replace_string_parts(self):
        buffers = {}
        block = {
            "type": "tool_call",
            "name": "python",
            "id": "tc-2",
            "args": {"code": "print(1)"},
            "index": 1,
        }
        BaseStreamHandler._process_tool_call_block(block, buffers)
        assert buffers[1]["args"] == {"code": "print(1)"}
        assert buffers[1]["args_parts"] == []


# ---------------------------------------------------------------------------
# _extract_tokens
# ---------------------------------------------------------------------------


class TestExtractTokens:
    """Test token extraction from various message formats."""

    class _DummyHandler(BaseStreamHandler):
        """Minimal concrete handler to test the non-abstract methods."""

        async def on_chat_token(self, token, ns_key): pass
        async def on_tool_start(self, tool_id, tool_name, tool_input, ns_key): pass
        async def on_tool_end(self, tool_id, tool_name, output, status): pass
        async def on_file_changed(self, record): pass
        async def on_thinking_token(self, token): pass
        async def on_thinking_end(self): pass
        async def on_message_complete(self, ns_key, tokens): pass
        async def on_human_message(self, ns_key): pass
        async def on_stream_error(self, error): pass
        async def on_interrupted(self): pass
        async def on_tool_rejected_cleanup(self): pass
        async def on_rejection_message(self): pass
        async def request_approval(self, action_requests, assistant_id): pass
        async def request_interview(self, interview_id, questions): pass
        async def on_approval_auto_approved(self): pass
        async def on_pre_stream(self): pass
        async def on_post_stream(self, tokens): pass
        async def on_interview_complete(self, interview_id): pass

    def setup_method(self):
        self.handler = self._DummyHandler()

    def _make_message(self, usage_metadata=None, response_metadata=None):
        class FakeMessage:
            pass
        msg = FakeMessage()
        msg.usage_metadata = usage_metadata
        msg.response_metadata = response_metadata
        return msg

    def test_usage_metadata_openai(self):
        msg = self._make_message(usage_metadata={
            "prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150
        })
        result = self.handler._extract_tokens(msg, 0)
        assert result == 150

    def test_usage_metadata_anthropic(self):
        msg = self._make_message(usage_metadata={
            "input_tokens": 200, "output_tokens": 75
        })
        result = self.handler._extract_tokens(msg, 0)
        assert result == 275

    def test_keeps_max_value(self):
        msg = self._make_message(usage_metadata={
            "input_tokens": 50, "output_tokens": 30
        })
        result = self.handler._extract_tokens(msg, 1000)
        assert result == 1000  # Previous max is higher

    def test_fallback_to_response_metadata(self):
        msg = self._make_message(
            usage_metadata=None,
            response_metadata={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}
        )
        result = self.handler._extract_tokens(msg, 0)
        assert result == 150

    def test_nested_usage_in_response_metadata(self):
        msg = self._make_message(
            usage_metadata=None,
            response_metadata={"usage": {"prompt_tokens": 80, "completion_tokens": 40, "total_tokens": 120}}
        )
        result = self.handler._extract_tokens(msg, 0)
        assert result == 120

    def test_no_metadata_returns_current(self):
        msg = self._make_message()
        result = self.handler._extract_tokens(msg, 42)
        assert result == 42
