"""
Anthropic provider integration for Visibe.

Tracks direct Anthropic SDK calls (messages.create / messages.stream) by
wrapping the client.  Also works with AsyncAnthropic for async code.

Two modes:
1. instrument(client) — wrap once, every call auto-sends its own trace
2. track(client, name) — context manager to group multiple calls into one trace
"""
import inspect
import logging
import time
import uuid
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional

from .base import TraceSummary
from ..utils import calculate_cost

logger = logging.getLogger("visibe.anthropic")

try:
    import anthropic as _anthropic_module
    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False


class AnthropicIntegration:
    """
    Anthropic provider integration for Visibe observability.

    Tracks every messages.create() and messages.stream() call with token
    usage, cost, model info, and tool calls.

    Usage (instrument mode — each call = one trace):
        from visibe import Visibe
        from anthropic import Anthropic

        obs = Visibe(api_key="sk_live_abc123")
        client = Anthropic()
        obs.instrument(client)

        response = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hello!"}],
        )

        obs.uninstrument(client)

    Usage (track mode — group multiple calls into one trace):
        with obs.track(client=client, name="multi-step-chat"):
            r1 = client.messages.create(...)
            r2 = client.messages.create(...)
    """

    def __init__(self, visibe_client):
        if not ANTHROPIC_AVAILABLE:
            raise ImportError(
                "Anthropic integration requires the anthropic package.\n"
                "Install with: pip install anthropic\n"
            )
        self.client = visibe_client
        self._instrumented_clients = {}  # id(client) -> {'create': fn, 'stream': fn}

        # Re-entry guard: messages.stream() internally calls messages.create(stream=True)
        # inside __enter__.  When this flag is True we're already tracking the stream,
        # so we pass the inner create() call straight through.
        self._insideStream = False

    def instrument(self, anthropic_client, name=None):
        """
        Instrument an Anthropic client so every messages.create() and
        messages.stream() call automatically sends its own trace.

        Args:
            anthropic_client: An anthropic.Anthropic or AsyncAnthropic instance
            name: Optional trace name shown in the Visibe dashboard.
        """
        client_id = id(anthropic_client)
        if client_id in self._instrumented_clients:
            logger.debug("Client is already instrumented")
            return

        original_create = anthropic_client.messages.create
        original_stream = getattr(anthropic_client.messages, 'stream', None)

        self._instrumented_clients[client_id] = {
            'create': original_create,
            'stream': original_stream,
        }

        integration = self
        visibe_client = self.client
        trace_name = name
        is_async = inspect.iscoroutinefunction(original_create)

        if is_async:
            async def wrapped_create(*args, **kwargs):
                # Re-entry: pass through if already inside a stream wrapper.
                if kwargs.get('stream') and integration._insideStream:
                    return await original_create(*args, **kwargs)
                if kwargs.get('stream'):
                    # Async streaming not yet supported — pass through.
                    logger.debug(
                        "Async streaming via create(stream=True) is not yet "
                        "tracked by Visibe; falling back to pass-through."
                    )
                    return await original_create(*args, **kwargs)
                return await AnthropicIntegration._execute_and_trace_async(
                    original_create, visibe_client, None, args, kwargs,
                    name=trace_name,
                )
        else:
            def wrapped_create(*args, **kwargs):
                # Re-entry: messages.stream().__enter__ calls create(stream=True).
                if kwargs.get('stream') and integration._insideStream:
                    return original_create(*args, **kwargs)
                if kwargs.get('stream'):
                    return integration._handle_streaming(
                        original_create, visibe_client, None, args, kwargs,
                        name=trace_name,
                    )
                return integration._execute_and_trace(
                    original_create, visibe_client, None, args, kwargs,
                    name=trace_name,
                )

        anthropic_client.messages.create = wrapped_create

        # Wrap messages.stream() (sync only — async handled via create).
        if original_stream is not None and not is_async:
            def wrapped_stream(*args, **kwargs):
                call_start = time.time()
                call_timestamp = datetime.now(timezone.utc)
                model = kwargs.get('model', args[0] if args else 'unknown')
                prompt_text = _extract_prompt(kwargs.get('messages', []))

                trace_id = str(uuid.uuid4())
                trace_name_stream = trace_name or f"anthropic-{model}"
                visibe_client.create_trace({
                    'trace_id': trace_id,
                    'name': trace_name_stream,
                    'framework': 'anthropic',
                    'model': model,
                    'started_at': call_timestamp.isoformat(),
                })

                return _AnthropicStreamManagerWrapper(
                    original_stream(*args, **kwargs),
                    integration=integration,
                    visibe_client=visibe_client,
                    group_tracker=None,
                    trace_id=trace_id,
                    model=model,
                    prompt_text=prompt_text,
                    call_start=call_start,
                    call_timestamp=call_timestamp,
                    name=trace_name,
                )

            anthropic_client.messages.stream = wrapped_stream

        logger.debug("Instrumented Anthropic client")

    def uninstrument(self, anthropic_client):
        """Restore the original messages.create() and messages.stream() methods."""
        client_id = id(anthropic_client)
        original = self._instrumented_clients.pop(client_id, None)
        if original:
            anthropic_client.messages.create = original['create']
            if original['stream'] is not None:
                anthropic_client.messages.stream = original['stream']
            logger.debug("Uninstrumented Anthropic client")
        else:
            logger.warning("Client was not instrumented")

    @contextmanager
    def track(self, client: Any, name: str = "anthropic-chat"):
        """
        Track multiple messages.create() calls grouped into a single trace.

        Args:
            client: An anthropic.Anthropic client instance
            name: Human-readable name for this trace

        Example:
            with tracker.track(client=client, name="conversation"):
                r1 = client.messages.create(...)
                r2 = client.messages.create(...)
        """
        tracker = _AnthropicGroupTracker(self.client, client, name, self)
        tracker.start()
        try:
            yield
        except Exception as e:
            tracker.record_error(e)
            raise
        finally:
            tracker.stop()

    def _execute_and_trace(self, original_create, visibe_client, group_tracker,
                           args, kwargs, name=None):
        """Execute a sync non-streaming messages.create() call and track it."""
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('model', args[0] if args else 'unknown')
        messages = kwargs.get('messages', [])
        prompt_text = _extract_prompt(messages)

        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            visibe_client.create_trace({
                'trace_id': trace_id,
                'name': name or f"anthropic-{model}",
                'framework': 'anthropic',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            })

        try:
            response = original_create(*args, **kwargs)
            duration_ms = int((time.time() - call_start) * 1000)

            input_tokens = 0
            output_tokens = 0
            if hasattr(response, 'usage') and response.usage:
                input_tokens = getattr(response.usage, 'input_tokens', 0) or 0
                output_tokens = getattr(response.usage, 'output_tokens', 0) or 0

            response_model = getattr(response, 'model', model)
            content = getattr(response, 'content', [])
            output_text = _extract_output_text(content)
            tools_used = _extract_tool_calls(content)
            cost = calculate_cost(response_model, input_tokens, output_tokens)

            call_data = {
                'model': response_model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)
            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )
            raise

    @staticmethod
    async def _execute_and_trace_async(original_create, visibe_client, group_tracker,
                                       args, kwargs, name=None):
        """Async version of _execute_and_trace for AsyncAnthropic clients."""
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('model', args[0] if args else 'unknown')
        messages = kwargs.get('messages', [])
        prompt_text = _extract_prompt(messages)

        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            visibe_client.create_trace({
                'trace_id': trace_id,
                'name': name or f"anthropic-{model}",
                'framework': 'anthropic',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            })

        try:
            response = await original_create(*args, **kwargs)
            duration_ms = int((time.time() - call_start) * 1000)

            input_tokens = 0
            output_tokens = 0
            if hasattr(response, 'usage') and response.usage:
                input_tokens = getattr(response.usage, 'input_tokens', 0) or 0
                output_tokens = getattr(response.usage, 'output_tokens', 0) or 0

            response_model = getattr(response, 'model', model)
            content = getattr(response, 'content', [])
            output_text = _extract_output_text(content)
            tools_used = _extract_tool_calls(content)
            cost = calculate_cost(response_model, input_tokens, output_tokens)

            call_data = {
                'model': response_model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)
            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )
            raise

    def _handle_streaming(self, original_create, visibe_client, group_tracker,
                          args, kwargs, name=None):
        """Handle messages.create(stream=True) calls."""
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('model', args[0] if args else 'unknown')
        messages = kwargs.get('messages', [])
        prompt_text = _extract_prompt(messages)

        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            visibe_client.create_trace({
                'trace_id': trace_id,
                'name': name or f"anthropic-{model}",
                'framework': 'anthropic',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            })

        try:
            stream = original_create(*args, **kwargs)
        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)
            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )
            raise

        return _AnthropicStreamWrapper(
            stream, visibe_client, group_tracker, trace_id, model, prompt_text,
            call_start, call_timestamp, name=name,
        )


# ------------------------------------------------------------------
# Stream wrappers
# ------------------------------------------------------------------

class _AnthropicStreamWrapper:
    """
    Wraps a Stream[RawMessageStreamEvent] (from messages.create(stream=True)).

    Collects content deltas and token usage from raw SSE events so we can
    send a trace span after the stream is exhausted.
    """

    def __init__(self, stream, visibe_client, group_tracker, trace_id, model,
                 prompt_text, call_start, call_timestamp, name=None):
        self._stream = stream
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name

        self._collected_content = []
        self._input_tokens = 0
        self._output_tokens = 0
        self._response_model = model
        self._tool_calls = []

    # -- context manager support (stream used as `with client.messages.create(stream=True) as s`) --

    def __enter__(self):
        if hasattr(self._stream, '__enter__'):
            self._stream.__enter__()
        return self

    def __exit__(self, *args):
        self._finalize()
        if hasattr(self._stream, '__exit__'):
            return self._stream.__exit__(*args)
        return False

    # -- iterator protocol (stream used as `for event in stream`) --

    def __iter__(self):
        return self

    def __next__(self):
        try:
            event = next(self._stream)
            self._process_event(event)
            return event
        except StopIteration:
            self._finalize()
            raise

    # -- text_stream convenience helper passthrough --

    @property
    def text_stream(self):
        """Yield text deltas; finalize on exhaustion."""
        for event in self:
            event_type = getattr(event, 'type', None)
            if event_type == 'content_block_delta':
                delta = getattr(event, 'delta', None)
                if delta and getattr(delta, 'type', None) == 'text_delta':
                    yield delta.text

    def get_final_message(self):
        """Passthrough to the underlying stream if available."""
        if hasattr(self._stream, 'get_final_message'):
            return self._stream.get_final_message()
        return None

    def _process_event(self, event):
        event_type = getattr(event, 'type', None)

        if event_type == 'message_start':
            msg = getattr(event, 'message', None)
            if msg:
                self._response_model = getattr(msg, 'model', self._model)
                usage = getattr(msg, 'usage', None)
                if usage:
                    self._input_tokens = getattr(usage, 'input_tokens', 0) or 0

        elif event_type == 'content_block_delta':
            delta = getattr(event, 'delta', None)
            if delta:
                if getattr(delta, 'type', None) == 'text_delta':
                    self._collected_content.append(getattr(delta, 'text', '') or '')
                elif getattr(delta, 'type', None) == 'input_json_delta':
                    # Tool input accumulation — handled at message_stop
                    pass

        elif event_type == 'message_delta':
            usage = getattr(event, 'usage', None)
            if usage:
                self._output_tokens = getattr(usage, 'output_tokens', 0) or 0

    def _finalize(self):
        if hasattr(self, '_finalized'):
            return
        self._finalized = True

        duration_ms = int((time.time() - self._call_start) * 1000)
        output_text = ''.join(self._collected_content)
        cost = calculate_cost(self._response_model, self._input_tokens, self._output_tokens)

        # Try to get tool calls from the final message if available
        tools_used = []
        if hasattr(self._stream, 'get_final_message'):
            try:
                final_msg = self._stream.get_final_message()
                if final_msg:
                    tools_used = _extract_tool_calls(getattr(final_msg, 'content', []))
                    if not output_text:
                        output_text = _extract_output_text(
                            getattr(final_msg, 'content', [])
                        )
            except Exception:
                pass

        call_data = {
            'model': self._response_model,
            'input_tokens': self._input_tokens,
            'output_tokens': self._output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(
                self._visibe_client, self._trace_id, call_data, name=self._name
            )


class _AnthropicStreamManagerWrapper:
    """
    Wraps a MessageStreamManager (from messages.stream()).

    Uses the re-entry guard on the integration so that the inner
    messages.create(stream=True) call — triggered by __enter__ — is
    passed straight through without being double-traced.

    On __exit__ the final message (with full usage) is read from the
    MessageStream and used to send the trace.
    """

    def __init__(self, original_manager, integration, visibe_client, group_tracker,
                 trace_id, model, prompt_text, call_start, call_timestamp, name=None):
        self._original = original_manager
        self._integration = integration
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._inner_stream = None

    def __enter__(self):
        # Set re-entry guard so the inner create(stream=True) call passes through.
        self._integration._insideStream = True
        try:
            self._inner_stream = self._original.__enter__()
        finally:
            self._integration._insideStream = False
        return self._inner_stream

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type is None and self._inner_stream is not None:
                self._finalize()
        except Exception:
            pass
        return self._original.__exit__(exc_type, exc_val, exc_tb)

    def _finalize(self):
        duration_ms = int((time.time() - self._call_start) * 1000)

        input_tokens = 0
        output_tokens = 0
        response_model = self._model
        output_text = ''
        tools_used = []

        try:
            final_msg = self._inner_stream.get_final_message()
            if final_msg:
                response_model = getattr(final_msg, 'model', self._model)
                usage = getattr(final_msg, 'usage', None)
                if usage:
                    input_tokens = getattr(usage, 'input_tokens', 0) or 0
                    output_tokens = getattr(usage, 'output_tokens', 0) or 0
                content = getattr(final_msg, 'content', [])
                output_text = _extract_output_text(content)
                tools_used = _extract_tool_calls(content)
        except Exception:
            pass

        cost = calculate_cost(response_model, input_tokens, output_tokens)

        call_data = {
            'model': response_model,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(
                self._visibe_client, self._trace_id, call_data, name=self._name
            )


# ------------------------------------------------------------------
# Group tracker (for track() context manager)
# ------------------------------------------------------------------

class _AnthropicGroupTracker:
    """Internal tracker for grouping multiple Anthropic calls into one trace."""

    def __init__(self, visibe_client, anthropic_client, name: str,
                 integration: 'AnthropicIntegration'):
        self.visibe_client = visibe_client
        self.anthropic_client = anthropic_client
        self.name = name
        self._integration = integration
        self.trace_id = str(uuid.uuid4())

        self.start_time = None
        self.end_time = None
        self.calls: List[Dict] = []
        self.errors: List[Dict] = []
        self._step_counter = 0

        self._saved_create = None
        self._saved_stream = None
        self._original_create = None

    def _next_step_id(self) -> str:
        self._step_counter += 1
        return f"step_{self._step_counter}"

    def start(self):
        self.start_time = datetime.now(timezone.utc)

        self.visibe_client.create_trace({
            'trace_id': self.trace_id,
            'name': self.name,
            'framework': 'anthropic',
            'started_at': self.start_time.isoformat(),
        })

        self._saved_create = self.anthropic_client.messages.create
        self._saved_stream = getattr(self.anthropic_client.messages, 'stream', None)

        # Use the real original (bypass instrument wrapper if active)
        client_id = id(self.anthropic_client)
        if client_id in self._integration._instrumented_clients:
            entry = self._integration._instrumented_clients[client_id]
            self._original_create = entry.get('create')
        else:
            self._original_create = self.anthropic_client.messages.create

        tracker = self
        original = self._original_create
        integration = self._integration
        is_async = inspect.iscoroutinefunction(original)

        if is_async:
            async def wrapped_create(*args, **kwargs):
                if kwargs.get('stream') and integration._insideStream:
                    return await original(*args, **kwargs)
                if kwargs.get('stream'):
                    logger.debug("Async streaming in track() not yet supported; passing through.")
                    return await original(*args, **kwargs)
                return await AnthropicIntegration._execute_and_trace_async(
                    original, tracker.visibe_client, tracker, args, kwargs
                )
        else:
            def wrapped_create(*args, **kwargs):
                if kwargs.get('stream') and integration._insideStream:
                    return original(*args, **kwargs)
                if kwargs.get('stream'):
                    return integration._handle_streaming(
                        original, tracker.visibe_client, tracker, args, kwargs
                    )
                return integration._execute_and_trace(
                    original, tracker.visibe_client, tracker, args, kwargs
                )

        self.anthropic_client.messages.create = wrapped_create

        # Wrap messages.stream() in track mode (sync only)
        if self._saved_stream is not None and not is_async:
            orig_stream = self._integration._instrumented_clients.get(
                id(self.anthropic_client), {}
            ).get('stream') or self._saved_stream

            def wrapped_stream(*args, **kwargs):
                call_start = time.time()
                call_timestamp = datetime.now(timezone.utc)
                model = kwargs.get('model', args[0] if args else 'unknown')
                prompt_text = _extract_prompt(kwargs.get('messages', []))

                return _AnthropicStreamManagerWrapper(
                    orig_stream(*args, **kwargs),
                    integration=integration,
                    visibe_client=tracker.visibe_client,
                    group_tracker=tracker,
                    trace_id=tracker.trace_id,
                    model=model,
                    prompt_text=prompt_text,
                    call_start=call_start,
                    call_timestamp=call_timestamp,
                    name=tracker.name,
                )

            self.anthropic_client.messages.stream = wrapped_stream

    def stop(self):
        self.end_time = datetime.now(timezone.utc)
        if self._saved_create:
            self.anthropic_client.messages.create = self._saved_create
        if self._saved_stream is not None:
            self.anthropic_client.messages.stream = self._saved_stream
        self._complete_trace()

    def add_call(self, call_data: Dict):
        self.calls.append(call_data)

        span = {
            'span_id': self._next_step_id(),
            'type': 'llm_call',
            'timestamp': call_data['timestamp'].isoformat(),
            'agent_name': 'anthropic',
            'model': call_data['model'],
            'provider': 'anthropic',
            'status': 'success',
            'description': f"LLM Call using {call_data['model']}",
            'input_tokens': call_data['input_tokens'],
            'output_tokens': call_data['output_tokens'],
            'cost': call_data['cost'],
            'duration_ms': call_data['duration_ms'],
            'input_text': (call_data.get('prompt') or '')[:1000],
            'output_text': (call_data.get('output') or '')[:1000],
        }
        self.visibe_client.queue_span(self.trace_id, span)

        for i, tool in enumerate(call_data.get('tools_used', []), start=1):
            tool_span = {
                'span_id': self._next_step_id(),
                'type': 'tool_call',
                'timestamp': call_data['timestamp'].isoformat(),
                'tool_name': tool['tool_name'],
                'agent_name': 'anthropic',
                'status': 'success',
                'input_text': tool.get('input', ''),
                'output_text': '',
            }
            self.visibe_client.queue_span(self.trace_id, tool_span)

    def record_error(self, exc: Exception):
        self.errors.append({
            'type': type(exc).__name__,
            'message': str(exc)[:500],
        })

        span = {
            'span_id': self._next_step_id(),
            'type': 'error',
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'description': f"Error: {type(exc).__name__}",
            'error_type': type(exc).__name__,
            'error_message': str(exc)[:500],
        }
        self.visibe_client.queue_span(self.trace_id, span)

    def _complete_trace(self):
        duration_ms = int(
            (self.end_time - self.start_time).total_seconds() * 1000
        )

        total_input = sum(c['input_tokens'] for c in self.calls)
        total_output = sum(c['output_tokens'] for c in self.calls)
        total_cost = round(sum(c['cost'] for c in self.calls), 6)

        prompt = self.calls[0]['prompt'] if self.calls else self.name
        model = self.calls[0]['model'] if self.calls else None
        status = 'failed' if self.errors else 'completed'

        summary = {
            'status': status,
            'ended_at': self.end_time.isoformat(),
            'duration_ms': duration_ms,
            'prompt': prompt,
            'model': model,
            'llm_call_count': len(self.calls),
            'total_cost': total_cost,
            'total_tokens': total_input + total_output,
            'total_input_tokens': total_input,
            'total_output_tokens': total_output,
        }

        success = self.visibe_client.complete_trace(self.trace_id, summary)
        if success:
            logger.info(f"Anthropic trace completed: {self.trace_id}")
        else:
            logger.error(f"Failed to complete Anthropic trace: {self.trace_id}")

        total_tool_calls = sum(len(c.get('tools_used', [])) for c in self.calls)
        trace_summary = TraceSummary(
            name=self.name,
            status=status,
            llm_calls=len(self.calls),
            tool_calls=total_tool_calls,
            total_tokens=total_input + total_output,
            total_cost=total_cost,
            duration_s=round(duration_ms / 1000, 1),
            agents=['anthropic'],
            sent=success,
        )
        print(f"[Visibe] {trace_summary}")


# ------------------------------------------------------------------
# Standalone trace helpers (instrument mode — one call = one trace)
# ------------------------------------------------------------------

def _send_call_as_trace(visibe_client, trace_id: str, call_data: dict,
                        name: str = None):
    """Send a single-call trace: span + complete (instrument mode)."""
    now = call_data['timestamp']
    end_time = now + timedelta(milliseconds=call_data['duration_ms'])

    span = {
        'span_id': 'step_1',
        'type': 'llm_call',
        'timestamp': now.isoformat(),
        'agent_name': 'anthropic',
        'model': call_data['model'],
        'provider': 'anthropic',
        'status': 'success',
        'description': f"LLM Call using {call_data['model']}",
        'input_tokens': call_data['input_tokens'],
        'output_tokens': call_data['output_tokens'],
        'cost': call_data['cost'],
        'duration_ms': call_data['duration_ms'],
        'input_text': (call_data.get('prompt') or '')[:1000],
        'output_text': (call_data.get('output') or '')[:1000],
    }
    visibe_client.queue_span(trace_id, span)

    for i, tool in enumerate(call_data.get('tools_used', []), start=2):
        tool_span = {
            'span_id': f'step_{i}',
            'type': 'tool_call',
            'timestamp': now.isoformat(),
            'tool_name': tool['tool_name'],
            'agent_name': 'anthropic',
            'status': 'success',
            'input_text': tool.get('input', ''),
            'output_text': '',
        }
        visibe_client.queue_span(trace_id, tool_span)

    summary = {
        'status': 'completed',
        'ended_at': end_time.isoformat(),
        'duration_ms': call_data['duration_ms'],
        'prompt': call_data.get('prompt') or 'No prompt captured',
        'model': call_data['model'],
        'llm_call_count': 1,
        'total_cost': call_data['cost'],
        'total_tokens': call_data['input_tokens'] + call_data['output_tokens'],
        'total_input_tokens': call_data['input_tokens'],
        'total_output_tokens': call_data['output_tokens'],
    }
    success = visibe_client.complete_trace(trace_id, summary)

    trace_summary = TraceSummary(
        name=name or call_data['model'],
        status='completed',
        llm_calls=1,
        tool_calls=len(call_data.get('tools_used', [])),
        total_tokens=call_data['input_tokens'] + call_data['output_tokens'],
        total_cost=call_data['cost'],
        duration_s=round(call_data['duration_ms'] / 1000, 1),
        agents=['anthropic'],
        sent=success,
    )
    print(f"[Visibe] {trace_summary}")


def _send_error_as_trace(visibe_client, trace_id: str, model: str,
                         prompt: str, exc: Exception,
                         timestamp: datetime, duration_ms: int,
                         name: str = None):
    """Send a failed trace: error span + complete as failed (instrument mode)."""
    end_time = timestamp + timedelta(milliseconds=duration_ms)

    error_span = {
        'span_id': 'step_1',
        'type': 'error',
        'timestamp': timestamp.isoformat(),
        'agent_name': 'anthropic',
        'model': model,
        'description': f"Error: {type(exc).__name__}",
        'error_type': type(exc).__name__,
        'error_message': str(exc)[:500],
        'duration_ms': duration_ms,
    }
    visibe_client.queue_span(trace_id, error_span)

    summary = {
        'status': 'failed',
        'ended_at': end_time.isoformat(),
        'duration_ms': duration_ms,
        'prompt': prompt or 'No prompt captured',
        'model': model,
        'llm_call_count': 0,
        'total_cost': 0.0,
        'total_tokens': 0,
        'total_input_tokens': 0,
        'total_output_tokens': 0,
    }
    visibe_client.complete_trace(trace_id, summary)

    trace_summary = TraceSummary(
        name=name or model,
        status='failed',
        llm_calls=0,
        tool_calls=0,
        total_tokens=0,
        total_cost=0.0,
        duration_s=round(duration_ms / 1000, 1),
        agents=['anthropic'],
        sent=False,
    )
    print(f"[Visibe] {trace_summary}")


# ------------------------------------------------------------------
# Extraction helpers
# ------------------------------------------------------------------

def _extract_prompt(messages: list) -> str:
    """Return the last user message text (truncated to 500 chars)."""
    for msg in reversed(messages or []):
        if not isinstance(msg, dict):
            continue
        if msg.get('role') != 'user':
            continue
        content = msg.get('content', '')
        if isinstance(content, str):
            return content[:500]
        if isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and block.get('type') == 'text':
                    return block.get('text', '')[:500]
    return ''


def _extract_output_text(content: list) -> str:
    """Extract text from Anthropic content blocks (truncated to 1000 chars)."""
    if not content:
        return ''
    parts = []
    for block in content:
        block_type = getattr(block, 'type', None) if not isinstance(block, dict) \
            else block.get('type')
        if block_type == 'text':
            text = getattr(block, 'text', None) if not isinstance(block, dict) \
                else block.get('text', '')
            if text:
                parts.append(text)
    return ' '.join(parts)[:1000]


def _extract_tool_calls(content: list) -> List[Dict]:
    """Extract tool_use blocks from Anthropic content."""
    tools = []
    if not content:
        return tools
    for block in content:
        block_type = getattr(block, 'type', None) if not isinstance(block, dict) \
            else block.get('type')
        if block_type == 'tool_use':
            if isinstance(block, dict):
                name = block.get('name', 'unknown')
                inp = block.get('input', {})
            else:
                name = getattr(block, 'name', 'unknown')
                inp = getattr(block, 'input', {})
            tools.append({
                'tool_name': name,
                'input': str(inp)[:500],
                'agent_name': 'anthropic',
                'status': 'success',
            })
    return tools
