# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
LLM Provider Abstraction

Supports:
- Anthropic (Claude)
- OpenAI (GPT-4, GPT-4o)
- Ollama (Llama, Mistral, DeepSeek, etc. - local/free)

All providers use a unified interface for easy switching.
Supports both synchronous and streaming responses.

Phase 2.1 (v2.1): Structured exception handling
Phase 3 (v2.1): LLM call metrics collection
"""

import json
import logging
import os
import sys
import time
import urllib.error
import urllib.request
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import Generator, Optional

from .exceptions import (
    ProviderAuthenticationError,
    ProviderConnectionError,
    ProviderContentFilterError,
    ProviderError,
    ProviderRateLimitError,
)
from .metrics import get_metrics_collector
from .observability import Span, Trace, calculate_cost
from .resilience import RetryableOperation, RetryConfig
from .secrets_detection import mask_secrets

logger = logging.getLogger(__name__)


def _wrap_provider_error(error: Exception, provider_name: str) -> ProviderError:
    """
    Convert provider-specific exceptions to our structured exception hierarchy.

    This ensures consistent error handling regardless of which provider is used.
    """
    error_str = str(error)
    error_type = type(error).__name__

    # Mask any secrets that might be in error messages
    error_str = mask_secrets(error_str)

    # Check for common error patterns
    error_lower = error_str.lower()

    # Authentication errors
    if any(
        kw in error_lower for kw in ["authentication", "invalid api key", "unauthorized", "401"]
    ):
        return ProviderAuthenticationError(
            f"{provider_name} authentication failed: {error_str}",
            provider_name=provider_name,
        )

    # Rate limit errors
    if any(kw in error_lower for kw in ["rate limit", "too many requests", "429", "quota"]):
        return ProviderRateLimitError(
            f"{provider_name} rate limit exceeded: {error_str}",
            provider_name=provider_name,
            retry_after=None,  # Could parse from headers if available
        )

    # Connection errors
    if any(
        kw in error_lower
        for kw in ["connection", "timeout", "network", "unreachable", "503", "502"]
    ):
        return ProviderConnectionError(
            f"{provider_name} connection error: {error_str}",
            provider_name=provider_name,
        )

    # Content filter errors
    if any(kw in error_lower for kw in ["content filter", "content policy", "safety", "blocked"]):
        return ProviderContentFilterError(
            f"{provider_name} content filtered: {error_str}",
            provider_name=provider_name,
        )

    # Generic provider error — include exception type for debug context
    return ProviderError(
        f"{provider_name} error [{error_type}]: {error_str}",
        provider_name=provider_name,
    )


class StreamEventType(str, Enum):
    """Types of events emitted during streaming."""

    TEXT = "text"  # Text chunk
    TOOL_START = "tool_start"  # Tool execution starting
    TOOL_END = "tool_end"  # Tool execution complete
    ERROR = "error"  # Error occurred
    DONE = "done"  # Stream complete


@dataclass
class StreamEvent:
    """
    A single event in a streaming response.

    For text chunks: type=TEXT, content=the text
    For tool progress: type=TOOL_START/TOOL_END, tool_name and tool_result
    For completion: type=DONE, usage contains token counts
    """

    type: StreamEventType
    content: str = ""
    tool_name: Optional[str] = None
    tool_input: Optional[dict] = None
    tool_result: Optional[str] = None
    usage: Optional[dict] = None
    error: Optional[str] = None


@dataclass
class ToolCall:
    """Represents a tool/function call from the model."""

    id: str
    name: str
    input: dict


@dataclass
class LLMResponse:
    """Unified response format across all providers."""

    text: Optional[str]
    tool_calls: list[ToolCall]
    stop_reason: str  # "end_turn", "tool_use", "max_tokens"
    usage: Optional[dict] = None  # Token usage if available


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

    @abstractmethod
    def chat(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
        retry_config: Optional[RetryConfig] = None,
    ) -> LLMResponse:
        """Send messages and get response."""
        pass

    def chat_stream(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
    ) -> Generator[StreamEvent, None, LLMResponse]:
        """
        Stream a response, yielding text chunks as they arrive.

        Yields:
            StreamEvent objects with type=TEXT for content chunks

        Returns:
            The complete LLMResponse when done (accessible via generator.value after StopIteration)

        Usage:
            gen = provider.chat_stream(messages, tools, system)
            try:
                for event in gen:
                    if event.type == StreamEventType.TEXT:
                        print(event.content, end="", flush=True)
            except StopIteration as e:
                response = e.value  # The complete LLMResponse

        Note: Default implementation falls back to non-streaming chat().
        Override in subclasses for true streaming support.
        """
        # Default: fall back to non-streaming
        response = self.chat(messages, tools, system, max_tokens, trace)

        # Yield the complete text as a single chunk
        if response.text:
            yield StreamEvent(type=StreamEventType.TEXT, content=response.text)

        yield StreamEvent(type=StreamEventType.DONE, usage=response.usage)

        return response

    @property
    @abstractmethod
    def name(self) -> str:
        """Human-readable provider name."""
        pass

    @property
    @abstractmethod
    def model_name(self) -> str:
        """The model identifier (for cost calculation)."""
        pass

    @property
    @abstractmethod
    def supports_tools(self) -> bool:
        """Whether this provider supports tool/function calling."""
        pass

    @property
    def supports_streaming(self) -> bool:
        """Whether this provider supports true streaming. Override if supported."""
        return False


# Model context window sizes (tokens). Prefix-matched, most specific first.
MODEL_CONTEXT_WINDOWS: dict[str, int] = {
    # Anthropic Claude
    "claude-opus": 200000,
    "claude-sonnet": 200000,
    "claude-haiku": 200000,
    "claude": 200000,
    # OpenAI
    "gpt-4o": 128000,
    "gpt-4o-mini": 128000,
    "gpt-4-turbo": 128000,
    "gpt-4": 8192,
    "gpt-3.5": 16385,
    # Gemini
    "gemini-2.5": 1048576,
    "gemini-2.0": 1048576,
    "gemini-1.5": 1048576,
    "gemini": 1048576,
    # Ollama / local models
    "llama3.2": 131072,
    "llama3.1": 131072,
    "llama3.3": 131072,
    "qwen3.5": 131072,
    "qwen2.5": 32768,
    "qwen": 32768,
    "deepseek-r1": 65536,
    "deepseek": 65536,
    "mistral": 32768,
    "gemma3": 8192,
    "gemma2": 8192,
    "phi3": 4096,
    "phi4": 16384,
    "smollm2": 2048,
    "tinyllama": 2048,
    "command-r": 131072,
}


def get_context_window(model_name: str) -> int:
    """Return the context window size for a model name (prefix-matched).

    Falls back to 8192 for unknown models.
    """
    name = model_name.lower()
    for prefix, window in MODEL_CONTEXT_WINDOWS.items():
        if name.startswith(prefix):
            return window
    return 8192


class AnthropicProvider(LLMProvider):
    """Claude via Anthropic API."""

    def __init__(self, model: str = "claude-sonnet-4-6"):
        from anthropic import Anthropic

        # SDK retries disabled — RetryableOperation handles all retries uniformly
        self.client = Anthropic(
            timeout=90.0,
            max_retries=0,
        )
        self.model = model

    @property
    def name(self) -> str:
        return f"Claude ({self.model.split('-')[1] if '-' in self.model else self.model})"

    @property
    def model_name(self) -> str:
        return self.model

    @property
    def supports_tools(self) -> bool:
        return True

    def chat(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
        retry_config: Optional[RetryConfig] = None,
    ) -> LLMResponse:

        start_time = time.time()
        prompt_tokens = 0
        completion_tokens = 0
        cost_usd = 0.0
        success = True
        error_message = None

        # Convert any stored LLMResponse objects back to proper format
        converted_messages = []
        for msg in messages:
            if msg["role"] == "assistant" and isinstance(msg.get("content"), LLMResponse):
                resp = msg["content"]
                content = []
                if resp.text:
                    content.append({"type": "text", "text": resp.text})
                for tc in resp.tool_calls:
                    content.append(
                        {"type": "tool_use", "id": tc.id, "name": tc.name, "input": tc.input}
                    )
                converted_messages.append({"role": "assistant", "content": content})
            else:
                converted_messages.append(msg)

        kwargs = {
            "model": self.model,
            "max_tokens": max_tokens,
            "system": system,
            "messages": converted_messages,
        }

        if tools:
            kwargs["tools"] = tools

        retry_config = retry_config or RetryConfig(timeout=120.0)

        try:
            with RetryableOperation(trace, "anthropic_chat", retry_config) as op:
                for attempt in op.attempts():
                    with attempt:
                        response = self.client.messages.create(**kwargs)

                        text = None
                        tool_calls = []

                        for block in response.content:
                            if hasattr(block, "text"):
                                text = (text or "") + block.text
                            elif block.type == "tool_use":
                                tool_calls.append(
                                    ToolCall(id=block.id, name=block.name, input=block.input)
                                )

                        # Preserve the full stop_reason from the API (end_turn, tool_use, max_tokens, stop_sequence, etc.)
                        stop_reason = response.stop_reason or "end_turn"

                        usage = None
                        if hasattr(response, "usage"):
                            usage = {
                                "input_tokens": response.usage.input_tokens,
                                "output_tokens": response.usage.output_tokens,
                            }
                            prompt_tokens = response.usage.input_tokens
                            completion_tokens = response.usage.output_tokens
                            cost_usd = calculate_cost(self.model, prompt_tokens, completion_tokens)
                            if attempt.span:
                                attempt.span.set_tokens(usage, self.model)

                        return LLMResponse(
                            text=text, tool_calls=tool_calls, stop_reason=stop_reason, usage=usage
                        )

        except ProviderError:
            success = False
            error_message = str(sys.exc_info()[1])
            raise
        except Exception as e:
            success = False
            error_message = str(e)
            # Raw provider exception escaped retries — wrap into our hierarchy
            raise _wrap_provider_error(e, "Anthropic") from e

        finally:
            duration_ms = (time.time() - start_time) * 1000
            try:
                metrics = get_metrics_collector()
                metrics.record_llm_call(
                    provider="anthropic",
                    model=self.model,
                    duration_ms=duration_ms,
                    success=success,
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    cost_usd=cost_usd,
                    error_message=error_message,
                    user_id=trace.user_id if trace else None,
                    trace_id=trace.trace_id if trace else None,
                )
            except Exception as metrics_error:
                logger.debug(f"Failed to record LLM metrics: {metrics_error}")

    @property
    def supports_streaming(self) -> bool:
        return True

    def chat_stream(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
    ) -> Generator[StreamEvent, None, LLMResponse]:
        """
        Stream a response from Claude.

        Note: When tools are present and the model wants to use them,
        we fall back to non-streaming since tool calls need to complete
        before we can continue.
        """
        # Convert messages (same as chat())
        converted_messages = []
        for msg in messages:
            if msg["role"] == "assistant" and isinstance(msg.get("content"), LLMResponse):
                resp = msg["content"]
                content = []
                if resp.text:
                    content.append({"type": "text", "text": resp.text})
                for tc in resp.tool_calls:
                    content.append(
                        {"type": "tool_use", "id": tc.id, "name": tc.name, "input": tc.input}
                    )
                converted_messages.append({"role": "assistant", "content": content})
            else:
                converted_messages.append(msg)

        kwargs = {
            "model": self.model,
            "max_tokens": max_tokens,
            "system": system,
            "messages": converted_messages,
        }

        if tools:
            kwargs["tools"] = tools

        # Start span for observability
        span = None
        if trace:
            span = Span.create_compat(name="anthropic_chat_stream", trace_id=trace.trace_id)

        span_finished = False

        def _finish_span(error=None):
            nonlocal span_finished
            if span and not span_finished:
                span_finished = True
                if error:
                    span.set_error(error)
                span.finish()
                if trace:
                    trace.add_span(span)

        try:
            # Use streaming API
            with self.client.messages.stream(**kwargs) as stream:
                collected_text = ""
                tool_calls = []
                usage = None

                for event in stream:
                    if hasattr(event, "type"):
                        if event.type == "content_block_delta":
                            if hasattr(event.delta, "text"):
                                chunk = event.delta.text
                                collected_text += chunk
                                yield StreamEvent(type=StreamEventType.TEXT, content=chunk)

                        elif event.type == "content_block_start":
                            if (
                                hasattr(event.content_block, "type")
                                and event.content_block.type == "tool_use"
                            ):
                                pass

                # Get the final message
                final_message = stream.get_final_message()

                for block in final_message.content:
                    if block.type == "tool_use":
                        tool_calls.append(ToolCall(id=block.id, name=block.name, input=block.input))

                if hasattr(final_message, "usage"):
                    usage = {
                        "input_tokens": final_message.usage.input_tokens,
                        "output_tokens": final_message.usage.output_tokens,
                    }
                    if span:
                        span.set_tokens(usage, self.model)

                # Preserve the full stop_reason from the API
                stop_reason = final_message.stop_reason or "end_turn"

                yield StreamEvent(type=StreamEventType.DONE, usage=usage)

                response = LLMResponse(
                    text=collected_text or None,
                    tool_calls=tool_calls,
                    stop_reason=stop_reason,
                    usage=usage,
                )

                _finish_span()
                return response

        except GeneratorExit:
            # Caller abandoned the generator — finish span cleanly and re-raise
            _finish_span()
            raise
        except Exception as e:
            wrapped_error = _wrap_provider_error(e, "Anthropic")
            _finish_span(error=wrapped_error)
            yield StreamEvent(type=StreamEventType.ERROR, error=str(wrapped_error))
            raise wrapped_error


class OpenAIProvider(LLMProvider):
    """GPT via OpenAI API."""

    def __init__(self, model: str = "gpt-4o", base_url: str = None, api_key: str = None):
        from openai import OpenAI

        # SDK retries disabled — RetryableOperation handles all retries uniformly
        self.client = OpenAI(
            timeout=90.0,
            max_retries=0,
            **({"base_url": base_url} if base_url else {}),
            **({"api_key": api_key} if api_key else {}),
        )
        self.model = model

    @property
    def name(self) -> str:
        return f"GPT ({self.model})"

    @property
    def model_name(self) -> str:
        return self.model

    @property
    def supports_tools(self) -> bool:
        return True

    def _convert_tools(self, tools: list) -> list:
        """Convert Anthropic tool format to OpenAI."""
        return [
            {
                "type": "function",
                "function": {
                    "name": t["name"],
                    "description": t["description"],
                    "parameters": t["input_schema"],
                },
            }
            for t in tools
        ]

    def _convert_messages(self, messages: list, system: str) -> list:
        """Convert message format for OpenAI."""
        converted = [{"role": "system", "content": system}]

        for msg in messages:
            role = msg["role"]
            content = msg["content"]

            if role == "user":
                if isinstance(content, list):
                    # Check for tool results — scan all items, not just content[0],
                    # since a mixed list could have text before the result block.
                    tool_results = [
                        item
                        for item in content
                        if isinstance(item, dict) and item.get("type") == "tool_result"
                    ]
                    if tool_results:
                        for result in tool_results:
                            converted.append(
                                {
                                    "role": "tool",
                                    "tool_call_id": result["tool_use_id"],
                                    "content": str(result["content"]),
                                }
                            )
                        continue
                try:
                    _content = content if isinstance(content, str) else json.dumps(content)
                except (TypeError, ValueError):
                    _content = str(content)
                converted.append({"role": "user", "content": _content})

            elif role == "assistant":
                if isinstance(content, LLMResponse):
                    msg_dict = {"role": "assistant"}
                    if content.text:
                        msg_dict["content"] = content.text
                    if content.tool_calls:
                        msg_dict["tool_calls"] = [
                            {
                                "id": tc.id,
                                "type": "function",
                                "function": {
                                    "name": tc.name,
                                    "arguments": json.dumps(tc.input, default=str),
                                },
                            }
                            for tc in content.tool_calls
                        ]
                    converted.append(msg_dict)
                else:
                    converted.append(
                        {
                            "role": "assistant",
                            "content": content if isinstance(content, str) else str(content),
                        }
                    )

        return converted

    def chat(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
        retry_config: Optional[RetryConfig] = None,
    ) -> LLMResponse:

        openai_messages = self._convert_messages(messages, system)

        # Sanitize: ensure all messages are JSON-serializable (tool results
        # can contain non-serializable objects like Memory instances).
        openai_messages = json.loads(json.dumps(openai_messages, default=str))

        kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "max_tokens": max_tokens,
        }

        if tools:
            kwargs["tools"] = self._convert_tools(tools)

        # Use retry logic if config provided
        retry_config = retry_config or RetryConfig(timeout=120.0)

        try:
            with RetryableOperation(trace, "openai_chat", retry_config) as op:
                for attempt in op.attempts():
                    with attempt:
                        response = self.client.chat.completions.create(**kwargs)

                        choice = response.choices[0]
                        message = choice.message

                        text = message.content
                        tool_calls = []

                        if message.tool_calls:
                            for tc in message.tool_calls:
                                try:
                                    args = json.loads(tc.function.arguments)
                                except (json.JSONDecodeError, TypeError):
                                    logger.warning(
                                        f"Malformed tool call arguments for {tc.function.name}: {tc.function.arguments!r}"
                                    )
                                    args = {}
                                tool_calls.append(
                                    ToolCall(id=tc.id, name=tc.function.name, input=args)
                                )

                        usage = None
                        if response.usage:
                            usage = {
                                "input_tokens": response.usage.prompt_tokens,
                                "output_tokens": response.usage.completion_tokens,
                            }
                            if attempt.span:
                                attempt.span.set_tokens(usage, self.model)

                        return LLMResponse(
                            text=text,
                            tool_calls=tool_calls,
                            stop_reason="tool_use" if tool_calls else "end_turn",
                            usage=usage,
                        )

        except ProviderError:
            raise
        except Exception as e:
            raise _wrap_provider_error(e, "OpenAI") from e

    @property
    def supports_streaming(self) -> bool:
        return True

    def chat_stream(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
    ) -> Generator[StreamEvent, None, LLMResponse]:
        """Stream a response from OpenAI."""
        openai_messages = self._convert_messages(messages, system)
        openai_messages = json.loads(json.dumps(openai_messages, default=str))

        kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "max_tokens": max_tokens,
            "stream": True,
            # stream_options: include_usage is required to receive token counts
            # in streamed responses. Without this, usage is always None.
            "stream_options": {"include_usage": True},
        }

        if tools:
            kwargs["tools"] = self._convert_tools(tools)

        # Start span for observability
        span = None
        if trace:
            span = Span.create_compat(name="openai_chat_stream", trace_id=trace.trace_id)

        span_finished = False

        def _finish_span(error=None):
            nonlocal span_finished
            if span and not span_finished:
                span_finished = True
                span.model = self.model
                if error:
                    span.set_error(error)
                span.finish()
                if trace:
                    trace.add_span(span)

        try:
            stream = self.client.chat.completions.create(**kwargs)

            collected_text = ""
            tool_calls_data = {}
            usage = None

            for chunk in stream:
                if not chunk.choices:
                    # Final usage-only chunk
                    if hasattr(chunk, "usage") and chunk.usage:
                        usage = {
                            "input_tokens": chunk.usage.prompt_tokens,
                            "output_tokens": chunk.usage.completion_tokens,
                        }
                        if span:
                            span.set_tokens(usage, self.model)
                    continue

                delta = chunk.choices[0].delta

                if delta.content:
                    collected_text += delta.content
                    yield StreamEvent(type=StreamEventType.TEXT, content=delta.content)

                if delta.tool_calls:
                    for tc in delta.tool_calls:
                        idx = tc.index
                        if idx not in tool_calls_data:
                            tool_calls_data[idx] = {"id": "", "name": "", "arguments": ""}
                        if tc.id:
                            tool_calls_data[idx]["id"] = tc.id
                        if tc.function:
                            if tc.function.name:
                                tool_calls_data[idx]["name"] = tc.function.name
                            if tc.function.arguments:
                                tool_calls_data[idx]["arguments"] += tc.function.arguments

                if hasattr(chunk, "usage") and chunk.usage:
                    usage = {
                        "input_tokens": chunk.usage.prompt_tokens,
                        "output_tokens": chunk.usage.completion_tokens,
                    }
                    if span:
                        span.set_tokens(usage, self.model)

            tool_calls = []
            for idx in sorted(tool_calls_data.keys()):
                tc_data = tool_calls_data[idx]
                if tc_data["id"] and tc_data["name"]:
                    try:
                        args = json.loads(tc_data["arguments"]) if tc_data["arguments"] else {}
                    except json.JSONDecodeError:
                        args = {}
                    tool_calls.append(ToolCall(id=tc_data["id"], name=tc_data["name"], input=args))

            yield StreamEvent(type=StreamEventType.DONE, usage=usage)
            _finish_span()

            return LLMResponse(
                text=collected_text or None,
                tool_calls=tool_calls,
                stop_reason="tool_use" if tool_calls else "end_turn",
                usage=usage,
            )

        except GeneratorExit:
            _finish_span()
            raise
        except Exception as e:
            wrapped_error = _wrap_provider_error(e, "OpenAI")
            _finish_span(error=wrapped_error)
            yield StreamEvent(type=StreamEventType.ERROR, error=str(wrapped_error))
            raise wrapped_error


class OllamaProvider(LLMProvider):
    """Local models via Ollama (Llama, Mistral, etc.).

    Uses Ollama's native HTTP API (urllib) with zero extra dependencies.
    Falls back to the OpenAI-compatible client if the openai package
    is installed, which enables tool-calling on supported models.
    """

    def __init__(
        self, model: str = "llama3.2", base_url: str = "http://localhost:11434", timeout: int = 120
    ):
        self.model = model
        self.base_url = base_url.rstrip("/")
        # Allow env override for CPU-only hosts where large models need more time
        self.timeout = int(os.environ.get("OLLAMA_TIMEOUT", timeout))

        # Try to use the OpenAI client for richer tool support
        self.client = None
        try:
            from openai import OpenAI

            # SDK retries disabled — RetryableOperation handles all retries uniformly
            self.client = OpenAI(
                base_url=f"{self.base_url}/v1",
                api_key="ollama",
                timeout=self.timeout,  # match the configured Ollama timeout (default 120s)
                max_retries=0,
            )
        except ImportError:
            pass  # Will use native urllib fallback

    @property
    def name(self) -> str:
        return f"Ollama ({self.model})"

    @property
    def model_name(self) -> str:
        return self.model

    @property
    def supports_tools(self) -> bool:
        """Whether this Ollama model supports tool/function calling.

        Checks a static known-good list first (fast, no network), then probes
        the Ollama /api/show endpoint for unknown models to read the model's
        declared capabilities.  Falls back to False if the probe fails.

        Known-good families (updated as Ollama expands tool support):
          Versioned: llama3.1, llama3.2, llama3.3
          Families:  mistral, qwen, deepseek, phi4, gemma3, command-r
        """
        base = self.model.lower().split(":")[0]  # strip tag (e.g. ':14b', ':latest')

        # Versioned prefixes: next char must be a non-digit separator or end-of-string
        versioned_prefixes = ["llama3.1", "llama3.2", "llama3.3"]
        for prefix in versioned_prefixes:
            if base == prefix:
                return True
            if base.startswith(prefix):
                next_char = base[len(prefix)] if len(base) > len(prefix) else ""
                if not next_char or next_char in ".-":  # e.g. llama3.2.1 or llama3.2-instruct
                    return True
                # Digit continuation means a different version (llama3.10 ≠ llama3.1)

        # Family prefixes — any suffix is acceptable.
        # gemma3 is intentionally excluded: as of Ollama 0.5, gemma3 does NOT
        # reliably support tool calling. phi4 and command-r are added proactively.
        family_prefixes = [
            "mistral",  # mistral, mistral-nemo, mistral-large, etc.
            "qwen",  # qwen, qwen2, qwen2.5, qwen3, etc.
            "deepseek",  # deepseek-r1, deepseek-v3, etc.
            "phi4",  # phi4, phi4-mini
            "command-r",  # Cohere Command-R family
        ]
        for prefix in family_prefixes:
            if base == prefix:
                return True
            # Prefix followed by digit or separator — e.g. "qwen2" starts with "qwen"
            if (
                len(base) > len(prefix)
                and base[: len(prefix)] == prefix
                and base[len(prefix)] in "0123456789.-"
            ):
                return True
            # Prefix followed by hyphen-word — e.g. "mistral-nemo"
            if base.startswith(prefix + "-"):
                return True

        # Unknown model: probe /api/show for capability declaration (Ollama ≥0.3)
        return self._probe_tool_support()

    def _probe_tool_support(self) -> bool:
        """Query /api/show to determine if the model declares tool support.

        Ollama ≥0.3 populates modelinfo.capabilities with a list that includes
        "tools" when the model supports function calling.  This is authoritative
        for custom or community models not in the static list above.
        """
        import json as _json
        import urllib.error
        import urllib.request

        try:
            payload = _json.dumps({"model": self.model}).encode()
            req = urllib.request.Request(
                f"{self.base_url}/api/show",
                data=payload,
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=3) as resp:
                data = _json.loads(resp.read())
            caps = data.get("capabilities", [])
            if isinstance(caps, list):
                return "tools" in caps
            # Older Ollama versions: check template for tool-call markers
            template = data.get("template", "")
            return "{{- if .Tools }}" in template or "tool_call" in template
        except Exception:
            return False  # Can't reach Ollama or model not pulled; be conservative

    def chat(
        self,
        messages: list,
        tools: list,
        system: str,
        max_tokens: int = 4096,
        trace: Optional[Trace] = None,
        retry_config: Optional[RetryConfig] = None,
    ) -> LLMResponse:

        # Build messages
        ollama_messages = [{"role": "system", "content": system}]

        for msg in messages:
            role = msg["role"]
            content = msg["content"]

            if role == "user":
                if isinstance(content, list):
                    # Tool results - format as text
                    parts = []
                    for item in content:
                        if isinstance(item, dict) and item.get("type") == "tool_result":
                            parts.append(f"[Tool Result]:\n{item['content']}")
                    content = "\n".join(parts) if parts else str(content)
                ollama_messages.append({"role": "user", "content": content})

            elif role == "assistant":
                if isinstance(content, LLMResponse):
                    text = content.text or ""
                    if content.tool_calls:
                        # Include tool call info in text for context
                        for tc in content.tool_calls:
                            text += f"\n[Called tool: {tc.name}]"
                    ollama_messages.append({"role": "assistant", "content": text})
                else:
                    ollama_messages.append({"role": "assistant", "content": str(content)})

        # Use retry logic
        retry_config = retry_config or RetryConfig(timeout=float(self.timeout))

        # ── Native urllib path (no openai package) ──
        if self.client is None:
            return self._chat_native(
                ollama_messages, max_tokens, retry_config, trace, timeout=self.timeout
            )

        # ── OpenAI client path (tool calling support) ──
        # Try with tools if supported
        if self.supports_tools and tools:
            try:
                with RetryableOperation(trace, "ollama_chat_tools", retry_config) as op:
                    for attempt in op.attempts():
                        with attempt:
                            openai_tools = [
                                {
                                    "type": "function",
                                    "function": {
                                        "name": t["name"],
                                        "description": t["description"],
                                        "parameters": t["input_schema"],
                                    },
                                }
                                for t in tools
                            ]

                            # Cap context window: small models on CPU benefit from
                            # a tighter window (2048) to reduce inference time.
                            _ctx_size = int(os.environ.get("OLLAMA_NUM_CTX", 0)) or max(
                                2048, max_tokens * 2
                            )
                            response = self.client.chat.completions.create(
                                model=self.model,
                                messages=ollama_messages,
                                tools=openai_tools,
                                max_tokens=max_tokens,
                                extra_body={"options": {"num_ctx": _ctx_size}},
                            )

                            choice = response.choices[0]
                            message = choice.message

                            tool_calls = []
                            if hasattr(message, "tool_calls") and message.tool_calls:
                                for tc in message.tool_calls:
                                    try:
                                        args = json.loads(tc.function.arguments)
                                    except (json.JSONDecodeError, TypeError):
                                        logger.warning(
                                            f"Malformed tool call arguments for {tc.function.name}: {tc.function.arguments!r}"
                                        )
                                        args = {}
                                    tool_calls.append(
                                        ToolCall(id=tc.id, name=tc.function.name, input=args)
                                    )

                            if attempt.span:
                                attempt.span.model = self.model

                            return LLMResponse(
                                text=message.content,
                                tool_calls=tool_calls,
                                stop_reason="tool_use" if tool_calls else "end_turn",
                            )

            except Exception as e:
                # Tool calling failed on this model — fall back to plain chat.
                # The caller will receive no tool_calls and can handle accordingly.
                logger.warning(
                    f"Ollama tool call failed for {self.model}: {e}. "
                    f"Falling back to plain chat — model will not see tool results."
                )

        # Plain chat (no tools)
        try:
            with RetryableOperation(trace, "ollama_chat", retry_config) as op:
                for attempt in op.attempts():
                    with attempt:
                        _ctx_size = int(os.environ.get("OLLAMA_NUM_CTX", 0)) or max(
                            2048, max_tokens * 2
                        )
                        response = self.client.chat.completions.create(
                            model=self.model,
                            messages=ollama_messages,
                            max_tokens=max_tokens,
                            extra_body={"options": {"num_ctx": _ctx_size}},
                        )

                        if attempt.span:
                            attempt.span.model = self.model

                        return LLMResponse(
                            text=response.choices[0].message.content,
                            tool_calls=[],
                            stop_reason="end_turn",
                        )

        except ProviderError:
            raise
        except Exception as e:
            raise _wrap_provider_error(e, "Ollama") from e

    def _chat_native(
        self,
        messages: list,
        max_tokens: int,
        retry_config: RetryConfig,
        trace: Optional[Trace] = None,
        timeout: int = 120,
    ) -> LLMResponse:
        """Ollama chat via native HTTP API — no openai package needed."""
        import urllib.error
        import urllib.request

        payload = json.dumps(
            {
                "model": self.model,
                "messages": messages,
                "stream": False,
                "options": {"num_predict": max_tokens},
            }
        ).encode()

        try:
            with RetryableOperation(trace, "ollama_chat_native", retry_config) as op:
                for attempt in op.attempts():
                    with attempt:
                        req = urllib.request.Request(
                            f"{self.base_url}/api/chat",
                            data=payload,
                            headers={"Content-Type": "application/json"},
                        )
                        try:
                            with urllib.request.urlopen(req, timeout=timeout) as resp:
                                data = json.loads(resp.read())
                        except urllib.error.URLError as e:
                            raise ConnectionError(
                                f"Cannot reach Ollama at {self.base_url}. "
                                f"Is it running?  Start with: ollama serve"
                            ) from e

                        text = data.get("message", {}).get("content", "")

                        if attempt.span:
                            attempt.span.model = self.model

                        return LLMResponse(
                            text=text,
                            tool_calls=[],
                            stop_reason="end_turn",
                        )

        except ProviderError:
            raise
        except Exception as e:
            raise _wrap_provider_error(e, "Ollama") from e


# Provider registry
class GeminiProvider(OpenAIProvider):
    """Google Gemini via OpenAI-compatible API (free tier available)."""

    def __init__(self, model: str = "gemini-2.5-flash"):
        api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY", "")
        if not api_key:
            raise ValueError(
                "Gemini requires GEMINI_API_KEY. Get one free at https://aistudio.google.com/apikey"
            )
        super().__init__(
            model=model,
            base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
            api_key=api_key,
        )

    @property
    def name(self) -> str:
        return f"Gemini ({self.model})"


PROVIDERS = {
    "anthropic": lambda cfg: AnthropicProvider(cfg.llm.anthropic_model),
    "claude": lambda cfg: AnthropicProvider(cfg.llm.anthropic_model),
    "openai": lambda cfg: OpenAIProvider(cfg.llm.openai_model),
    "gpt": lambda cfg: OpenAIProvider(cfg.llm.openai_model),
    "gpt-4o": lambda cfg: OpenAIProvider("gpt-4o"),
    "gpt-4o-mini": lambda cfg: OpenAIProvider("gpt-4o-mini"),
    "gemini": lambda cfg: GeminiProvider(os.environ.get("GEMINI_MODEL", "gemini-2.5-flash")),
    "google": lambda cfg: GeminiProvider(os.environ.get("GEMINI_MODEL", "gemini-2.5-flash")),
    "ollama": lambda cfg: OllamaProvider(cfg.llm.ollama_model, cfg.llm.ollama_base_url),
    "llama": lambda cfg: OllamaProvider("llama3.2", cfg.llm.ollama_base_url),
    "llama3.2": lambda cfg: OllamaProvider("llama3.2", cfg.llm.ollama_base_url),
    "mistral": lambda cfg: OllamaProvider("mistral", cfg.llm.ollama_base_url),
    "deepseek": lambda cfg: OllamaProvider("deepseek-r1:14b", cfg.llm.ollama_base_url),
    "qwen": lambda cfg: OllamaProvider("qwen3.5:27b", cfg.llm.ollama_base_url),
    "qwen2.5": lambda cfg: OllamaProvider("qwen2.5:14b", cfg.llm.ollama_base_url),
    "qwen3.5": lambda cfg: OllamaProvider("qwen3.5:27b", cfg.llm.ollama_base_url),
    "qwen3.5:27b": lambda cfg: OllamaProvider("qwen3.5:27b", cfg.llm.ollama_base_url),
    "qwen3.5:35b": lambda cfg: OllamaProvider("qwen3.5:35b", cfg.llm.ollama_base_url),
    "qwen3.5:122b": lambda cfg: OllamaProvider("qwen3.5:122b", cfg.llm.ollama_base_url),
}

# ── Lightweight models for background tasks (planning, extraction) ──
# These are small enough to run on a Raspberry Pi or similar.
LIGHTWEIGHT_MODELS = {
    "qwen2.5:0.5b": lambda cfg: OllamaProvider("qwen2.5:0.5b", cfg.llm.ollama_base_url),
    "qwen2.5:1.5b": lambda cfg: OllamaProvider("qwen2.5:1.5b", cfg.llm.ollama_base_url),
    "qwen2.5:3b": lambda cfg: OllamaProvider("qwen2.5:3b", cfg.llm.ollama_base_url),
    "gemma2:2b": lambda cfg: OllamaProvider("gemma2:2b", cfg.llm.ollama_base_url),
    "gemma3:1b": lambda cfg: OllamaProvider("gemma3:1b", cfg.llm.ollama_base_url),
    "smollm2:135m": lambda cfg: OllamaProvider("smollm2:135m", cfg.llm.ollama_base_url),
    "smollm2:360m": lambda cfg: OllamaProvider("smollm2:360m", cfg.llm.ollama_base_url),
    "smollm2:1.7b": lambda cfg: OllamaProvider("smollm2:1.7b", cfg.llm.ollama_base_url),
    "qwen3.5:27b": lambda cfg: OllamaProvider("qwen3.5:27b", cfg.llm.ollama_base_url),
    "qwen3.5:35b": lambda cfg: OllamaProvider("qwen3.5:35b", cfg.llm.ollama_base_url),
    "phi3:mini": lambda cfg: OllamaProvider("phi3:mini", cfg.llm.ollama_base_url),
    "tinyllama": lambda cfg: OllamaProvider("tinyllama", cfg.llm.ollama_base_url),
    # Cloud lightweight options — gpt-4o-mini is also in PROVIDERS, both are intentional:
    # PROVIDERS is for main agent use, LIGHTWEIGHT_MODELS is for background task resolution.
    "gpt-4o-mini": lambda cfg: OpenAIProvider("gpt-4o-mini"),
    "claude-haiku": lambda cfg: AnthropicProvider("claude-haiku-4-5-20251001"),
}

# Aliases for provider names (for tests and normalization)
PROVIDER_ALIASES = {
    "claude": "anthropic",
    "gpt": "openai",
    "google": "gemini",
    "llama": "ollama",
    "llama3.2": "ollama",
    "mistral": "ollama",
    "deepseek": "ollama",
    "qwen": "ollama",
    "qwen2.5": "ollama",
    "qwen3.5": "ollama",
}


def get_provider(name: str, config) -> LLMProvider:
    """Get a provider by name, resolving aliases automatically.

    Accepts both canonical names (e.g. 'anthropic') and common aliases
    (e.g. 'claude', 'gpt', 'llama') as defined in PROVIDER_ALIASES.
    """
    name = name.lower()
    # Resolve alias → canonical name
    name = PROVIDER_ALIASES.get(name, name)
    if name not in PROVIDERS:
        available = ", ".join(sorted(set(list(PROVIDERS.keys()) + list(PROVIDER_ALIASES.keys()))))
        raise ValueError(f"Unknown provider: {name}. Available: {available}")
    return PROVIDERS[name](config)


def get_available_providers() -> list[str]:
    """Get list of available provider names."""
    return sorted(PROVIDERS.keys())


def get_lightweight_provider(config) -> Optional[LLMProvider]:
    """
    Get a lightweight LLM provider for background tasks.

    Used by the planner and memory extractor so they don't burn
    the main (expensive) provider on maintenance work.

    Resolution order:
        1. Explicit config: lightweight_model + lightweight_provider
        2. Auto-detect: check Ollama for small models already pulled
        3. None: fall back to main provider (caller handles this)

    Returns:
        LLMProvider or None if no lightweight model is available
    """
    model = getattr(config.llm, "lightweight_model", "")
    provider_name = getattr(config.llm, "lightweight_provider", "")

    # 1. Explicit config
    if model:
        # Check LIGHTWEIGHT_MODELS registry first
        if model in LIGHTWEIGHT_MODELS:
            try:
                provider = LIGHTWEIGHT_MODELS[model](config)
                logger.info(f"Lightweight provider: {provider.name}")
                return provider
            except Exception as e:
                logger.warning(f"Failed to create lightweight provider {model}: {e}")

        # Try as a generic Ollama model
        if provider_name in ("ollama", ""):
            try:
                provider = OllamaProvider(model, config.llm.ollama_base_url)
                logger.info(f"Lightweight provider (Ollama): {provider.name}")
                return provider
            except Exception as e:
                logger.warning(f"Failed to create Ollama lightweight provider {model}: {e}")

        # Try as OpenAI model
        if provider_name == "openai":
            try:
                provider = OpenAIProvider(model)
                logger.info(f"Lightweight provider (OpenAI): {provider.name}")
                return provider
            except Exception as e:
                logger.warning(f"Failed to create OpenAI lightweight provider {model}: {e}")

        # Try as Anthropic model
        if provider_name == "anthropic":
            try:
                provider = AnthropicProvider(model)
                logger.info(f"Lightweight provider (Anthropic): {provider.name}")
                return provider
            except Exception as e:
                logger.warning(f"Failed to create Anthropic lightweight provider {model}: {e}")

    # 2. Auto-detect: check if Ollama has any small models pulled
    try:
        base_url = config.llm.ollama_base_url.rstrip("/")
        req = urllib.request.Request(f"{base_url}/api/tags", method="GET")
        with urllib.request.urlopen(req, timeout=2) as resp:
            data = json.loads(resp.read().decode())
            pulled_models = {m["name"] for m in data.get("models", [])}

            # Prefer smaller models first
            preferred_order = [
                "smollm2:135m",
                "smollm2:360m",
                "tinyllama",
                "qwen2.5:0.5b",
                "qwen2.5:1.5b",
                "smollm2:1.7b",
                "gemma3:1b",
                "phi3:mini",
                "qwen2.5:3b",
                "qwen3.5:35b",
                "qwen3.5:27b",
            ]
            for candidate in preferred_order:
                if candidate in pulled_models:
                    provider = OllamaProvider(candidate, base_url)
                    logger.info(f"Lightweight provider (auto-detected): {provider.name}")
                    return provider
    except Exception:
        pass  # Ollama not running or no small models

    # 3. Nothing available
    logger.debug("No lightweight provider available, will use main provider")
    return None


def _extract_param_size(model_name: str, model_info: Optional[dict] = None) -> float:
    """
    Extract parameter count (in billions) from Ollama model info or tag.

    Tries model_info["details"]["parameter_size"] first (e.g. "7.6B"),
    then falls back to parsing the tag (e.g. qwen2.5:7b -> 7.0).
    Returns 0.0 if size is unknown.
    """
    # Try Ollama metadata first
    if model_info:
        try:
            size_str = model_info.get("details", {}).get("parameter_size", "")
            if size_str:
                # e.g. "7.6B", "1.5B", "360M"
                size_str = size_str.strip().upper()
                if size_str.endswith("B"):
                    return float(size_str[:-1])
                if size_str.endswith("M"):
                    return float(size_str[:-1]) / 1000.0
        except (ValueError, TypeError, AttributeError):
            pass

    # Fall back to parsing the tag (e.g. "qwen2.5:7b" -> 7.0)
    if ":" in model_name:
        tag = model_name.split(":")[-1].lower()
        # Strip common suffixes like "-instruct", "-chat", etc.
        for suffix in ("-instruct", "-chat", "-code", "-fp16", "-q4_0", "-q8_0"):
            tag = tag.replace(suffix, "")
        tag = tag.strip()
        if tag.endswith("b"):
            try:
                return float(tag[:-1])
            except ValueError:
                pass
        if tag.endswith("m"):
            try:
                return float(tag[:-1]) / 1000.0
            except ValueError:
                pass

    return 0.0


def get_best_ollama_model(config) -> Optional["OllamaProvider"]:
    """
    Find the best (largest tool-capable) Ollama model for fallback use.

    Priority:
      1. OLLAMA_FALLBACK_MODEL env var (explicit override)
      2. Largest pulled model that supports tool calling
      3. None if no suitable model found

    Returns an OllamaProvider instance or None.
    """
    base_url = config.llm.ollama_base_url.rstrip("/")

    # 1. Explicit override
    override = os.environ.get("OLLAMA_FALLBACK_MODEL")
    if override:
        try:
            provider = OllamaProvider(override, base_url)
            logger.info(f"Ollama fallback (env override): {provider.name}")
            return provider
        except Exception as e:
            logger.warning(f"OLLAMA_FALLBACK_MODEL={override} failed: {e}")

    # 2. Auto-detect: query Ollama for pulled models, pick largest tool-capable
    try:
        req = urllib.request.Request(f"{base_url}/api/tags", method="GET")
        with urllib.request.urlopen(req, timeout=3) as resp:
            data = json.loads(resp.read().decode())

        candidates = []
        for model_info in data.get("models", []):
            name = model_info.get("name", "")
            if not name:
                continue
            # Check tool support via static list (no network probe)
            probe = OllamaProvider.__new__(OllamaProvider)
            probe.model = name
            probe.base_url = base_url
            probe.client = None
            if not probe.supports_tools:
                continue
            size = _extract_param_size(name, model_info)
            candidates.append((size, name, model_info))

        if candidates:
            # Sort by parameter size descending, pick largest
            candidates.sort(key=lambda c: c[0], reverse=True)
            best_name = candidates[0][1]
            provider = OllamaProvider(best_name, base_url)
            logger.info(
                f"Ollama fallback (auto-detected): {provider.name} "
                f"({candidates[0][0]:.1f}B params, {len(candidates)} candidates)"
            )
            return provider

    except Exception as e:
        logger.debug(f"Ollama fallback discovery failed: {e}")

    return None
