"""
Ollama / Local Model Proxy.

Discovers local model servers (Ollama, LM Studio, any OpenAI-compatible)
and reports available models to Neura Cloud for routing.
"""
import asyncio
from typing import Dict, Any, List, Optional

import httpx

from .config import RuntimeConfig


class OllamaProxy:
    """Discovers and proxies local LLM endpoints."""

    def __init__(self, config: RuntimeConfig):
        self.config = config
        self._models: Dict[str, Dict[str, Any]] = {}  # model_id -> model info
        self._endpoints: Dict[str, str] = {}  # model_id -> base_url

    async def discover(self) -> List[Dict[str, Any]]:
        """Discover available local models from all endpoints."""
        self._models.clear()
        self._endpoints.clear()

        # Try Ollama
        await self._discover_ollama()

        # Try LM Studio
        await self._discover_openai_compatible(
            self.config.lm_studio_url,
            provider="lm_studio",
        )

        return list(self._models.values())

    async def _discover_ollama(self):
        """Discover models from Ollama's /api/tags endpoint."""
        try:
            async with httpx.AsyncClient(timeout=5) as client:
                resp = await client.get(f"{self.config.ollama_url}/api/tags")
                if resp.status_code == 200:
                    data = resp.json()
                    for model in data.get("models", []):
                        model_id = f"local:ollama:{model['name']}"
                        self._models[model_id] = {
                            "id": model_id,
                            "name": model["name"],
                            "provider": "ollama",
                            "size": model.get("size", 0),
                            "parameter_size": model.get("details", {}).get("parameter_size", ""),
                            "quantization": model.get("details", {}).get("quantization_level", ""),
                        }
                        self._endpoints[model_id] = self.config.ollama_url
        except (httpx.ConnectError, httpx.TimeoutException):
            pass

    async def _discover_openai_compatible(self, base_url: str, provider: str):
        """Discover models from an OpenAI-compatible /v1/models endpoint."""
        try:
            async with httpx.AsyncClient(timeout=5) as client:
                resp = await client.get(f"{base_url}/v1/models")
                if resp.status_code == 200:
                    data = resp.json()
                    for model in data.get("data", []):
                        model_id = f"local:{provider}:{model['id']}"
                        self._models[model_id] = {
                            "id": model_id,
                            "name": model["id"],
                            "provider": provider,
                            "owned_by": model.get("owned_by", ""),
                        }
                        self._endpoints[model_id] = base_url
        except (httpx.ConnectError, httpx.TimeoutException):
            pass

    async def chat_completion(
        self,
        model_id: str,
        messages: List[Dict[str, str]],
        stream: bool = False,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Proxy a chat completion request to the local model.

        Supports both Ollama and OpenAI-compatible endpoints.
        """
        base_url = self._endpoints.get(model_id)
        if not base_url:
            return {"error": f"Model {model_id} not found locally"}

        # Extract actual model name from our namespaced ID
        # Format: local:<provider>:<model_name>
        parts = model_id.split(":", 2)
        model_name = parts[2] if len(parts) == 3 else model_id
        provider = parts[1] if len(parts) >= 2 else "unknown"

        if provider == "ollama":
            return await self._ollama_chat(base_url, model_name, messages, stream, **kwargs)
        else:
            return await self._openai_chat(base_url, model_name, messages, stream, **kwargs)

    async def _ollama_chat(
        self,
        base_url: str,
        model: str,
        messages: List[Dict[str, str]],
        stream: bool,
        **kwargs,
    ) -> Dict[str, Any]:
        """Chat via Ollama's OpenAI-compatible endpoint."""
        try:
            async with httpx.AsyncClient(timeout=300) as client:
                resp = await client.post(
                    f"{base_url}/v1/chat/completions",
                    json={
                        "model": model,
                        "messages": messages,
                        "stream": stream,
                        **kwargs,
                    },
                )
                return resp.json()
        except Exception as e:
            return {"error": str(e)}

    async def _openai_chat(
        self,
        base_url: str,
        model: str,
        messages: List[Dict[str, str]],
        stream: bool,
        **kwargs,
    ) -> Dict[str, Any]:
        """Chat via OpenAI-compatible endpoint."""
        try:
            async with httpx.AsyncClient(timeout=300) as client:
                resp = await client.post(
                    f"{base_url}/v1/chat/completions",
                    json={
                        "model": model,
                        "messages": messages,
                        "stream": stream,
                        **kwargs,
                    },
                )
                return resp.json()
        except Exception as e:
            return {"error": str(e)}

    def get_models(self) -> List[Dict[str, Any]]:
        """Return cached model list."""
        return list(self._models.values())
