"""
Unit tests for _AnthropicStreamWrapper and _AnthropicStreamManagerWrapper.
No anthropic package or network required.
"""
import types
import time
from datetime import datetime, timezone

from visibe.integrations.anthropic import (
    _AnthropicStreamWrapper,
    _AnthropicStreamManagerWrapper,
    AnthropicIntegration,
)


# ------------------------------------------------------------------
# Minimal mock helpers
# ------------------------------------------------------------------

class _MockVisibe:
    def __init__(self):
        self._spans = []
        self._trace_header = None
        self._summary = None

    def create_trace(self, data):
        self._trace_header = data

    def queue_span(self, trace_id, span):
        self._spans.append(span)

    def complete_trace(self, trace_id, summary):
        self._summary = summary


def _make_event(type_, **fields):
    return types.SimpleNamespace(type=type_, **fields)


def _make_stream_wrapper(events, mock=None, group_tracker=None):
    """Build a wrapper around a list of raw events."""
    if mock is None:
        mock = _MockVisibe()

    class _FakeStream:
        def __init__(self):
            self._iter = iter(events)

        def __next__(self):
            return next(self._iter)

        def __iter__(self):
            return self

    return _AnthropicStreamWrapper(
        stream=_FakeStream(),
        visibe_client=mock,
        group_tracker=group_tracker,
        trace_id='trace-123',
        model='claude-3-5-sonnet-20241022',
        prompt_text='Hello',
        call_start=time.time(),
        call_timestamp=datetime.now(timezone.utc),
        name='test',
    ), mock


# ------------------------------------------------------------------
# _AnthropicStreamWrapper._process_event
# ------------------------------------------------------------------

def test_stream_wrapper_captures_input_tokens_from_message_start():
    usage = types.SimpleNamespace(input_tokens=80)
    msg = types.SimpleNamespace(model='claude-3-5-sonnet-20241022', usage=usage)
    event = _make_event('message_start', message=msg)
    wrapper, _ = _make_stream_wrapper([])
    wrapper._process_event(event)
    assert wrapper._input_tokens == 80


def test_stream_wrapper_captures_model_from_message_start():
    usage = types.SimpleNamespace(input_tokens=10)
    msg = types.SimpleNamespace(model='claude-opus-4-5', usage=usage)
    event = _make_event('message_start', message=msg)
    wrapper, _ = _make_stream_wrapper([])
    wrapper._process_event(event)
    assert wrapper._response_model == 'claude-opus-4-5'


def test_stream_wrapper_accumulates_text_deltas():
    delta1 = types.SimpleNamespace(type='text_delta', text='Hello ')
    delta2 = types.SimpleNamespace(type='text_delta', text='world')
    e1 = _make_event('content_block_delta', delta=delta1)
    e2 = _make_event('content_block_delta', delta=delta2)
    wrapper, _ = _make_stream_wrapper([])
    wrapper._process_event(e1)
    wrapper._process_event(e2)
    assert ''.join(wrapper._collected_content) == 'Hello world'


def test_stream_wrapper_captures_output_tokens_from_message_delta():
    usage = types.SimpleNamespace(output_tokens=42)
    event = _make_event('message_delta', usage=usage)
    wrapper, _ = _make_stream_wrapper([])
    wrapper._process_event(event)
    assert wrapper._output_tokens == 42


def test_stream_wrapper_ignores_unknown_event_types():
    event = _make_event('ping')
    wrapper, _ = _make_stream_wrapper([])
    wrapper._process_event(event)  # must not raise
    assert wrapper._input_tokens == 0
    assert wrapper._output_tokens == 0


def test_stream_wrapper_skips_input_json_delta_without_crashing():
    delta = types.SimpleNamespace(type='input_json_delta', partial_json='{"q":')
    event = _make_event('content_block_delta', delta=delta)
    wrapper, _ = _make_stream_wrapper([])
    wrapper._process_event(event)  # must not raise
    assert wrapper._collected_content == []


# ------------------------------------------------------------------
# _AnthropicStreamWrapper iteration + finalize (no group tracker)
# ------------------------------------------------------------------

def test_stream_wrapper_iteration_exhausts_stream():
    usage_in = types.SimpleNamespace(input_tokens=50)
    msg = types.SimpleNamespace(model='claude-3-5-sonnet-20241022', usage=usage_in)
    usage_out = types.SimpleNamespace(output_tokens=20)

    events = [
        _make_event('message_start', message=msg),
        _make_event('content_block_delta',
                    delta=types.SimpleNamespace(type='text_delta', text='Hi!')),
        _make_event('message_delta', usage=usage_out),
    ]
    wrapper, mock = _make_stream_wrapper(events)

    collected = list(wrapper)  # exhausts stream → triggers _finalize
    assert len(collected) == 3

    # _finalize must have sent a trace
    assert mock._summary is not None
    assert mock._summary['llm_call_count'] == 1
    assert mock._summary['total_input_tokens'] == 50
    assert mock._summary['total_output_tokens'] == 20
    assert mock._summary['prompt'] == 'Hello'
    # output goes into the span, not the summary
    llm_spans = [s for s in mock._spans if s.get('type') == 'llm_call']
    assert llm_spans and llm_spans[0]['output_text'] == 'Hi!'


def test_stream_wrapper_finalize_only_called_once():
    """Double-exhausting must not double-send the trace."""
    events = [_make_event('message_delta', usage=types.SimpleNamespace(output_tokens=5))]
    wrapper, mock = _make_stream_wrapper(events)

    list(wrapper)        # first exhaustion
    wrapper._finalize()  # explicit second call

    # complete_trace (summary) should only have been set once
    assert mock._summary is not None


def test_stream_wrapper_finalize_sends_to_group_tracker():
    class _FakeTracker:
        def __init__(self):
            self.calls = []

        def add_call(self, data):
            self.calls.append(data)

    tracker = _FakeTracker()
    events = [
        _make_event('message_delta', usage=types.SimpleNamespace(output_tokens=10)),
    ]
    wrapper, mock = _make_stream_wrapper(events, group_tracker=tracker)
    list(wrapper)

    assert len(tracker.calls) == 1
    assert tracker.calls[0]['output_tokens'] == 10
    # complete_trace not called (group tracker handles it)
    assert mock._summary is None


# ------------------------------------------------------------------
# _AnthropicStreamWrapper.text_stream
# ------------------------------------------------------------------

def test_stream_wrapper_text_stream_yields_text_only():
    events = [
        _make_event('message_start',
                    message=types.SimpleNamespace(
                        model='claude-3-5-sonnet-20241022',
                        usage=types.SimpleNamespace(input_tokens=10))),
        _make_event('content_block_delta',
                    delta=types.SimpleNamespace(type='text_delta', text='Part1')),
        _make_event('content_block_delta',
                    delta=types.SimpleNamespace(type='text_delta', text='Part2')),
        _make_event('message_delta',
                    usage=types.SimpleNamespace(output_tokens=5)),
    ]
    wrapper, _ = _make_stream_wrapper(events)
    texts = list(wrapper.text_stream)
    assert texts == ['Part1', 'Part2']


# ------------------------------------------------------------------
# _AnthropicStreamManagerWrapper
# ------------------------------------------------------------------

def _make_integration():
    integration = AnthropicIntegration.__new__(AnthropicIntegration)
    integration._instrumented_clients = {}
    integration._insideStream = False
    return integration


class _FakeFinalMessage:
    def __init__(self, input_tokens=30, output_tokens=15, text='Answer', model='claude-3-5-sonnet-20241022'):
        self.model = model
        self.usage = types.SimpleNamespace(input_tokens=input_tokens, output_tokens=output_tokens)
        self.content = [types.SimpleNamespace(type='text', text=text)]


class _FakeInnerStream:
    def __init__(self, final_msg):
        self._final_msg = final_msg

    def get_final_message(self):
        return self._final_msg


class _FakeManager:
    """Mimics anthropic.MessageStreamManager."""
    def __init__(self, inner_stream):
        self._inner = inner_stream
        self._entered = False
        self._exited = False

    def __enter__(self):
        self._entered = True
        return self._inner

    def __exit__(self, *args):
        self._exited = True
        return False


def _make_manager_wrapper(final_msg=None, mock=None, group_tracker=None):
    if mock is None:
        mock = _MockVisibe()
    if final_msg is None:
        final_msg = _FakeFinalMessage()
    inner = _FakeInnerStream(final_msg)
    mgr = _FakeManager(inner)
    integration = _make_integration()

    wrapper = _AnthropicStreamManagerWrapper(
        original_manager=mgr,
        integration=integration,
        visibe_client=mock,
        group_tracker=group_tracker,
        trace_id='trace-mgr',
        model='claude-3-5-sonnet-20241022',
        prompt_text='Test prompt',
        call_start=time.time(),
        call_timestamp=datetime.now(timezone.utc),
        name='mgr-test',
    )
    return wrapper, mock, integration


def test_manager_wrapper_enter_sets_and_clears_inside_stream():
    wrapper, _, integration = _make_manager_wrapper()
    assert not integration._insideStream
    with wrapper:
        # _insideStream was True during __enter__, cleared by the time body runs
        pass
    assert not integration._insideStream


def test_manager_wrapper_enter_clears_flag_even_on_exception():
    integration = _make_integration()
    mock = _MockVisibe()

    class _BrokenManager:
        def __enter__(self):
            raise RuntimeError("boom")
        def __exit__(self, *a):
            return False

    wrapper = _AnthropicStreamManagerWrapper(
        original_manager=_BrokenManager(),
        integration=integration,
        visibe_client=mock,
        group_tracker=None,
        trace_id='t',
        model='m',
        prompt_text='p',
        call_start=time.time(),
        call_timestamp=datetime.now(timezone.utc),
    )
    try:
        wrapper.__enter__()
    except RuntimeError:
        pass
    assert not integration._insideStream  # flag must be cleared


def test_manager_wrapper_exit_reads_final_message():
    final = _FakeFinalMessage(input_tokens=40, output_tokens=20, text='Done')
    wrapper, mock, _ = _make_manager_wrapper(final_msg=final)
    with wrapper:
        pass

    assert mock._summary is not None
    assert mock._summary['total_input_tokens'] == 40
    assert mock._summary['total_output_tokens'] == 20
    assert mock._summary['prompt'] == 'Test prompt'
    # output goes into the span, not the summary
    llm_spans = [s for s in mock._spans if s.get('type') == 'llm_call']
    assert llm_spans and llm_spans[0]['output_text'] == 'Done'


def test_manager_wrapper_exit_sends_to_group_tracker():
    class _FakeTracker:
        def __init__(self):
            self.calls = []

        def add_call(self, d):
            self.calls.append(d)

    tracker = _FakeTracker()
    final = _FakeFinalMessage(input_tokens=25, output_tokens=12)
    wrapper, mock, _ = _make_manager_wrapper(final_msg=final, group_tracker=tracker)
    with wrapper:
        pass

    assert len(tracker.calls) == 1
    assert tracker.calls[0]['input_tokens'] == 25
    assert mock._summary is None  # group tracker, not direct trace


def test_manager_wrapper_exit_skips_finalize_on_exception():
    """If an exception propagates, _finalize should not run (exc_type is set)."""
    final = _FakeFinalMessage()
    wrapper, mock, _ = _make_manager_wrapper(final_msg=final)
    try:
        with wrapper:
            raise ValueError("user error")
    except ValueError:
        pass
    # No trace sent when the body raised
    assert mock._summary is None
