"""Local model provider — supports Ollama and LM Studio via OpenAI-compatible API."""

from __future__ import annotations

import json

from hatchdx.agent.providers.base import (
    AgentResponse,
    ProviderError,
    ToolCall,
)
from hatchdx.agent.providers.openai import _openai_messages, mcp_tools_to_openai


# ---------------------------------------------------------------------------
# Default endpoints
# ---------------------------------------------------------------------------

_DEFAULT_BASE_URLS = {
    "ollama": "http://localhost:11434/v1",
    "lm-studio": "http://localhost:1234/v1",
}

DEFAULT_BASE_URL = "http://localhost:11434/v1"


def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Local models are free — always returns 0."""
    return 0.0


# ---------------------------------------------------------------------------
# Provider
# ---------------------------------------------------------------------------


class LocalProvider:
    """Local model provider for Ollama and LM Studio (OpenAI-compatible API).

    Uses the OpenAI SDK with a custom base_url pointed at a local server.
    No API key is required — a dummy key is used to satisfy the SDK.
    """

    def __init__(
        self,
        model: str,
        base_url: str | None = None,
    ) -> None:
        self.model = model
        self._base_url = base_url or DEFAULT_BASE_URL
        self._client = None
        self._supports_tools: bool | None = None

    def _get_client(self):
        """Lazy-init the OpenAI client pointed at local server."""
        if self._client is not None:
            return self._client

        try:
            import openai
        except ImportError:
            raise ProviderError(
                "OpenAI SDK not installed (required for local provider).\n\n"
                "Install it with:\n"
                "  uv add hatchdx --extra openai"
            )

        self._client = openai.AsyncOpenAI(
            api_key="not-needed",
            base_url=self._base_url,
        )
        return self._client

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict],
        system: str,
        max_tokens: int,
        temperature: float,
    ) -> AgentResponse:
        """Send a chat request to the local model."""
        client = self._get_client()
        openai_messages = _openai_messages(messages, system)

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

        # Only include tools if the model supports them
        if tools and self._supports_tools is not False:
            openai_tools = mcp_tools_to_openai(tools)
            kwargs["tools"] = openai_tools

        try:
            import openai as openai_module

            response = await client.chat.completions.create(**kwargs)
            # If we got here with tools, they're supported
            if tools and self._supports_tools is None:
                self._supports_tools = True

        except Exception as e:
            error_str = str(e).lower()
            # Detect tool calling not supported
            if tools and ("tool" in error_str or "function" in error_str):
                self._supports_tools = False
                # Retry without tools
                kwargs.pop("tools", None)
                try:
                    response = await client.chat.completions.create(**kwargs)
                except Exception as retry_e:
                    raise ProviderError(
                        f"Local model error: {retry_e}\n\n"
                        f"Is your local model server running at {self._base_url}?"
                    ) from retry_e
            elif "connect" in error_str or "refused" in error_str:
                raise ProviderError(
                    f"Cannot connect to local model server at {self._base_url}\n\n"
                    "Make sure your model server is running:\n"
                    "  Ollama: ollama serve\n"
                    "  LM Studio: Start the local server in LM Studio"
                ) from e
            else:
                raise ProviderError(f"Local model error: {e}") from e

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

        text = message.content

        tool_calls = []
        if message.tool_calls:
            for tc in message.tool_calls:
                try:
                    arguments = json.loads(tc.function.arguments)
                except (json.JSONDecodeError, TypeError):
                    arguments = {}
                tool_calls.append(
                    ToolCall(
                        id=tc.id,
                        name=tc.function.name,
                        arguments=arguments,
                    )
                )

        stop_reason = "end_turn"
        if choice.finish_reason == "tool_calls":
            stop_reason = "tool_use"
        elif choice.finish_reason == "length":
            stop_reason = "max_tokens"

        usage = {"input": 0, "output": 0}
        if response.usage:
            usage["input"] = response.usage.prompt_tokens or 0
            usage["output"] = response.usage.completion_tokens or 0

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