"""Ollama LLM plugin — uses Ollama's OpenAI-compatible API.

Ollama serves local LLMs at ``http://localhost:11434/v1`` and speaks the
OpenAI chat completions protocol.  This plugin reuses the ``openai`` Python
package pointed at the Ollama endpoint, so it supports both ``generate`` and
``generate_with_tools`` (for models that support tool calling, e.g. Llama 3).

Configuration via environment variables or settings.json:
  - ``OLLAMA_BASE_URL`` — defaults to ``http://localhost:11434``
  - ``OLLAMA_MODEL`` — the model to use (e.g. ``llama3``, ``mistral``)
"""

from __future__ import annotations

import logging
import os

from vastian.config import Settings
from vastian.providers.llm_provider import BaseLLMProvider

logger = logging.getLogger(__name__)

_DEFAULT_OLLAMA_URL = "http://localhost:11434"


class OllamaProvider(BaseLLMProvider):
    """Ollama LLM provider using the OpenAI-compatible API."""

    def __init__(self, settings: Settings) -> None:
        # Use regular attributes (not @property) — avoids PyInstaller frozen-module
        # descriptor issues and keeps parity with OpenAI/Anthropic providers.
        self.settings = settings
        model = os.environ.get("OLLAMA_MODEL", "") or "llama3"
        self.default_model = model
        self.fast_model = model

        from openai import AsyncOpenAI

        base_url = os.environ.get("OLLAMA_BASE_URL", _DEFAULT_OLLAMA_URL)
        # Ollama's OpenAI-compatible endpoint is at /v1
        if not base_url.endswith("/v1"):
            base_url = base_url.rstrip("/") + "/v1"

        self.client = AsyncOpenAI(
            base_url=base_url,
            api_key="ollama",  # Ollama doesn't need a real key
        )
        self.last_usage: dict[str, int] = {}
        logger.info("OllamaProvider: using model %s at %s", model, base_url)

    def _extract_usage(self, response: object, model: str) -> None:
        """Extract token usage from response."""
        usage = getattr(response, "usage", None)
        if usage:
            self.last_usage = {
                "provider": "ollama",
                "model": model,
                "input_tokens": getattr(usage, "prompt_tokens", 0) or 0,
                "output_tokens": getattr(usage, "completion_tokens", 0) or 0,
            }
        else:
            self.last_usage = {}

    async def generate(
        self,
        system_prompt: str,
        messages: list[dict[str, str]],
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
    ) -> str:
        model = model or self.default_model
        full_messages = [{"role": "system", "content": system_prompt}] + messages

        response = await self.client.chat.completions.create(
            model=model,
            messages=full_messages,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        self._extract_usage(response, model)
        return response.choices[0].message.content or ""

    async def generate_with_tools(
        self,
        system_prompt: str,
        messages: list[dict[str, str]],
        tools: list[dict],
        model: str | None = None,
        temperature: float = 0.7,
    ) -> dict:
        model = model or self.default_model
        full_messages = [{"role": "system", "content": system_prompt}] + messages

        response = await self.client.chat.completions.create(
            model=model,
            messages=full_messages,
            tools=tools,
            temperature=temperature,
        )
        self._extract_usage(response, model)
        choice = response.choices[0]

        result: dict = {"content": choice.message.content or ""}
        if choice.message.tool_calls:
            result["tool_calls"] = [
                {
                    "id": tc.id,
                    "function": tc.function.name,
                    "arguments": tc.function.arguments,
                }
                for tc in choice.message.tool_calls
            ]
        return result


def register() -> None:
    """Register the Ollama LLM plugin with the loader."""
    from vastian.loader import register as loader_register

    loader_register("ollama_llm", OllamaProvider)
