"""Ollama function-calling backend for ORBIT Assist.

Uses Ollama's OpenAI-compatible API via the openai Python client.
Supports tool/function calling with models like Gemma 4 and Llama 3.
Falls back to prompt-based tool parsing if the model doesn't support
native function calling.
"""

from __future__ import annotations

import json
import logging

from orbit.assist.tools import TOOL_DECLARATIONS, execute_tool

logger = logging.getLogger(__name__)

# Maximum size (chars) of a single tool result before truncation
_MAX_TOOL_RESULT_CHARS = 12_000


def _convert_tool_declarations() -> list[dict]:
    """Convert Gemini-style function declarations to OpenAI tool format."""
    tools = []
    for decl in TOOL_DECLARATIONS:
        tool = {
            "type": "function",
            "function": {
                "name": decl["name"],
                "description": decl.get("description", ""),
                "parameters": decl.get("parameters", {"type": "object", "properties": {}}),
            },
        }
        tools.append(tool)
    return tools


class OllamaBackend:
    """Ollama function-calling backend via OpenAI-compatible API.

    Implements the same interface as GeminiBackend: run_turn() and reset().
    """

    def __init__(self, model: str, base_url: str = "http://localhost:11434/v1"):
        self.model = model
        self.base_url = base_url
        self.history: list[dict] = []
        self._tools = _convert_tool_declarations()

    def run_turn(
        self,
        user_message: str,
        system_prompt: str,
        on_tool_call: callable | None = None,
    ) -> str:
        """Run one turn of the agentic loop.

        May involve multiple tool calls before returning a final text response.
        """
        try:
            from openai import OpenAI
        except ImportError:
            return (
                "The openai Python package is required for Ollama assist. "
                "Install it with: pip install openai"
            )

        client = OpenAI(base_url=self.base_url, api_key="ollama")

        self.history.append({"role": "user", "content": user_message})

        max_iterations = 15
        for _iteration in range(max_iterations):
            messages = [{"role": "system", "content": system_prompt}] + self.history

            try:
                response = client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    tools=self._tools,
                    temperature=0.1,
                    max_tokens=4096,
                    timeout=120,
                )
            except Exception as e:
                error_msg = f"Ollama call failed: {e}"
                logger.debug(error_msg)
                self.history.append({"role": "assistant", "content": error_msg})
                return error_msg

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

            # Check for tool calls
            if message.tool_calls:
                # Add assistant message with tool calls to history
                self.history.append({
                    "role": "assistant",
                    "content": message.content or "",
                    "tool_calls": [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.function.name,
                                "arguments": tc.function.arguments,
                            },
                        }
                        for tc in message.tool_calls
                    ],
                })

                # Execute each tool call
                for tc in message.tool_calls:
                    tool_name = tc.function.name
                    try:
                        tool_args = json.loads(tc.function.arguments)
                    except json.JSONDecodeError:
                        tool_args = {}

                    logger.debug("Tool call: %s(%s)", tool_name, tool_args)
                    result = execute_tool(tool_name, tool_args)

                    if on_tool_call:
                        on_tool_call(tool_name, tool_args, result)

                    # Truncate large results
                    result_str = json.dumps(result, default=str)
                    if len(result_str) > _MAX_TOOL_RESULT_CHARS:
                        result_str = result_str[:_MAX_TOOL_RESULT_CHARS] + "...(truncated)"

                    self.history.append({
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "content": result_str,
                    })
            else:
                # Final text response
                final_text = message.content or "(No response)"
                self.history.append({"role": "assistant", "content": final_text})
                return final_text

        return "Reached maximum reasoning steps. Here's a summary of what I found."

    def reset(self) -> None:
        """Clear conversation history."""
        self.history.clear()
