#!/usr/bin/env python3
"""
Ollama Backend - HTTP client for local Ollama server.

WO-ECHO-AGENT-LOOP-01: Provides tool-calling support via Ollama's API.

Supports:
- Text generation
- Tool calling (when model supports it)
- Streaming responses
- Fallback to markdown-block tool parsing for models without native support

Usage:
    from echo_cog.backend_ollama import chat_with_tools, OllamaConfig

    config = OllamaConfig(model="qwen2.5-coder:7b")
    response = chat_with_tools(
        messages=[{"role": "user", "content": "Search for the Se formula"}],
        tools=[...],  # OpenAI-compatible tool schemas
        config=config,
    )
"""

from __future__ import annotations

import json
import os
import re
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import requests


# Configuration
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
DEFAULT_MODEL = os.environ.get("ECHO_AGENT_MODEL", "qwen2.5-coder:7b")
DEFAULT_TEMPERATURE = float(os.environ.get("ECHO_AGENT_TEMPERATURE", "0.1"))
DEFAULT_TIMEOUT = int(os.environ.get("ECHO_AGENT_TIMEOUT", "120"))


@dataclass
class OllamaConfig:
    """Configuration for Ollama backend."""
    model: str = DEFAULT_MODEL
    temperature: float = DEFAULT_TEMPERATURE
    timeout: int = DEFAULT_TIMEOUT
    url: str = OLLAMA_URL
    num_ctx: int = 4096  # Context window
    num_predict: int = 1024  # Max tokens to generate
    use_native_tools: bool = True  # Try native tool calling first
    fallback_to_markdown: bool = True  # Fall back to markdown parsing if native fails


@dataclass
class OllamaResponse:
    """Response from Ollama chat completion."""
    content: str
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
    done: bool = False
    model: str = ""
    total_duration_ns: int = 0
    eval_count: int = 0
    error: Optional[str] = None

    @property
    def duration_ms(self) -> float:
        return self.total_duration_ns / 1_000_000

    def to_dict(self) -> Dict[str, Any]:
        return {
            "content": self.content,
            "tool_calls": self.tool_calls,
            "done": self.done,
            "model": self.model,
            "duration_ms": self.duration_ms,
            "eval_count": self.eval_count,
            "error": self.error,
        }


def check_ollama_health(url: str = OLLAMA_URL) -> Tuple[bool, str]:
    """
    Check if Ollama server is healthy.

    Returns:
        (is_healthy, message)
    """
    try:
        response = requests.get(f"{url}/api/version", timeout=5)
        if response.status_code == 200:
            version = response.json().get("version", "unknown")
            return True, f"Ollama v{version} running"
        return False, f"Ollama returned status {response.status_code}"
    except requests.exceptions.ConnectionError:
        return False, f"Cannot connect to Ollama at {url}"
    except requests.exceptions.Timeout:
        return False, "Ollama health check timed out"
    except Exception as e:
        return False, f"Health check failed: {e}"


def list_models(url: str = OLLAMA_URL) -> List[str]:
    """List available models in Ollama."""
    try:
        response = requests.get(f"{url}/api/tags", timeout=10)
        if response.status_code == 200:
            models = response.json().get("models", [])
            return [m.get("name", "") for m in models]
        return []
    except Exception:
        return []


def parse_tool_calls_from_markdown(text: str) -> List[Dict[str, Any]]:
    """
    Parse tool calls from markdown code blocks.

    This is the fallback for models that don't support native tool calling.
    Looks for blocks like:
        ```tool_call
        {"name": "codebase", "arguments": {"query": "Se formula"}}
        ```

    Also handles:
        ```json
        {"tool_call": {"name": "...", ...}}
        ```

    Returns:
        List of parsed tool call dicts with 'name' and 'arguments'
    """
    calls = []

    # Pattern 1: ```tool_call blocks
    tool_call_pattern = r'```tool_call\s*\n?(.*?)\n?```'
    matches = re.findall(tool_call_pattern, text, re.DOTALL | re.IGNORECASE)

    for match in matches:
        try:
            data = json.loads(match.strip())
            if "name" in data:
                calls.append({
                    "name": data["name"],
                    "arguments": data.get("arguments", {}),
                })
        except json.JSONDecodeError:
            continue

    # Pattern 2: Inline tool call JSON
    if not calls:
        inline_pattern = r'\{"name":\s*"(\w+)"[^}]*"arguments":\s*(\{[^}]*\})\s*\}'
        matches = re.findall(inline_pattern, text, re.DOTALL)
        for name, args_str in matches:
            try:
                args = json.loads(args_str)
                calls.append({"name": name, "arguments": args})
            except json.JSONDecodeError:
                calls.append({"name": name, "arguments": {}})

    # Pattern 3: Function call syntax
    if not calls:
        func_pattern = r'<function_call>\s*(\{.*?\})\s*</function_call>'
        matches = re.findall(func_pattern, text, re.DOTALL)
        for match in matches:
            try:
                data = json.loads(match)
                if "name" in data:
                    calls.append({
                        "name": data["name"],
                        "arguments": data.get("arguments", {}),
                    })
            except json.JSONDecodeError:
                continue

    return calls


def chat_with_tools(
    messages: List[Dict[str, Any]],
    tools: List[Dict[str, Any]],
    config: Optional[OllamaConfig] = None,
) -> OllamaResponse:
    """
    Call Ollama with tool definitions.

    Supports two modes:
    1. Native tool calling (for models that support it)
    2. Markdown fallback (parse ```tool_call blocks from response)

    If native tool calling fails with 400 error (unsupported), automatically
    falls back to markdown parsing mode.

    Args:
        messages: List of message dicts with 'role' and 'content'
        tools: List of OpenAI-compatible tool schemas
        config: OllamaConfig (uses defaults if not provided)

    Returns:
        OllamaResponse with content and parsed tool_calls
    """
    if config is None:
        config = OllamaConfig()

    # Build request payload
    payload = {
        "model": config.model,
        "messages": messages,
        "stream": False,
        "options": {
            "temperature": config.temperature,
            "num_ctx": config.num_ctx,
            "num_predict": config.num_predict,
        },
    }

    # Try native tool calling first if enabled
    use_native = config.use_native_tools and tools
    if use_native:
        payload["tools"] = tools

    try:
        response = requests.post(
            f"{config.url}/api/chat",
            json=payload,
            timeout=config.timeout,
        )

        # Check for 400 error (model doesn't support tools)
        if response.status_code == 400 and use_native and config.fallback_to_markdown:
            # Retry without tools parameter - model will use markdown blocks
            del payload["tools"]
            response = requests.post(
                f"{config.url}/api/chat",
                json=payload,
                timeout=config.timeout,
            )

        response.raise_for_status()
        data = response.json()
    except requests.exceptions.ConnectionError:
        return OllamaResponse(
            content="",
            error=f"Cannot connect to Ollama at {config.url}",
        )
    except requests.exceptions.Timeout:
        return OllamaResponse(
            content="",
            error=f"Ollama request timed out after {config.timeout}s",
        )
    except requests.exceptions.HTTPError as e:
        return OllamaResponse(
            content="",
            error=f"Ollama HTTP error: {e}",
        )
    except json.JSONDecodeError:
        return OllamaResponse(
            content="",
            error="Invalid JSON response from Ollama",
        )

    # Parse response
    message = data.get("message", {})
    content = message.get("content", "")

    # Check for native tool calls first
    tool_calls = []
    native_tool_calls = message.get("tool_calls", [])

    if native_tool_calls:
        # Ollama native tool calling
        for tc in native_tool_calls:
            func = tc.get("function", {})
            tool_calls.append({
                "name": func.get("name", ""),
                "arguments": func.get("arguments", {}),
            })
    else:
        # Fall back to markdown parsing
        tool_calls = parse_tool_calls_from_markdown(content)

    # Check for done signal
    done = False
    if '{"done": true' in content.lower() or '"done":true' in content.lower():
        done = True
    if not tool_calls and content and not done:
        # If no tool calls and content present, might be final answer
        # Check if it looks like a summary/answer
        if len(content) > 50 and not any(p in content for p in ['```', '"name":', 'tool_call']):
            done = True

    return OllamaResponse(
        content=content,
        tool_calls=tool_calls,
        done=done,
        model=data.get("model", config.model),
        total_duration_ns=data.get("total_duration", 0),
        eval_count=data.get("eval_count", 0),
    )


def chat_simple(
    prompt: str,
    system: Optional[str] = None,
    config: Optional[OllamaConfig] = None,
) -> OllamaResponse:
    """
    Simple chat without tools.

    Args:
        prompt: User prompt
        system: Optional system prompt
        config: OllamaConfig

    Returns:
        OllamaResponse
    """
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    return chat_with_tools(messages, tools=[], config=config)


def pull_model(model: str, url: str = OLLAMA_URL) -> Tuple[bool, str]:
    """
    Pull a model from Ollama registry.

    Returns:
        (success, message)
    """
    try:
        response = requests.post(
            f"{url}/api/pull",
            json={"name": model},
            timeout=600,  # Models can be large
            stream=True,
        )

        # Stream progress
        for line in response.iter_lines():
            if line:
                try:
                    data = json.loads(line)
                    status = data.get("status", "")
                    if "error" in data:
                        return False, data["error"]
                except json.JSONDecodeError:
                    continue

        return True, f"Model {model} pulled successfully"
    except Exception as e:
        return False, f"Failed to pull model: {e}"


# Convenience function for quick testing
def quick_chat(prompt: str, model: str = DEFAULT_MODEL) -> str:
    """Quick one-shot chat, returns just the content string."""
    config = OllamaConfig(model=model)
    response = chat_simple(prompt, config=config)
    if response.error:
        return f"Error: {response.error}"
    return response.content
