"""Thin client for interacting with a local Ollama instance."""

from __future__ import annotations

import json
import logging
import re
from typing import Any, Dict, List, Optional

import httpx

logger = logging.getLogger(__name__)


class OllamaClient:
    """Wrapper around the Ollama /api/chat endpoint."""

    def __init__(self, host: str, model: str, timeout_s: int = 60) -> None:
        self.host = host.rstrip("/")
        self.model = model
        self.timeout = timeout_s

    def chat_json(
        self, messages: List[Dict[str, Any]], retry_strict: bool = True
    ) -> Dict[str, Any]:
        """
        Call Ollama's chat API expecting a JSON response payload.

        Auto-retries once with strict schema prompt if parse fails.
        Uses low temperature for consistent structured output.
        """
        payload = {
            "model": self.model,
            "messages": messages,
            "format": "json",
            "stream": False,
            "options": {
                "temperature": 0.1,  # Low temp for consistency
                "top_p": 0.9,
                "num_ctx": 4096,
                "num_predict": 256,  # Limit response length
            },
        }

        with httpx.Client(timeout=self.timeout) as client:
            response = client.post(f"{self.host}/api/chat", json=payload)
            response.raise_for_status()
            data = response.json()

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

        # FIX: Ollama can return content as dict if JSON is valid
        # Convert dict to JSON string so we can parse it properly
        if isinstance(content, dict):
            content = json.dumps(content)
        elif not isinstance(content, str):
            logger.warning(f"Ollama returned unexpected content type: {type(content)}")
            return {"_parse_error": True, "_raw": str(content)}

        content = content.strip()
        if not content:
            logger.warning("Ollama returned empty content")
            return {"_parse_error": True, "_raw": ""}

        # Try parsing
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            # Try salvaging from code fences or embedded JSON
            salvaged = self._salvage_json(content)
            if salvaged:
                return salvaged

            # Last resort: strict reprompt (one retry)
            if retry_strict:
                logger.warning("JSON parse failed, retrying with strict prompt (preserving context)")
                # FIX: Preserve original context by appending strict instruction to last message
                retry_messages = messages.copy()
                if retry_messages and retry_messages[-1].get("role") == "user":
                    retry_messages[-1] = {
                        "role": "user",
                        "content": retry_messages[-1]["content"] + '\n\nRespond ONLY with valid JSON: {"decision":"approve|reject|unsure","suggested_domain":"example.com or null","confidence":0.00,"reasoning":"..."}',
                    }
                else:
                    # Fallback if no user message
                    retry_messages.append({
                        "role": "user",
                        "content": 'Respond ONLY with valid JSON: {"decision":"approve|reject|unsure","suggested_domain":"example.com or null","confidence":0.00,"reasoning":"..."}',
                    })
                return self.chat_json(retry_messages, retry_strict=False)  # No double-retry

            logger.error(f"JSON parse failed after retry: {content[:200]}")
            return {"_parse_error": True, "_raw": content}

    def _salvage_json(self, text: str) -> Optional[Dict[str, Any]]:
        """Extract JSON from markdown fences or prose."""
        # Strip code fences
        cleaned = re.sub(
            r"^```(?:json)?\s*|\s*```$", "", text.strip(), flags=re.I | re.M
        )

        # Find first {...}
        start = cleaned.find("{")
        end = cleaned.rfind("}")
        if start != -1 and end != -1 and end > start:
            try:
                return json.loads(cleaned[start : end + 1])
            except json.JSONDecodeError:
                pass

        return None

    def get_model_info(self) -> Dict[str, Any]:
        """Fetch metadata about the configured model (name, digest, etc.)."""
        payload = {"name": self.model}
        with httpx.Client(timeout=self.timeout) as client:
            response = client.post(f"{self.host}/api/show", json=payload)
            response.raise_for_status()
            return response.json()
