"""Local AI Synthesizer — Directly interfaces with Ollama to generate code and configs."""

import json
import re
from typing import Any

try:
    import httpx
except ImportError:
    httpx = None  # type: ignore

from nikame.exceptions import NikameError
from nikame.utils.logger import get_logger

_log = get_logger("ai_synthesizer")


class LocalOllamaClient:
    """Client for interacting with local Ollama instance."""

    def __init__(self, base_url: str = "http://localhost:11434", model: str = "llama3"):
        self.base_url = base_url.rstrip("/")
        self.model = model
        
        if httpx is None:
            raise NikameError(
                "The 'httpx' package is required for the AI Copilot. "
                "Please install it using 'pip install httpx'."
            )

    def is_available(self) -> bool:
        """Check if Ollama server is running."""
        try:
            response = httpx.get(f"{self.base_url}/", timeout=2.0)
            return response.status_code == 200
        except httpx.RequestError:
            return False

    def ensure_model(self) -> None:
        """Ensure the target model is available locally. If not, attempt to pull it."""
        try:
            response = httpx.get(f"{self.base_url}/api/tags", timeout=5.0)
            response.raise_for_status()
            models = [m["name"] for m in response.json().get("models", [])]
            
            # Ollama models often have the :tag suffix. If user asks for llama3, match llama3:latest
            has_model = any(m == self.model or m.startswith(f"{self.model}:") for m in models)
            
            if not has_model:
                _log.info(f"Model '{self.model}' not found locally. Pulling from Ollama registry...")
                with httpx.stream("POST", f"{self.base_url}/api/pull", json={"name": self.model}) as r:
                    r.raise_for_status()
                    for line in r.iter_lines():
                        if line:
                            data = json.loads(line)
                            if "status" in data:
                                _log.debug(f"Pulling {self.model}: {data['status']}")
                _log.info(f"Model '{self.model}' downloaded successfully.")
                
        except httpx.RequestError as e:
            raise NikameError(f"Failed to communicate with Ollama: {e}")

    def generate(self, prompt: str, system: str = "", format_json: bool = False, temperature: float = 0.2) -> str:
        """Generate text from Ollama."""
        payload: dict[str, Any] = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
            }
        }
        if system:
            payload["system"] = system
            
        if format_json:
            payload["format"] = "json"

        try:
            response = httpx.post(f"{self.base_url}/api/generate", json=payload, timeout=120.0)
            response.raise_for_status()
            data = response.json()
            return data.get("response", "")
        except httpx.RequestError as e:
            raise NikameError(f"Ollama generation failed: {e}")
            

class AISynthesizer:
    """High-level synthesizer for generating configs and code."""

    def __init__(self, model: str = "llama3"):
        self.client = LocalOllamaClient(model=model)
        
    def check_readiness(self) -> None:
        """Verify the local AI engine is ready to synthesize."""
        if not self.client.is_available():
            raise NikameError(
                "Ollama is not running locally. Please install Ollama "
                "(https://ollama.com) and ensure the daemon is running on port 11434."
            )
        self.client.ensure_model()

    def extract_json(self, response: str) -> dict[str, Any]:
        """Extract a JSON object from the LLM response."""
        try:
            return json.loads(response) # type: ignore
        except json.JSONDecodeError:
            # Fallback regex extraction if the model wrapped it in markdown
            match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', response, re.DOTALL)
            if match:
                try:
                    return json.loads(match.group(1)) # type: ignore
                except json.JSONDecodeError:
                    pass
            raise NikameError("Failed to parse valid JSON from AI response.")
            
    def extract_python(self, response: str) -> str:
        """Extract Python code block from the LLM response."""
        match = re.search(r'```(?:python)?\s*(.*?)\s*```', response, re.DOTALL)
        if match:
            return match.group(1).strip()
        # If no markdown block, assume the whole response is code
        return response.strip()

    def draft_config(self, system_prompt: str, user_request: str) -> dict[str, Any]:
        """Generate a NikameConfig blueprint."""
        self.check_readiness()
        response = self.client.generate(
            prompt=user_request,
            system=system_prompt,
            format_json=True,
            temperature=0.1
        )
        return self.extract_json(response)

    def synthesize_code(self, system_prompt: str, request: str, context: str = "") -> str:
        """Generate valid Python code."""
        self.check_readiness()
        prompt = f"Context:\n{context}\n\nTask:\n{request}" if context else request
        response = self.client.generate(
            prompt=prompt,
            system=system_prompt,
            temperature=0.2
        )
        return self.extract_python(response)
