"""
Purpose: Unified LLM provider abstraction with factory pattern for OpenAI, Anthropic, Gemini, Groq, Grok, Mistral, OpenRouter, and OpenOnion
LLM-Note:
  Dependencies: imports from [abc, typing, dataclasses, json, os, base64, openai, anthropic, requests, pathlib, yaml, pydantic, .usage, .exceptions] | imported by [agent.py, llm_do.py, conftest.py] | tested by [tests/test_llm.py, tests/test_llm_do.py, tests/test_real_*.py, tests/test_billing_error_agent.py]
  Data flow: Agent/llm_do calls create_llm(model, api_key) → factory routes to provider class → Provider.__init__() validates API key → Agent calls complete(messages, tools) OR structured_complete(messages, output_schema) → provider converts to native format → calls API → parses response → returns LLMResponse(content, tool_calls, raw_response) OR Pydantic model instance
  State/Effects: reads environment variables (OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY/GOOGLE_API_KEY, GROQ_API_KEY, OPENROUTER_API_KEY, XAI_API_KEY, OPENONION_API_KEY) | reads OPENONION_API_KEY from env / .env / ~/.co/keys.env | makes HTTP requests to LLM APIs | no caching or persistence
  Integration: exposes create_llm(model, api_key), LLM abstract base class, OpenAILLM, AnthropicLLM, GeminiLLM, GroqLLM, GrokLLM, OpenRouterLLM, OpenOnionLLM, LLMResponse, ToolCall dataclasses | providers implement complete() and structured_complete() | OpenAI message format is lingua franca | tool calling uses OpenAI schema converted per-provider
  Performance: stateless (no caching) | synchronous (no streaming) | default max_tokens=8192 for Anthropic (required) | each call hits API
  Errors: raises ValueError for missing API keys, unknown models, invalid parameters | provider-specific errors bubble up (openai.APIError, anthropic.APIError, etc.) | OpenOnionLLM transforms 402 errors to InsufficientCreditsError with formatted message and typed attributes | Pydantic ValidationError for invalid structured output

Unified LLM provider abstraction layer for ConnectOnion framework.

This module provides a consistent interface for interacting with multiple LLM providers
(OpenAI, Anthropic, Google Gemini, Groq, Grok, Mistral, OpenRouter, and ConnectOnion managed keys)
through a common API.

Architecture Overview
--------------------
The module follows a factory pattern with provider-specific implementations:

1. **Abstract Base Class (LLM)**:
   - Defines the contract all providers must implement
   - Two core methods: complete() for text, structured_complete() for Pydantic models
   - Ensures consistent interface across all providers

2. **Provider Implementations**:
   - OpenAILLM: Native OpenAI API with responses.parse() for structured output
   - AnthropicLLM: Claude API with tool calling workaround for structured output
   - GeminiLLM: Google Gemini with response_schema for structured output
   - GroqLLM: Groq via OpenAI-compatible endpoint
   - GrokLLM: xAI Grok via OpenAI-compatible endpoint
   - MistralLLM: Mistral AI via OpenAI-compatible endpoint
   - OpenRouterLLM: OpenRouter via OpenAI-compatible endpoint
   - OpenOnionLLM: Managed keys using OpenAI-compatible proxy endpoint

3. **Factory Function (create_llm)**:
   - Routes model names to appropriate providers
   - Handles API key initialization
   - Returns configured provider instance

Key Design Decisions
-------------------
- **Structured Output**: Each provider uses its native structured output API when available
  * OpenAI: responses.parse() with text_format parameter
  * Anthropic: Forced tool calling with schema validation
  * Gemini: response_schema with JSON MIME type
  * OpenOnion: Proxies to OpenAI with fallback

- **Tool Calling**: OpenAI format used as the common schema, converted per-provider
  * All providers return ToolCall dataclasses with (name, arguments, id)
  * Enables consistent agent behavior across providers

- **Message Format**: OpenAI's message format (role/content) is the lingua franca
  * Providers convert to their native format internally
  * Simplifies Agent integration

- **Parameter Passing**: **kwargs pattern for runtime parameters
  * temperature, max_tokens, etc. flow through to provider APIs
  * Allows provider-specific features without bloating base interface

Data Flow
---------
Agent/llm_do → create_llm(model) → Provider.__init__(api_key)
           ↓
Provider.complete(messages, tools, **kwargs)
           ↓
Convert messages → Call native API → Parse response
           ↓
Return LLMResponse(content, tool_calls, raw_response)

For structured output:
Provider.structured_complete(messages, output_schema, **kwargs)
           ↓
Use native structured API → Validate with Pydantic
           ↓
Return Pydantic model instance

Dependencies
-----------
- openai: OpenAI and OpenOnion provider implementations
- anthropic: Claude provider implementation
- google.generativeai: Gemini provider implementation
- pydantic: Structured output validation
- requests: OpenOnion authentication checks
- yaml: OpenOnion config file parsing

Integration Points
-----------------
Imported by:
  - agent.py: Agent class uses LLM for reasoning
  - llm_do.py: One-shot function uses LLM directly
  - conftest.py: Test fixtures

Tested by:
  - tests/test_llm.py: Unit tests with mocked APIs
  - tests/test_llm_do.py: Integration tests
  - tests/test_real_*.py: Real API integration tests

Environment Variables
--------------------
Required (pick one):
  - OPENAI_API_KEY: For OpenAI models
  - ANTHROPIC_API_KEY: For Claude models
  - GEMINI_API_KEY or GOOGLE_API_KEY: For Gemini models
  - OPENONION_API_KEY: For co/ managed keys (from .env or ~/.co/keys.env)
  - GROQ_API_KEY: For groq/ prefixed models
  - OPENROUTER_API_KEY: For openrouter/ prefixed models
  - OPENROUTER_HTTP_REFERER: Optional attribution header for OpenRouter
  - OPENROUTER_X_TITLE: Optional app title header for OpenRouter
  - XAI_API_KEY: For grok/ prefixed models (xAI)
  - MISTRAL_API_KEY: For mistral/ prefixed models

Optional:
  - OPENONION_DEV: Use localhost:8000 for OpenOnion (development)
  - ENVIRONMENT=development: Same as OPENONION_DEV

Error Handling
-------------
- ValueError: Missing API keys, unknown models, invalid parameters
- Provider-specific errors: Bubble up from native SDKs (openai.APIError, etc.)
- Structured output errors: Pydantic ValidationError if response doesn't match schema

Performance Considerations
-------------------------
- Default max_tokens: 8192 for Anthropic (required), configurable for others
- No caching: Each call is stateless (Agent maintains conversation history)
- No streaming: Currently synchronous only (streaming planned for future)

Example Usage
------------
Basic completion:
    >>> from connectonion.llm import create_llm
    >>> llm = create_llm(model="gpt-4o-mini")
    >>> response = llm.complete([{"role": "user", "content": "Hello"}])
    >>> print(response.content)

Structured output:
    >>> from pydantic import BaseModel
    >>> class Answer(BaseModel):
    ...     value: int
    >>> llm = create_llm(model="gpt-4o-mini")
    >>> result = llm.structured_complete(
    ...     [{"role": "user", "content": "What is 2+2?"}],
    ...     Answer
    ... )
    >>> print(result.value)  # 4

With tools:
    >>> tools = [{"name": "search", "description": "Search the web", "parameters": {...}}]
    >>> response = llm.complete(messages, tools=tools)
    >>> if response.tool_calls:
    ...     print(response.tool_calls[0].name)  # "search"
"""

from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional, Type
from dataclasses import dataclass
import json
import os
import base64
import logging

logger = logging.getLogger(__name__)
import openai
import anthropic
# google-genai not needed - using OpenAI-compatible endpoint instead
import requests
from pathlib import Path
from pydantic import BaseModel


@dataclass
class ToolCall:
    """Represents a tool call from the LLM.

    Attributes:
        name: The function name to call
        arguments: Dict of arguments to pass to the function
        id: Unique identifier for this tool call
        extra_content: Provider-specific metadata (e.g., Gemini 3 thought_signature).
            Must be echoed back in the assistant message for models that require it.
            See: https://ai.google.dev/gemini-api/docs/thinking#openai-sdk
    """
    name: str
    arguments: Dict[str, Any]
    id: str
    extra_content: Optional[Dict[str, Any]] = None


# Import TokenUsage from usage module
from .usage import TokenUsage, calculate_cost
from .exceptions import InsufficientCreditsError, LLMConnectionError, ProviderServiceError


@dataclass
class LLMResponse:
    """Response from LLM including content and tool calls."""
    content: Optional[str]
    tool_calls: List[ToolCall]
    raw_response: Any
    usage: Optional[TokenUsage] = None


class LLM(ABC):
    """Abstract base class for LLM providers."""

    @abstractmethod
    def complete(self, messages: List[Dict[str, str]], tools: Optional[List[Dict[str, Any]]] = None) -> LLMResponse:
        """Complete a conversation with optional tool support."""
        pass

    @abstractmethod
    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel]) -> BaseModel:
        """Get structured Pydantic output matching the schema.

        Args:
            messages: Conversation messages in OpenAI format
            output_schema: Pydantic model class defining the expected output structure

        Returns:
            Instance of output_schema with parsed and validated data

        Raises:
            ValueError: If the LLM fails to generate valid structured output
        """
        pass


class OpenAILLM(LLM):
    """OpenAI LLM implementation."""
    
    def __init__(self, api_key: Optional[str] = None, model: str = "o4-mini", **kwargs):
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError("OpenAI API key required. Set OPENAI_API_KEY environment variable or pass api_key parameter.")
        
        self.client = openai.OpenAI(api_key=self.api_key)
        self.model = model
    
    def complete(self, messages: List[Dict[str, str]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation with optional tool support."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs  # Pass through user kwargs (max_tokens, temperature, etc.)
        }

        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        response = self.client.chat.completions.create(**api_kwargs)
        message = response.choices[0].message

        # Parse tool calls if present
        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments),
                    id=tc.id
                ))

        # Extract token usage
        input_tokens = response.usage.prompt_tokens
        output_tokens = response.usage.completion_tokens
        cached_tokens = response.usage.prompt_tokens_details.cached_tokens if response.usage.prompt_tokens_details else 0
        cost = calculate_cost(self.model, input_tokens, output_tokens, cached_tokens)

        return LLMResponse(
            content=message.content,
            tool_calls=tool_calls,
            raw_response=response,
            usage=TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cached_tokens=cached_tokens,
                cost=cost,
            ),
        )

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using OpenAI's native responses.parse API.

        Uses the new OpenAI responses.parse() endpoint with text_format parameter
        for guaranteed schema adherence.
        """
        response = self.client.responses.parse(
            model=self.model,
            input=messages,
            text_format=output_schema,
            **kwargs  # Pass through temperature, max_tokens, etc.
        )

        # Handle edge cases
        if response.status == "incomplete":
            if response.incomplete_details.reason == "max_output_tokens":
                raise ValueError("Response incomplete: maximum output tokens reached")
            elif response.incomplete_details.reason == "content_filter":
                raise ValueError("Response incomplete: content filtered")

        # Check for refusal
        if response.output and len(response.output) > 0:
            first_content = response.output[0].content[0] if response.output[0].content else None
            if first_content and hasattr(first_content, 'type') and first_content.type == "refusal":
                raise ValueError(f"Model refused to respond: {first_content.refusal}")

        # Return the parsed Pydantic object
        return response.output_parsed


class AnthropicLLM(LLM):
    """Anthropic Claude LLM implementation."""
    
    def __init__(self, api_key: Optional[str] = None, model: str = "claude-3-5-sonnet-20241022", max_tokens: int = 8192, **kwargs):
        self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
        if not self.api_key:
            raise ValueError("Anthropic API key required. Set ANTHROPIC_API_KEY environment variable or pass api_key parameter.")

        self.client = anthropic.Anthropic(api_key=self.api_key)
        self.model = model
        self.max_tokens = max_tokens  # Anthropic requires max_tokens (default 8192)
    
    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation with optional tool support."""
        # Convert messages to Anthropic format
        anthropic_messages = self._convert_messages(messages)

        api_kwargs = {
            "model": self.model,
            "messages": anthropic_messages,
            "max_tokens": self.max_tokens,  # Required by Anthropic
            **kwargs  # User can override max_tokens via kwargs
        }

        # Add tools if provided
        if tools:
            api_kwargs["tools"] = self._convert_tools(tools)

        response = self.client.messages.create(**api_kwargs)
        
        # Parse tool calls if present
        tool_calls = []
        content = ""
        
        for block in response.content:
            if block.type == "text":
                content += block.text
            elif block.type == "tool_use":
                tool_calls.append(ToolCall(
                    name=block.name,
                    arguments=block.input,
                    id=block.id
                ))

        # Extract token usage - Anthropic uses input_tokens/output_tokens
        input_tokens = response.usage.input_tokens
        output_tokens = response.usage.output_tokens
        cached_tokens = getattr(response.usage, 'cache_read_input_tokens', 0) or 0
        cache_write_tokens = getattr(response.usage, 'cache_creation_input_tokens', 0) or 0
        cost = calculate_cost(self.model, input_tokens, output_tokens, cached_tokens, cache_write_tokens)

        return LLMResponse(
            content=content if content else None,
            tool_calls=tool_calls,
            raw_response=response,
            usage=TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cached_tokens=cached_tokens,
                cache_write_tokens=cache_write_tokens,
                cost=cost,
            ),
        )

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using tool calling method.

        Anthropic doesn't have native Pydantic support yet, so we use a tool calling
        workaround: create a dummy tool with the Pydantic schema and force its use.
        """
        # Convert messages to Anthropic format
        anthropic_messages = self._convert_messages(messages)

        # Create a tool with the Pydantic schema as input_schema
        tool = {
            "name": "return_structured_output",
            "description": "Returns the structured output based on the user's request",
            "input_schema": output_schema.model_json_schema()
        }

        # Set max_tokens with safe default
        api_kwargs = {
            "model": self.model,
            "max_tokens": self.max_tokens,
            "messages": anthropic_messages,
            "tools": [tool],
            "tool_choice": {"type": "tool", "name": "return_structured_output"},
            **kwargs  # User can override max_tokens, temperature, etc.
        }

        # Force the model to use this tool
        response = self.client.messages.create(**api_kwargs)

        # Extract structured data from tool call
        for block in response.content:
            if block.type == "tool_use" and block.name == "return_structured_output":
                # Validate and return as Pydantic model
                return output_schema.model_validate(block.input)

        raise ValueError("No structured output received from Claude")

    def _convert_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Convert OpenAI-style messages to Anthropic format."""
        anthropic_messages = []
        i = 0
        
        while i < len(messages):
            msg = messages[i]
            
            # Skip system messages (will be handled separately)
            if msg["role"] == "system":
                i += 1
                continue
            
            # Handle assistant messages with tool calls
            if msg["role"] == "assistant" and msg.get("tool_calls"):
                content_blocks = []
                if msg.get("content"):
                    content_blocks.append({
                        "type": "text",
                        "text": msg["content"]
                    })
                
                for tc in msg["tool_calls"]:
                    content_blocks.append({
                        "type": "tool_use",
                        "id": tc["id"],
                        "name": tc["function"]["name"],
                        "input": json.loads(tc["function"]["arguments"]) if isinstance(tc["function"]["arguments"], str) else tc["function"]["arguments"]
                    })
                
                anthropic_messages.append({
                    "role": "assistant",
                    "content": content_blocks
                })
                
                # Now collect all the tool responses that follow immediately
                i += 1
                tool_results = []
                while i < len(messages) and messages[i]["role"] == "tool":
                    tool_msg = messages[i]
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": tool_msg["tool_call_id"],
                        "content": tool_msg["content"]
                    })
                    i += 1
                
                # Add all tool results in a single user message
                if tool_results:
                    anthropic_messages.append({
                        "role": "user",
                        "content": tool_results
                    })
            
            # Handle tool role messages that aren't immediately after assistant tool calls
            elif msg["role"] == "tool":
                # This shouldn't happen in normal flow, but handle it just in case
                anthropic_messages.append({
                    "role": "user",
                    "content": [{
                        "type": "tool_result",
                        "tool_use_id": msg["tool_call_id"],
                        "content": msg["content"]
                    }]
                })
                i += 1
            
            # Handle user messages
            elif msg["role"] == "user":
                if isinstance(msg.get("content"), list):
                    # This is already a structured message
                    anthropic_msg = {
                        "role": "user",
                        "content": []
                    }
                    for item in msg["content"]:
                        if item.get("type") == "tool_result":
                            anthropic_msg["content"].append({
                                "type": "tool_result",
                                "tool_use_id": item["tool_call_id"],
                                "content": item["content"]
                            })
                    anthropic_messages.append(anthropic_msg)
                else:
                    # Regular text message
                    anthropic_messages.append({
                        "role": "user",
                        "content": msg["content"]
                    })
                i += 1
            
            # Handle regular assistant messages
            elif msg["role"] == "assistant":
                anthropic_messages.append({
                    "role": "assistant",
                    "content": msg["content"]
                })
                i += 1
            
            else:
                i += 1
        
        return anthropic_messages
    
    def _convert_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Convert OpenAI-style tools to Anthropic format."""
        anthropic_tools = []
        
        for tool in tools:
            # Tools already in our internal format
            anthropic_tool = {
                "name": tool["name"],
                "description": tool.get("description", ""),
                "input_schema": tool.get("parameters", {
                    "type": "object",
                    "properties": {},
                    "required": []
                })
            }
            anthropic_tools.append(anthropic_tool)
        
        return anthropic_tools


class GeminiLLM(LLM):
    """Google Gemini LLM implementation using OpenAI-compatible endpoint."""

    def __init__(self, api_key: Optional[str] = None, model: str = "gemini-2.0-flash-exp", **kwargs):
        self.api_key = api_key or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
        if not self.api_key:
            raise ValueError("Gemini API key required. Set GEMINI_API_KEY environment variable or pass api_key parameter. (GOOGLE_API_KEY is also supported for backward compatibility)")

        # Use Gemini's OpenAI-compatible endpoint
        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
        )
        self.model = model
    
    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation using Gemini's OpenAI-compatible endpoint."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }

        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        response = self.client.chat.completions.create(**api_kwargs)
        message = response.choices[0].message

        # Parse tool calls if present
        # Preserve extra_content for providers that need it (e.g., Gemini 3 thought_signature)
        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                extra = getattr(tc, 'extra_content', None)
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments,
                    id=tc.id,
                    extra_content=extra
                ))

        # Extract token usage (OpenAI-compatible format)
        usage = None
        if hasattr(response, 'usage') and response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            cached_tokens = 0
            if hasattr(response.usage, 'prompt_tokens_details') and response.usage.prompt_tokens_details:
                cached_tokens = getattr(response.usage.prompt_tokens_details, 'cached_tokens', 0) or 0
            cost = calculate_cost(self.model, input_tokens, output_tokens, cached_tokens)
            usage = TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cached_tokens=cached_tokens,
                cost=cost,
            )

        return LLMResponse(
            content=message.content,
            tool_calls=tool_calls,
            raw_response=response,
            usage=usage,
        )

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using Gemini's OpenAI-compatible endpoint with beta.chat.completions.parse."""
        completion = self.client.beta.chat.completions.parse(
            model=self.model,
            messages=messages,
            response_format=output_schema,
            **kwargs
        )
        return completion.choices[0].message.parsed


class GroqLLM(LLM):
    """Groq LLM implementation using OpenAI-compatible endpoint."""

    def __init__(self, api_key: Optional[str] = None, model: str = "groq/llama-3.3-70b-versatile", **kwargs):
        self.api_key = api_key or os.getenv("GROQ_API_KEY")
        if not self.api_key:
            raise ValueError("Groq API key required. Set GROQ_API_KEY environment variable or pass api_key parameter.")

        self.model = model.removeprefix("groq/")
        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url="https://api.groq.com/openai/v1"
        )

    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation using Groq's OpenAI-compatible endpoint."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }

        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        try:
            response = self.client.chat.completions.create(**api_kwargs)
        except openai.APIError as e:
            raise ValueError(f"Groq API Error: {str(e)}") from e

        message = response.choices[0].message

        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments,
                    id=tc.id
                ))

        usage = None
        if hasattr(response, 'usage') and response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            cost = calculate_cost(self.model, input_tokens, output_tokens)
            usage = TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost=cost,
            )

        return LLMResponse(content=message.content, tool_calls=tool_calls, raw_response=response, usage=usage)

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using JSON-mode + schema validation.

        Uses chat.completions with JSON response format for compatibility with
        OpenAI-like providers that may not support beta parse endpoints.
        """
        schema_json = json.dumps(output_schema.model_json_schema(), indent=2)
        schema_instruction = (
            "Return ONLY valid JSON (no markdown) that matches this JSON Schema:\n"
            f"{schema_json}"
        )

        structured_messages = [{"role": "system", "content": schema_instruction}, *messages]

        completion = self.client.chat.completions.create(
            model=self.model,
            messages=structured_messages,
            response_format={"type": "json_object"},
            **kwargs,
        )
        content = completion.choices[0].message.content or "{}"
        return output_schema.model_validate_json(content)


class GrokLLM(LLM):
    """Grok (xAI) LLM implementation using OpenAI-compatible endpoint."""

    def __init__(self, api_key: Optional[str] = None, model: str = "grok/grok-4", **kwargs):
        self.api_key = api_key or os.getenv("XAI_API_KEY")
        if not self.api_key:
            raise ValueError("Grok API key required. Set XAI_API_KEY environment variable or pass api_key parameter.")

        self.model = model.removeprefix("grok/")
        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url="https://api.x.ai/v1"
        )

    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation using Grok's OpenAI-compatible endpoint."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }

        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        response = self.client.chat.completions.create(**api_kwargs)
        message = response.choices[0].message

        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments,
                    id=tc.id
                ))

        usage = None
        if hasattr(response, 'usage') and response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            cost = calculate_cost(self.model, input_tokens, output_tokens)
            usage = TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost=cost,
            )

        return LLMResponse(content=message.content, tool_calls=tool_calls, raw_response=response, usage=usage)

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using JSON-mode + schema validation."""
        schema_json = json.dumps(output_schema.model_json_schema(), indent=2)
        schema_instruction = (
            "Return ONLY valid JSON (no markdown) that matches this JSON Schema:\n"
            f"{schema_json}"
        )

        structured_messages = [{"role": "system", "content": schema_instruction}, *messages]

        completion = self.client.chat.completions.create(
            model=self.model,
            messages=structured_messages,
            response_format={"type": "json_object"},
            **kwargs,
        )
        content = completion.choices[0].message.content or "{}"
        return output_schema.model_validate_json(content)


class OpenRouterLLM(LLM):
    """OpenRouter LLM implementation using OpenAI-compatible endpoint."""

    def __init__(self, api_key: Optional[str] = None, model: str = "openrouter/openai/gpt-4o-mini", **kwargs):
        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
        if not self.api_key:
            raise ValueError("OpenRouter API key required. Set OPENROUTER_API_KEY environment variable or pass api_key parameter.")

        self.model = model.removeprefix("openrouter/")

        # OpenRouter recommends these optional headers for request attribution.
        default_headers = {}
        if os.getenv("OPENROUTER_HTTP_REFERER"):
            default_headers["HTTP-Referer"] = os.getenv("OPENROUTER_HTTP_REFERER")
        if os.getenv("OPENROUTER_X_TITLE"):
            default_headers["X-Title"] = os.getenv("OPENROUTER_X_TITLE")

        client_kwargs = {
            "api_key": self.api_key,
            "base_url": "https://openrouter.ai/api/v1",
        }
        if default_headers:
            client_kwargs["default_headers"] = default_headers

        self.client = openai.OpenAI(**client_kwargs)

    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation using OpenRouter's OpenAI-compatible endpoint."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }

        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        response = self.client.chat.completions.create(**api_kwargs)
        message = response.choices[0].message

        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments,
                    id=tc.id
                ))

        usage = None
        if hasattr(response, 'usage') and response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            cost = calculate_cost(self.model, input_tokens, output_tokens)
            usage = TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost=cost,
            )

        return LLMResponse(content=message.content, tool_calls=tool_calls, raw_response=response, usage=usage)

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using JSON-mode + schema validation.

        Uses chat.completions with JSON response format for compatibility with
        OpenAI-like providers that may not support beta parse endpoints.
        """
        schema_json = json.dumps(output_schema.model_json_schema(), indent=2)
        schema_instruction = (
            "Return ONLY valid JSON (no markdown) that matches this JSON Schema:\n"
            f"{schema_json}"
        )

        structured_messages = [{"role": "system", "content": schema_instruction}, *messages]

        completion = self.client.chat.completions.create(
            model=self.model,
            messages=structured_messages,
            response_format={"type": "json_object"},
            **kwargs,
        )
        content = completion.choices[0].message.content or "{}"
        return output_schema.model_validate_json(content)


class MistralLLM(LLM):
    """Mistral AI LLM implementation using OpenAI-compatible endpoint."""

    def __init__(self, api_key: Optional[str] = None, model: str = "mistral/mistral-large-latest", **kwargs):
        self.api_key = api_key or os.getenv("MISTRAL_API_KEY")
        if not self.api_key:
            raise ValueError("Mistral API key required. Set MISTRAL_API_KEY environment variable or pass api_key parameter.")

        self.model = model.removeprefix("mistral/")
        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url="https://api.mistral.ai/v1"
        )

    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation using Mistral's OpenAI-compatible endpoint."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }

        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        response = self.client.chat.completions.create(**api_kwargs)
        message = response.choices[0].message

        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments,
                    id=tc.id
                ))

        usage = None
        if hasattr(response, 'usage') and response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            cost = calculate_cost(self.model, input_tokens, output_tokens)
            usage = TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost=cost,
            )

        return LLMResponse(content=message.content, tool_calls=tool_calls, raw_response=response, usage=usage)

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using JSON-mode + schema validation.

        Uses chat.completions with JSON response format for compatibility with
        OpenAI-like providers that may not support beta parse endpoints.
        """
        schema_json = json.dumps(output_schema.model_json_schema(), indent=2)
        schema_instruction = (
            "Return ONLY valid JSON (no markdown) that matches this JSON Schema:\n"
            f"{schema_json}"
        )

        structured_messages = [{"role": "system", "content": schema_instruction}, *messages]

        completion = self.client.chat.completions.create(
            model=self.model,
            messages=structured_messages,
            response_format={"type": "json_object"},
            **kwargs,
        )
        content = completion.choices[0].message.content or "{}"
        return output_schema.model_validate_json(content)


# Model registry mapping model names to providers
MODEL_REGISTRY = {
    # OpenAI models
    "gpt-4o": "openai",
    "gpt-4o-mini": "openai",
    "gpt-4-turbo": "openai",
    "gpt-3.5-turbo": "openai",
    "o1": "openai",
    "o1-mini": "openai",
    "o1-preview": "openai",
    "o4-mini": "openai",  # Testing placeholder
    
    # Anthropic Claude models
    "claude-3-5-sonnet": "anthropic",
    "claude-3-5-sonnet-20241022": "anthropic",
    "claude-3-5-sonnet-latest": "anthropic",
    "claude-3-5-haiku": "anthropic",
    "claude-3-5-haiku-20241022": "anthropic",
    "claude-3-5-haiku-latest": "anthropic",
    "claude-3-haiku-20240307": "anthropic",
    "claude-3-opus-20240229": "anthropic",
    "claude-3-opus-latest": "anthropic",
    "claude-3-sonnet-20240229": "anthropic",
    
    # Claude 4 models
    "claude-opus-4.1": "anthropic",
    "claude-opus-4-1-20250805": "anthropic",
    "claude-opus-4-1": "anthropic",  # Alias
    "claude-opus-4": "anthropic",
    "claude-opus-4-20250514": "anthropic",
    "claude-opus-4-0": "anthropic",  # Alias
    "claude-sonnet-4": "anthropic",
    "claude-sonnet-4-20250514": "anthropic",
    "claude-sonnet-4-0": "anthropic",  # Alias
    "claude-3-7-sonnet-latest": "anthropic",
    "claude-3-7-sonnet-20250219": "anthropic",
    
    # Google Gemini models
    "gemini-3-pro-preview": "google",
    "gemini-3-pro-image-preview": "google",
    "gemini-2.5-pro": "google",
    "gemini-2.5-flash": "google",
    "gemini-2.0-flash-exp": "google",
    "gemini-2.0-flash-thinking-exp": "google",
    "gemini-1.5-pro": "google",
    "gemini-1.5-pro-002": "google",
    "gemini-1.5-pro-001": "google",
    "gemini-1.5-flash": "google",
    "gemini-1.5-flash-002": "google",
    "gemini-1.5-flash-001": "google",
    "gemini-1.5-flash-8b": "google",
    "gemini-1.0-pro": "google",
}


class OpenOnionLLM(LLM):
    """OpenOnion managed keys LLM implementation using OpenAI-compatible API."""

    def __init__(self, api_key: Optional[str] = None, model: str = "co/o4-mini", **kwargs):
        # For co/ models, api_key is actually the auth token
        # Framework auto-loads .env, so OPENONION_API_KEY will be in environment
        self.auth_token = api_key or os.getenv("OPENONION_API_KEY")
        if not self.auth_token:
            raise ValueError(
                "OPENONION_API_KEY not found in environment.\n"
                "Run 'co init' to get started or set OPENONION_API_KEY in your .env file."
            )

        # Strip co/ prefix - it's only for client-side routing
        self.model = model.removeprefix("co/")

        # Determine base URL for OpenAI-compatible endpoint
        if os.getenv("OPENONION_DEV") or os.getenv("ENVIRONMENT") == "development":
            self.base_url = "http://localhost:8000/v1"
        else:
            self.base_url = "https://oo.openonion.ai/v1"

        # Use OpenAI client with OpenOnion endpoint
        self.client = openai.OpenAI(
            base_url=self.base_url,
            api_key=self.auth_token
        )

    def complete(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> LLMResponse:
        """Complete a conversation with optional tool support using OpenAI-compatible API."""
        api_kwargs = {
            "model": self.model,
            "messages": messages,
            **kwargs  # Pass through user kwargs (temperature, max_tokens, etc.)
        }

        # Add tools if provided
        if tools:
            api_kwargs["tools"] = [{"type": "function", "function": tool} for tool in tools]
            api_kwargs["tool_choice"] = "auto"

        try:
            response = self.client.chat.completions.create(**api_kwargs)
        except openai.APIStatusError as e:
            if e.status_code == 402:
                raise InsufficientCreditsError(e) from e
            elif e.status_code == 503:
                raise ProviderServiceError(e) from e
            logger.error(f"APIStatusError: status={e.status_code}, message={e.message}, body={getattr(e, 'body', None)}")
            raise
        except (openai.APITimeoutError, openai.APIConnectionError) as e:
            raise LLMConnectionError(e, model=f"co/{self.model}", base_url=self.base_url) from e
        except Exception as e:
            logger.error(f"LLM error: {type(e).__name__}: {e}")
            raise

        message = response.choices[0].message

        # Parse tool calls if present
        # Preserve extra_content for providers that need it (e.g., Gemini 3 thought_signature)
        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            for tc in message.tool_calls:
                extra = getattr(tc, 'extra_content', None)
                tool_calls.append(ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments,
                    id=tc.id,
                    extra_content=extra
                ))

        # Extract token usage (OpenAI-compatible format)
        usage = None
        if hasattr(response, 'usage') and response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            cached_tokens = 0
            if hasattr(response.usage, 'prompt_tokens_details') and response.usage.prompt_tokens_details:
                cached_tokens = getattr(response.usage.prompt_tokens_details, 'cached_tokens', 0) or 0
            # Use the underlying model for pricing (without co/ prefix)
            cost = calculate_cost(self.model, input_tokens, output_tokens, cached_tokens)
            usage = TokenUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cached_tokens=cached_tokens,
                cost=cost,
            )

        return LLMResponse(
            content=message.content,
            tool_calls=tool_calls,
            raw_response=response,
            usage=usage,
        )

    def structured_complete(self, messages: List[Dict], output_schema: Type[BaseModel], **kwargs) -> BaseModel:
        """Get structured Pydantic output using OpenAI-compatible chat completions API.

        Uses beta.chat.completions.parse() which routes through /v1/chat/completions,
        allowing proper provider routing for Gemini, OpenAI, and other models.
        """
        completion = self.client.beta.chat.completions.parse(
            model=self.model,
            messages=messages,
            response_format=output_schema,
            **kwargs
        )
        return completion.choices[0].message.parsed

    def get_balance(self) -> Optional[float]:
        """Fetch current account balance from OpenOnion API.

        Makes a GET request to /api/v1/auth/me endpoint to retrieve the user's
        current balance. This is called once at agent startup to display balance
        in the banner.

        Returns:
            Balance in USD (e.g., 4.22 for $4.22), or None if request fails

        Note:
            - Fast timeout (5s) to avoid hanging on network issues
            - Only called for co/ models (OpenOnion managed keys)
            - Returns None on any error (network, auth, etc.)
            - ~200ms typical latency, acceptable for startup
        """
        import requests

        # Build auth endpoint URL (strip /v1 suffix)
        auth_url = f"{self.base_url.rstrip('/v1')}/api/v1/auth/me"

        # 15s timeout: balance ~200ms typical, but proxy/VPN users need more headroom
        # Network errors return None — balance is non-critical banner info
        try:
            response = requests.get(
                auth_url,
                headers={"Authorization": f"Bearer {self.auth_token}"},
                timeout=15
            )
        except (requests.exceptions.Timeout, requests.exceptions.ConnectionError):
            return None

        if response.status_code == 200:
            data = response.json()
            return data.get("balance_usd")

        return None


def create_llm(model: str, api_key: Optional[str] = None, **kwargs) -> LLM:
    """Factory function to create the appropriate LLM based on model name.
    
    Args:
        model: The model name (e.g., "gpt-4o", "claude-3-5-sonnet", "gemini-1.5-pro")
        api_key: Optional API key to override environment variable
        **kwargs: Additional arguments to pass to the LLM constructor
    
    Returns:
        An LLM instance for the specified model
    
    Raises:
        ValueError: If the model is not recognized
    """
    # Check if it's a co/ model (OpenOnion managed keys)
    if model.startswith("co/"):
        return OpenOnionLLM(api_key=api_key, model=model, **kwargs)

    # Explicit provider prefixes for OpenAI-compatible third-party providers
    if model.startswith("groq/"):
        return GroqLLM(api_key=api_key, model=model, **kwargs)
    if model.startswith("openrouter/"):
        return OpenRouterLLM(api_key=api_key, model=model, **kwargs)
    if model.startswith("grok/"):
        return GrokLLM(api_key=api_key, model=model, **kwargs)
    if model.startswith("mistral/"):
        return MistralLLM(api_key=api_key, model=model, **kwargs)
    
    # Get provider from registry
    provider = MODEL_REGISTRY.get(model)
    
    if not provider:
        # Try to infer provider from model name
        if model.startswith("gpt") or model.startswith("o"):
            provider = "openai"
        elif model.startswith("claude"):
            provider = "anthropic"
        elif model.startswith("gemini"):
            provider = "google"
        else:
            raise ValueError(f"Unknown model '{model}'")
    
    # Create the appropriate LLM
    if provider == "openai":
        return OpenAILLM(api_key=api_key, model=model, **kwargs)
    elif provider == "anthropic":
        return AnthropicLLM(api_key=api_key, model=model, **kwargs)
    elif provider == "google":
        return GeminiLLM(api_key=api_key, model=model, **kwargs)
    else:
        raise ValueError(f"Provider '{provider}' not implemented")
