"""
AutoGen integration for Visibe (v0.7.5+).

Tracks multi-agent conversations, LLM calls, tool usage, and costs
by wrapping the model client's create() method.
"""
import logging
import time
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from ..utils import calculate_cost
from .base import TraceSummary

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

# Check if AutoGen is available
try:
    from autogen_agentchat.agents import AssistantAgent
    AUTOGEN_AVAILABLE = True
except ImportError:
    AUTOGEN_AVAILABLE = False


class AutoGenIntegration:
    """
    AutoGen integration for Visibe observability (v0.7.5+).

    Tracks LLM calls, tool usage, and token costs by wrapping the model client.

    Usage:
        from visibe import Visibe
        from visibe.integrations.autogen import AutoGenIntegration
        from autogen_agentchat.agents import AssistantAgent
        from autogen_ext.models.openai import OpenAIChatCompletionClient
        import asyncio

        obs = Visibe(api_key="sk_live_abc123")
        tracker = AutoGenIntegration(obs)

        model_client = OpenAIChatCompletionClient(
            model="gpt-4o-mini",
            api_key="your-openai-key"
        )

        assistant = AssistantAgent("assistant", model_client=model_client)

        async def main():
            with tracker.track(model_client=model_client, name="my-workflow"):
                result = await assistant.run(task="Hello!")

        asyncio.run(main())
    """

    def __init__(self, client):
        """
        Initialize AutoGen integration.

        Args:
            client: Visibe client instance

        Raises:
            ImportError: If AutoGen is not installed
        """
        if not AUTOGEN_AVAILABLE:
            raise ImportError(
                "AutoGen integration requires autogen-agentchat to be installed.\n"
                "Install with: pip install autogen-agentchat autogen-ext[openai]\n"
                "Docs: https://docs.visibe.ai/integrations/autogen"
            )

        self.client = client
        self._instrumented_clients = {}

    def instrument(self, model_client: Any, name: Optional[str] = None,
                   agents: Optional[List[Any]] = None):
        """Instrument a model client so every call sends its own trace."""
        client_id = id(model_client)
        if client_id in self._instrumented_clients:
            if not name:
                logger.debug("Model client already instrumented — skipping")
                return
            logger.debug(f"Model client already instrumented, updating trace name to: {name}")
            return

        original_create = getattr(model_client, 'create', None)
        original_create_stream = getattr(model_client, 'create_stream', None)

        self._instrumented_clients[client_id] = {
            'create': original_create,
            'create_stream': original_create_stream,
            'name': name,
            'agents': agents or [],
        }

        integration = self
        trace_name = name or "autogen"

        def _init_tracker() -> '_AutoGenTracker':
            tracker = _AutoGenTracker(
                integration.client, model_client, trace_name, agents
            )
            tracker.start_time = datetime.now(timezone.utc)

            tracker.client.create_trace({
                "trace_id": tracker.trace_id,
                "name": trace_name,
                "framework": "autogen",
                "started_at": tracker.start_time.isoformat(),
                "model": tracker._get_model_name(),
            })

            tracker._start_usage = tracker._get_current_usage()

            for agent in tracker.agents:
                agent_name = getattr(agent, 'name', 'unknown')
                step = {
                    "type": "agent_start",
                    "step_id": tracker._next_step_id(),
                    "timestamp": tracker.start_time.isoformat(),
                    "agent_name": agent_name,
                    "description": f"Agent started: {agent_name}",
                }
                tracker.steps.append(step)
                tracker._queue_span(step)

            return tracker

        if original_create is not None:
            async def wrapped_create(*args, **kwargs):
                tracker = _init_tracker()
                try:
                    wrapped = tracker._wrap_create(original_create)
                    return await wrapped(*args, **kwargs)
                finally:
                    tracker.stop()

            model_client.create = wrapped_create

        if original_create_stream is not None:
            async def wrapped_create_stream(*args, **kwargs):
                tracker = _init_tracker()
                try:
                    wrapped = tracker._wrap_create_stream(original_create_stream)
                    async for chunk in wrapped(*args, **kwargs):
                        yield chunk
                finally:
                    tracker.stop()

            model_client.create_stream = wrapped_create_stream

    def uninstrument(self, model_client: Any):
        """Restore the original model client methods."""
        client_id = id(model_client)
        entry = self._instrumented_clients.pop(client_id, None)
        if entry is None:
            logger.warning("Model client was not instrumented")
            return

        if entry.get('create') is not None:
            model_client.create = entry['create']
        if entry.get('create_stream') is not None:
            model_client.create_stream = entry['create_stream']

    @contextmanager
    def track(self, model_client: Any, name: str = "autogen-conversation",
              agents: Optional[List[Any]] = None):
        """
        Track an AutoGen conversation by wrapping the model client.

        Args:
            model_client: ChatCompletionClient instance (e.g. OpenAIChatCompletionClient)
            name: Name for this trace
            agents: Optional list of agents for metadata

        Yields:
            None

        Example:
            >>> with tracker.track(model_client=model_client, name="support-chat"):
            ...     result = await assistant.run(task="Help me")
        """
        tracker = _AutoGenTracker(self.client, model_client, name, agents)
        tracker.start()
        try:
            yield
        finally:
            tracker.stop()


class _AutoGenTracker:
    """Internal tracker for AutoGen conversations."""

    def __init__(self, client, model_client: Any, name: str,
                 agents: Optional[List[Any]] = None):
        self.client = client
        self.model_client = model_client
        self.name = name or "autogen"
        self.agents = agents or []
        self.trace_id = str(uuid.uuid4())

        # Tracking data
        self.start_time = None
        self.end_time = None
        self.llm_calls = []
        self.steps = []
        self.tools_used = []
        self.total_tokens = 0
        self.total_cost = 0.0
        self.errors = []
        self._step_counter = 0

        # Pending tool calls (waiting for results)
        self._pending_tool_calls = {}

        # Map system prompt signature to agent name
        self._agent_signatures = {}

        # Task info (extracted from first user message)
        self._task_description = None

        # Original methods (for restoration)
        self._original_create = None
        self._original_create_stream = None

        self._build_agent_signatures()

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

    def _default_agent_name(self) -> str:
        """Get the default agent name from the agents list."""
        if self.agents:
            return getattr(self.agents[0], 'name', 'autogen')
        return 'autogen'

    def _build_agent_signatures(self) -> None:
        """Build a signature map from agent system prompts."""
        if not self.agents:
            return

        for agent in self.agents:
            name = getattr(agent, 'name', None)
            system_message = getattr(agent, 'system_message', None)
            if system_message is None:
                system_message = getattr(agent, 'system_prompt', None)
            if name and system_message:
                signature = str(system_message)[:200]
                self._agent_signatures[signature] = name

    def _detect_calling_agent(self, messages) -> str:
        """Detect which agent initiated this LLM call from system prompts."""
        if not messages or not self._agent_signatures:
            return self._default_agent_name()

        for msg in messages:
            if isinstance(msg, dict):
                role = msg.get('role', '')
                if isinstance(role, str) and role.lower() == 'system':
                    content = msg.get('content', '')
                    signature = str(content)[:200]
                    if signature in self._agent_signatures:
                        return self._agent_signatures[signature]
                continue

            msg_type = getattr(msg, 'type', '')
            source = getattr(msg, 'source', '')
            if (
                isinstance(msg_type, str) and 'system' in msg_type.lower()
            ) or (
                isinstance(source, str) and source.lower() == 'system'
            ):
                content = getattr(msg, 'content', '')
                signature = str(content)[:200]
                if signature in self._agent_signatures:
                    return self._agent_signatures[signature]

        return self._default_agent_name()

    def _queue_span(self, span_data: Dict[str, Any]) -> None:
        """Queue a span to the backend. Normalizes step_id -> span_id for API."""
        span = dict(span_data)
        if "step_id" in span and "span_id" not in span:
            span["span_id"] = span.pop("step_id")
        self.client.queue_span(self.trace_id, span)

    def start(self):
        """Start tracking by wrapping model client methods."""
        self.start_time = datetime.now(timezone.utc)

        # Create trace header on backend (span-based API)
        self.client.create_trace({
            "trace_id": self.trace_id,
            "name": self.name,
            "framework": "autogen",
            "started_at": self.start_time.isoformat(),
            "model": self._get_model_name(),
        })

        # Snapshot initial usage from model client (reliable fallback)
        self._start_usage = self._get_current_usage()

        # Record agent start steps and queue them
        for agent in self.agents:
            agent_name = getattr(agent, 'name', 'unknown')
            step = {
                "type": "agent_start",
                "step_id": self._next_step_id(),
                "timestamp": self.start_time.isoformat(),
                "agent_name": agent_name,
                "description": f"Agent started: {agent_name}",
            }
            self.steps.append(step)
            self._queue_span(step)

        # Wrap create()
        if hasattr(self.model_client, 'create'):
            self._original_create = self.model_client.create
            self.model_client.create = self._wrap_create(self._original_create)
            logger.debug("Wrapped model_client.create()")

        # Wrap create_stream()
        if hasattr(self.model_client, 'create_stream'):
            self._original_create_stream = self.model_client.create_stream
            self.model_client.create_stream = self._wrap_create_stream(
                self._original_create_stream
            )
            logger.debug("Wrapped model_client.create_stream()")

    def stop(self):
        """Stop tracking and send trace to backend."""
        self.end_time = datetime.now(timezone.utc)

        # Restore original methods
        if self._original_create:
            self.model_client.create = self._original_create
        if self._original_create_stream:
            self.model_client.create_stream = self._original_create_stream

        # Fallback: if per-call wrapping didn't capture anything,
        # use model client's built-in usage accounting
        if not self.llm_calls:
            end_usage = self._get_current_usage()
            delta_prompt = end_usage[0] - self._start_usage[0]
            delta_completion = end_usage[1] - self._start_usage[1]

            if delta_prompt > 0 or delta_completion > 0:
                model = self._get_model_name()
                cost = calculate_cost(model, delta_prompt, delta_completion)
                duration_sec = (
                    self.end_time - self.start_time
                ).total_seconds()

                agent_name = self._default_agent_name()
                self.llm_calls.append({
                    "agent_name": agent_name,
                    "model": model,
                    "input_tokens": delta_prompt,
                    "output_tokens": delta_completion,
                    "total_tokens": delta_prompt + delta_completion,
                    "cost": cost,
                    "duration": duration_sec
                })
                self.total_tokens = delta_prompt + delta_completion
                self.total_cost = cost

                step = {
                    "type": "llm_call",
                    "step_id": self._next_step_id(),
                    "timestamp": self.start_time.isoformat(),
                    "agent_name": agent_name,
                    "model": model,
                    "description": f"LLM Call using {model}",
                    "input_tokens": delta_prompt,
                    "output_tokens": delta_completion,
                    "cost": cost,
                    "duration_ms": int(duration_sec * 1000),
                    "status": "success",
                }
                self.steps.append(step)
                self._queue_span(step)

                logger.debug(
                    f"Fallback tracking: {delta_prompt} input, "
                    f"{delta_completion} output tokens"
                )

        # Calculate duration
        duration_ms = int(
            (self.end_time - self.start_time).total_seconds() * 1000
        )

        # Build agent names list
        agent_names = []
        for agent in self.agents:
            agent_names.append(getattr(agent, 'name', 'unknown'))

        # Build cost breakdown per agent (aggregate by agent name)
        agent_costs = {}
        for call in self.llm_calls:
            name = call["agent_name"]
            if name not in agent_costs:
                agent_costs[name] = {
                    "agent_name": name,
                    "model": call["model"],
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "cost": 0.0,
                }
            agent_costs[name]["input_tokens"] += call["input_tokens"]
            agent_costs[name]["output_tokens"] += call["output_tokens"]
            agent_costs[name]["cost"] += call["cost"]

        cost_breakdown = []
        for entry in agent_costs.values():
            entry["cost"] = round(entry["cost"], 6)
            cost_breakdown.append(entry)

        # Build total_tokens (matches CrewAI/LangChain format)
        total_tokens = sum(c["total_tokens"] for c in self.llm_calls)
        total_input = sum(c["input_tokens"] for c in self.llm_calls)
        total_output = sum(c["output_tokens"] for c in self.llm_calls)

        # Determine the actual prompt text (clean task, not raw messages)
        prompt_text = self._task_description or self.name

        # Complete trace with summary (span-based API — spans already queued)
        status = "failed" if self.errors else "completed"
        model = self.llm_calls[0]['model'] if self.llm_calls else None
        summary = {
            "status": status,
            "ended_at": self.end_time.isoformat(),
            "duration_ms": duration_ms,
            "prompt": prompt_text,
            "total_cost": round(self.total_cost, 6),
            "total_tokens": total_tokens,
            "total_input_tokens": total_input,
            "total_output_tokens": total_output,
            "model": model,
        }

        success = self.client.complete_trace(self.trace_id, summary)

        if success:
            logger.info(f"AutoGen trace completed: {self.trace_id}")
        else:
            logger.error(f"Failed to complete AutoGen trace: {self.trace_id}")

        # Print TraceSummary to stdout for developer feedback
        agent_names = [getattr(a, 'name', 'unknown') for a in self.agents]
        trace_summary = TraceSummary(
            name=self.name,
            status=status,
            llm_calls=len(self.llm_calls),
            tool_calls=len(self.tools_used),
            total_tokens=total_tokens,
            total_cost=round(self.total_cost, 6),
            duration_s=round(duration_ms / 1000, 1),
            agents=agent_names,
            sent=success,
        )
        print(f"[Visibe] {trace_summary}")

    def _wrap_create(self, original_method):
        """Wrap model_client.create() to track LLM calls and tool usage."""
        tracker = self

        async def wrapped_create(*args, **kwargs):
            call_start = time.time()
            call_timestamp = datetime.now(timezone.utc)

            # Extract input messages for metadata
            input_text = ""
            agent_name = tracker._default_agent_name()
            if args:
                messages = args[0]
                input_text = tracker._format_messages(messages)
                agent_name = tracker._detect_calling_agent(messages)
                # Extract task description from first user message
                tracker._extract_task_description(messages)
                # Check for tool execution results in input messages
                tracker._process_tool_results(
                    messages, call_timestamp, agent_name=agent_name
                )

            try:
                result = await original_method(*args, **kwargs)

                duration_sec = time.time() - call_start
                duration_ms = int(duration_sec * 1000)

                # Extract usage from CreateResult
                if hasattr(result, 'usage') and result.usage:
                    input_tokens = result.usage.prompt_tokens
                    output_tokens = result.usage.completion_tokens
                else:
                    input_tokens = 0
                    output_tokens = 0

                # Extract output text and detect tool calls
                output_text = ""
                if hasattr(result, 'content'):
                    if isinstance(result.content, str):
                        output_text = result.content
                    elif isinstance(result.content, list):
                        # Function/tool calls from LLM
                        tracker._process_tool_requests(
                            result.content, call_timestamp,
                            input_tokens, output_tokens,
                            agent_name=agent_name
                        )
                        output_text = tracker._format_tool_calls(
                            result.content
                        )

                model = tracker._get_model_name()
                cost = calculate_cost(model, input_tokens, output_tokens)

                # Track LLM call
                tracker._track_llm_call(
                    model=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    duration=duration_sec,
                    agent_name=agent_name
                )

                # Record step with full detail and queue span
                step = {
                    "type": "llm_call",
                    "step_id": tracker._next_step_id(),
                    "timestamp": call_timestamp.isoformat(),
                    "agent_name": agent_name,
                    "model": model,
                    "description": f"LLM Call using {model}",
                    "input_text": input_text[:1000],
                    "output_text": output_text[:1000],
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "cost": cost,
                    "duration_ms": duration_ms,
                    "status": "success",
                }
                tracker.steps.append(step)
                tracker._queue_span(step)

                return result

            except Exception as e:
                duration_ms = int((time.time() - call_start) * 1000)

                tracker.errors.append({
                    "timestamp": call_timestamp.isoformat(),
                    "agent_name": agent_name,
                    "error_type": type(e).__name__,
                    "message": str(e)[:500]
                })

                step = {
                    "type": "error",
                    "step_id": tracker._next_step_id(),
                    "timestamp": call_timestamp.isoformat(),
                    "agent_name": agent_name,
                    "description": f"LLM Error: {type(e).__name__}",
                    "error_type": type(e).__name__,
                    "error_message": str(e)[:500],
                    "input_text": input_text[:1000],
                    "duration_ms": duration_ms,
                }
                tracker.steps.append(step)
                tracker._queue_span(step)

                logger.error(f"Error in model client create: {e}")
                raise

        return wrapped_create

    def _wrap_create_stream(self, original_method):
        """Wrap model_client.create_stream() to track streaming LLM calls."""
        tracker = self

        async def wrapped_create_stream(*args, **kwargs):
            call_start = time.time()
            call_timestamp = datetime.now(timezone.utc)

            input_text = ""
            agent_name = tracker._default_agent_name()
            if args:
                messages = args[0]
                input_text = tracker._format_messages(messages)
                agent_name = tracker._detect_calling_agent(messages)
                tracker._extract_task_description(messages)
                tracker._process_tool_results(
                    messages, call_timestamp, agent_name=agent_name
                )

            try:
                output_text = ""
                async for chunk in original_method(*args, **kwargs):
                    if hasattr(chunk, 'usage') and chunk.usage:
                        duration_sec = time.time() - call_start
                        duration_ms = int(duration_sec * 1000)

                        input_tokens = chunk.usage.prompt_tokens
                        output_tokens = chunk.usage.completion_tokens
                        model = tracker._get_model_name()
                        cost = calculate_cost(
                            model, input_tokens, output_tokens
                        )

                        if hasattr(chunk, 'content'):
                            if isinstance(chunk.content, str):
                                output_text = chunk.content
                            elif isinstance(chunk.content, list):
                                tracker._process_tool_requests(
                                    chunk.content, call_timestamp,
                                    input_tokens, output_tokens,
                                    agent_name=agent_name
                                )
                                output_text = tracker._format_tool_calls(
                                    chunk.content
                                )

                        tracker._track_llm_call(
                            model=model,
                            input_tokens=input_tokens,
                            output_tokens=output_tokens,
                            duration=duration_sec,
                            agent_name=agent_name
                        )

                        step = {
                            "type": "llm_call",
                            "step_id": tracker._next_step_id(),
                            "timestamp": call_timestamp.isoformat(),
                            "agent_name": agent_name,
                            "model": model,
                            "description": (
                                f"LLM Call (streaming) using {model}"
                            ),
                            "input_text": input_text[:1000],
                            "output_text": output_text[:1000],
                            "input_tokens": input_tokens,
                            "output_tokens": output_tokens,
                            "cost": cost,
                            "duration_ms": duration_ms,
                            "status": "success",
                        }
                        tracker.steps.append(step)
                        tracker._queue_span(step)

                    elif isinstance(chunk, str):
                        output_text += chunk

                    yield chunk

            except Exception as e:
                duration_ms = int((time.time() - call_start) * 1000)
                tracker.errors.append({
                    "timestamp": call_timestamp.isoformat(),
                    "agent_name": agent_name,
                    "error_type": type(e).__name__,
                    "message": str(e)[:500]
                })

                step = {
                    "type": "error",
                    "step_id": tracker._next_step_id(),
                    "timestamp": call_timestamp.isoformat(),
                    "agent_name": agent_name,
                    "description": f"LLM Error: {type(e).__name__}",
                    "error_type": type(e).__name__,
                    "error_message": str(e)[:500],
                    "input_text": input_text[:1000],
                    "duration_ms": duration_ms,
                }
                tracker.steps.append(step)
                tracker._queue_span(step)

                logger.error(f"Error in model client create_stream: {e}")
                raise

        return wrapped_create_stream

    # ------------------------------------------------------------------
    # Tool tracking
    # ------------------------------------------------------------------

    def _process_tool_requests(self, content_list, timestamp,
                               input_tokens, output_tokens,
                               agent_name: Optional[str] = None):
        """Record tool call requests from LLM output (FunctionCall objects)."""
        agent_name = agent_name or self._default_agent_name()
        for item in content_list:
            if hasattr(item, 'name') and hasattr(item, 'arguments'):
                call_id = getattr(item, 'id', str(uuid.uuid4().hex[:8]))
                self._pending_tool_calls[call_id] = {
                    "tool_name": item.name,
                    "input": str(item.arguments)[:500],
                    "request_time": timestamp,
                    "agent_name": agent_name,
                    "total_tokens": input_tokens + output_tokens,
                    "estimated_cost": round(calculate_cost(
                        self._get_model_name(), input_tokens, output_tokens
                    ), 6),
                }


    def _process_tool_results(self, messages, current_timestamp,
                              agent_name: Optional[str] = None):
        """Extract tool execution results from input messages."""
        default_agent = agent_name or self._default_agent_name()
        for msg in messages:
            # FunctionExecutionResultMessage has content = List[FunctionExecutionResult]
            if not hasattr(msg, 'content') or not isinstance(msg.content, list):
                continue

            for result in msg.content:
                if not hasattr(result, 'call_id'):
                    continue

                call_id = result.call_id
                tool_name = getattr(result, 'name', 'unknown')
                output = getattr(result, 'content', '')
                is_error = getattr(result, 'is_error', False)

                # Match with pending request — use its agent, else fall back
                pending = self._pending_tool_calls.pop(call_id, None)
                result_agent = (
                    pending["agent_name"]
                    if pending and pending.get("agent_name")
                    else default_agent
                )
                request_time = (
                    pending["request_time"] if pending else current_timestamp
                )
                execution_ms = int(
                    (current_timestamp - request_time).total_seconds() * 1000
                )

                status = "failed" if is_error else "success"

                tool_entry = {
                    "tool_name": tool_name,
                    "agent_name": result_agent,
                    "execution_time_ms": execution_ms,
                    "status": status,
                    "input": (
                        pending["input"] if pending
                        else ""
                    ),
                    "output": str(output)[:500],
                }
                if pending:
                    tool_entry["total_tokens"] = pending["total_tokens"]
                    tool_entry["estimated_cost"] = pending["estimated_cost"]

                self.tools_used.append(tool_entry)

                # Record step and queue span
                step = {
                    "type": "tool_call",
                    "step_id": self._next_step_id(),
                    "timestamp": current_timestamp.isoformat(),
                    "tool_name": tool_name,
                    "agent_name": result_agent,
                    "description": f"Tool: {tool_name}",
                    "status": status,
                    "input_text": (
                        pending["input"] if pending
                        else ""
                    ),
                    "output_text": str(output)[:500],
                    "duration_ms": execution_ms,
                }
                self.steps.append(step)
                self._queue_span(step)

                # Track tool errors and queue error span
                if is_error:
                    self.errors.append({
                        "timestamp": current_timestamp.isoformat(),
                        "agent_name": result_agent,
                        "tool_name": tool_name,
                        "error_type": "TOOL_ERROR",
                        "message": str(output)[:500],
                    })
                    error_step = {
                        "type": "error",
                        "step_id": self._next_step_id(),
                        "timestamp": current_timestamp.isoformat(),
                        "agent_name": result_agent,
                        "tool_name": tool_name,
                        "description": f"Tool Error: {tool_name}",
                        "error_type": "TOOL_ERROR",
                        "error_message": str(output)[:500],
                    }
                    self.steps.append(error_step)
                    self._queue_span(error_step)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _extract_task_description(self, messages):
        """Extract task description from the first user message."""
        if self._task_description:
            return
        for msg in messages:
            msg_type = getattr(msg, 'type', type(msg).__name__)
            if 'user' in msg_type.lower():
                content = getattr(msg, 'content', '')
                if content:
                    self._task_description = str(content)[:500]
                    return

    def _track_llm_call(self, model: str, input_tokens: int,
                        output_tokens: int, duration: float,
                        agent_name: Optional[str] = None):
        """Track a single LLM call."""
        total_tokens = input_tokens + output_tokens
        cost = calculate_cost(model, input_tokens, output_tokens)
        agent_name = agent_name or self._default_agent_name()

        self.llm_calls.append({
            "agent_name": agent_name,
            "model": model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": total_tokens,
            "cost": cost,
            "duration": duration
        })

        self.total_tokens += total_tokens
        self.total_cost += cost

        logger.debug(
            f"LLM call: model={model}, "
            f"tokens={total_tokens}, cost=${cost:.6f}"
        )

    def _format_messages(self, messages) -> str:
        """Format LLM messages into a readable string for metadata."""
        parts = []
        try:
            if not messages:
                return ""
            for msg in messages:
                role = getattr(msg, 'type', type(msg).__name__)
                content = ""
                if hasattr(msg, 'content'):
                    if isinstance(msg.content, list):
                        # FunctionExecutionResultMessage or similar
                        content = str(msg.content)
                    else:
                        content = str(msg.content)
                elif hasattr(msg, 'text'):
                    content = str(msg.text)
                else:
                    content = str(msg)
                parts.append(f"[{role}]: {content}")
        except Exception:
            return str(messages)[:1000]
        return "\n".join(parts)

    def _format_tool_calls(self, content_list) -> str:
        """Format tool call list into readable string."""
        parts = []
        for item in content_list:
            if hasattr(item, 'name') and hasattr(item, 'arguments'):
                parts.append(
                    f"tool_name='{item.name}' "
                    f"arguments={item.arguments}"
                )
        return "\n".join(parts) if parts else str(content_list)

    def _get_current_usage(self) -> tuple:
        """Get current cumulative usage from the model client."""
        try:
            usage = self.model_client.total_usage()
            return (usage.prompt_tokens, usage.completion_tokens)
        except Exception:
            return (0, 0)

    def _get_model_name(self) -> str:
        """Extract model name from the model client."""
        # Try exact model name from create args (most precise)
        try:
            model = self.model_client._create_args.get('model')
            if model and isinstance(model, str):
                return model
        except Exception:
            pass

        # Fall back to model_info family
        try:
            info = self.model_client.model_info
            if isinstance(info, dict) and 'family' in info:
                return str(info['family'])
        except Exception:
            pass

        return 'unknown'
