"""
Ollama implementation for dazllm
"""

import json
import re
from typing import Type

import keyring
import requests
from jsonschema import ValidationError, validate
from pydantic import BaseModel

from .core import (
    ConfigurationError,
    Conversation,
    DazLlmError,
    Llm,
    LlmResponse,
    _run_sync,
)


class LlmOllama(Llm):
    """Ollama implementation"""

    # Models that don't support the 'format' parameter and need special handling
    MODELS_WITHOUT_FORMAT_SUPPORT = [
        "gpt-oss:20b",  # Known issue: uses "Harmony" response format, unreliable with JSON schema
        # Add more models here as needed
    ]

    def __init__(self, model: str):
        super().__init__(model)
        self.base_url = self._get_base_url()
        self.headers = {"Content-Type": "application/json"}
        self._ensure_model_available()
        self._supports_format = self._check_format_support()

    @staticmethod
    def default_model() -> str:
        """Default model for Ollama"""
        return "mistral-small"

    @staticmethod
    def default_for_type(model_type: str) -> str:
        """Get default model for a given type"""
        defaults = {
            "local_small": "phi3",
            "local_medium": "mistral-small",
            "local_large": "qwen3:32b",
            "paid_cheap": None,
            "paid_best": None,
        }
        return defaults.get(model_type)

    @staticmethod
    def capabilities() -> set[str]:
        """Return set of capabilities this provider supports"""
        return {"chat", "structured"}

    @staticmethod
    def supported_models() -> list[str]:
        """Return list of models this provider supports"""
        url = keyring.get_password("dazllm", "ollama_url") or "http://127.0.0.1:11434"
        response = requests.get(f"{url}/api/tags")
        response.raise_for_status()
        models = response.json().get("models", [])
        return [model["name"] for model in models] if models else []

    @staticmethod
    def check_config():
        """Check if Ollama is properly configured"""
        try:
            base_url = LlmOllama._get_base_url_static()
            response = requests.get(f"{base_url}/api/version")
            response.raise_for_status()
        except Exception as e:
            raise ConfigurationError(f"Ollama not accessible: {e}") from e

    def _get_base_url(self) -> str:
        """Get Ollama base URL from keyring or default"""
        return self._get_base_url_static()

    @staticmethod
    def _get_base_url_static() -> str:
        """Static version of _get_base_url"""
        return keyring.get_password("dazllm", "ollama_url") or "http://127.0.0.1:11434"

    def _ensure_model_available(self):
        """Ensure model is available, pull if necessary"""
        if not self._is_model_available() and not self._pull_model():
            raise ConfigurationError(f"Failed to pull model '{self.model}' from Ollama registry")

    def _is_model_available(self) -> bool:
        """Check if model is available locally"""
        try:
            response = requests.get(f"{self.base_url}/api/tags", headers=self.headers)
            response.raise_for_status()
            models = response.json().get("models", [])
            return any(model["name"].startswith(self.model) for model in models)
        except (requests.exceptions.RequestException, ValueError, KeyError):
            return False

    def _pull_model(self) -> bool:
        """Pull model from Ollama registry"""
        try:
            response = requests.post(
                f"{self.base_url}/api/pull",
                json={"model": self.model},
                headers=self.headers,
            )
            response.raise_for_status()
            return True
        except (requests.exceptions.RequestException, ValueError):
            return False

    def _check_format_support(self) -> bool:
        """Check if the model supports the 'format' parameter"""
        # Check against our list of known models without format support
        for model_pattern in self.MODELS_WITHOUT_FORMAT_SUPPORT:
            if self.model == model_pattern or self.model.startswith(model_pattern.split(":")[0]):
                return False
        return True

    def _normalize_conversation(self, conversation: Conversation) -> list:
        """Convert conversation to Ollama message format"""
        return [{"role": "user", "content": conversation}] if isinstance(conversation, str) else conversation

    def chat(self, conversation: Conversation, force_json: bool = False) -> LlmResponse:
        """Sync wrapper; blocks on chat_async."""
        return _run_sync(self.chat_async(conversation, force_json))

    def chat_structured(
        self, conversation: Conversation, schema: Type[BaseModel], context_size: int = 0
    ) -> LlmResponse:
        """Sync wrapper; blocks on chat_structured_async."""
        return _run_sync(self.chat_structured_async(conversation, schema, context_size))

    def _chat_structured_with_format(
        self,
        messages: list,
        schema: Type[BaseModel],
        schema_json: dict,
        context_size: int,
    ) -> BaseModel:
        """Standard structured output using format parameter"""
        system_message = {
            "role": "system",
            "content": (
                f"All responses should be strictly in JSON obeying this schema: {schema_json} "
                "with no accompanying text or delimiters. Do not include the schema in the output. "
                "We want the shortest possible output with no explanations. If there is source code or "
                "other technical output, pay very close attention to proper escaping so the result is valid JSON."
            ),
        }

        conv_with_system = [system_message] + messages
        attempts = 20

        while attempts > 0:
            data = {
                "model": self.model,
                "messages": conv_with_system,
                "stream": False,
                "format": schema_json,
            }
            if context_size > 0:
                data["options"] = {"num_ctx": context_size}

            try:
                response = requests.post(f"{self.base_url}/api/chat", json=data, headers=self.headers)
                response.raise_for_status()
                content = response.json()["message"]["content"]
                parsed_content = self._find_json(content)
                validate(instance=parsed_content, schema=schema_json)
                return schema(**parsed_content)

            except requests.exceptions.RequestException as e:
                raise DazLlmError(f"Ollama API error: {e}") from e
            except json.JSONDecodeError:
                conv_with_system.append(
                    {
                        "role": "system",
                        "content": (
                            "The previous response was not valid JSON. "
                            "Please ensure the output is valid JSON strictly "
                            "following the schema."
                        ),
                    }
                )
            except ValidationError as e:
                conv_with_system.append(
                    {
                        "role": "system",
                        "content": (
                            f"Your previous output did not adhere to the JSON schema "
                            f"because: {e}. Please generate a response that strictly "
                            "follows the schema without any extra text or formatting."
                        ),
                    }
                )
            except KeyError as e:
                raise DazLlmError(f"Unexpected Ollama response structure: {e}") from e

            attempts -= 1

        raise DazLlmError("Failed to get valid structured response after multiple attempts")

    async def chat_async(self, conversation: Conversation, force_json: bool = False) -> LlmResponse:
        import asyncio

        messages = self._normalize_conversation(conversation)
        data = {"model": self.model, "messages": messages, "stream": False}
        if force_json:
            data["format"] = "json"

        def _call_sync():
            try:
                response = requests.post(f"{self.base_url}/api/chat", json=data, headers=self.headers)
                response.raise_for_status()
                content = response.json()["message"]["content"]
                return LlmResponse(value=content, output="", provider=self.provider_name)
            except requests.exceptions.RequestException as e:  # noqa: BLE001
                raise DazLlmError(f"Ollama API error: {e}") from e
            except KeyError as e:  # noqa: BLE001
                raise DazLlmError(f"Unexpected Ollama response structure: {e}") from e

        return await asyncio.to_thread(_call_sync)

    async def chat_structured_async(
        self, conversation: Conversation, schema: Type[BaseModel], context_size: int = 0
    ) -> LlmResponse:
        """Chat with structured output using Pydantic schema"""
        import asyncio

        def _call_sync():
            messages = self._normalize_conversation(conversation)
            schema_json = schema.model_json_schema()

            if self._supports_format:
                # Use the standard format parameter approach
                result = self._chat_structured_with_format(messages, schema, schema_json, context_size)
            else:
                # Use manual schema injection for models without format support
                result = self._chat_structured_without_format(messages, schema, schema_json, context_size)

            return LlmResponse(value=result, output="", provider=self.provider_name)

        return await asyncio.to_thread(_call_sync)

    def _chat_structured_without_format(
        self,
        messages: list,
        schema: Type[BaseModel],
        schema_json: dict,
        context_size: int,
    ) -> BaseModel:
        """Structured output for models without format support"""
        # Inject schema instructions directly into the prompt
        schema_str = json.dumps(schema_json, indent=2)

        system_message = {
            "role": "system",
            "content": (
                "You must respond with valid JSON that matches the following schema exactly. "
                "Output ONLY the JSON, wrapped in ```json code blocks. "
                "Do not include any other text, explanations, or the schema itself.\n\n"
                f"Required JSON Schema:\n{schema_str}"
            ),
        }

        # Add schema reminder to the last user message
        modified_messages = messages.copy()
        if modified_messages and modified_messages[-1]["role"] == "user":
            modified_messages[-1] = {
                "role": "user",
                "content": (
                    f"{modified_messages[-1]['content']}\n\n"
                    f"Remember: Return your response as valid JSON matching this schema:\n```json\n{schema_str}\n```"
                ),
            }

        conv_with_system = [system_message] + modified_messages
        attempts = 20

        while attempts > 0:
            data = {
                "model": self.model,
                "messages": conv_with_system,
                "stream": False,
            }
            if context_size > 0:
                data["options"] = {"num_ctx": context_size}

            try:
                response = requests.post(f"{self.base_url}/api/chat", json=data, headers=self.headers)
                response.raise_for_status()
                content = response.json()["message"]["content"]

                # Try multiple parsing strategies
                parsed_content = self._parse_json_with_multiple_strategies(content)
                validate(instance=parsed_content, schema=schema_json)
                return schema(**parsed_content)

            except requests.exceptions.RequestException as e:
                raise DazLlmError(f"Ollama API error: {e}") from e
            except (json.JSONDecodeError, ValueError) as e:
                conv_with_system.append(
                    {
                        "role": "system",
                        "content": (
                            f"Invalid JSON in previous response. Error: {e}. "
                            "Please output valid JSON wrapped in ```json blocks."
                        ),
                    }
                )
            except ValidationError as e:
                conv_with_system.append(
                    {
                        "role": "system",
                        "content": (
                            f"JSON doesn't match schema. Error: {e}. Please fix and "
                            "return valid JSON matching the schema."
                        ),
                    }
                )
            except KeyError as e:
                raise DazLlmError(f"Unexpected Ollama response structure: {e}") from e

            attempts -= 1

        raise DazLlmError("Failed to get valid structured response after multiple attempts")

    def _find_json(self, text: str) -> dict:
        """Extract JSON from text response"""
        if "```json" in text:
            start = text.index("```json") + len("```json")
            end = text.index("```", start)
            json_text = text[start:end].strip()
        else:
            json_text = text
        return json.loads(json_text)

    def _parse_json_with_multiple_strategies(self, text: str) -> dict:
        """Parse JSON with multiple strategies"""
        # Strategy 1: Try to find JSON in markdown code blocks
        if "```json" in text:
            try:
                return self._find_json(text)
            except (json.JSONDecodeError, ValueError):
                pass

        # Strategy 2: Try to parse the entire text as JSON
        try:
            return json.loads(text.strip())
        except json.JSONDecodeError:
            pass

        # Strategy 3: Look for JSON-like structures with regex
        json_patterns = [
            r"\{[^{}]*\}",  # Simple single-level JSON
            r"\{.*\}",  # Any JSON object
            r"\[.*\]",  # JSON array
        ]

        for pattern in json_patterns:
            matches = re.findall(pattern, text, re.DOTALL)
            for match in matches:
                try:
                    parsed = json.loads(match)
                    if isinstance(parsed, dict):
                        return parsed
                except json.JSONDecodeError:
                    continue

        # Strategy 4: Clean up common issues and retry
        cleaned = text.strip()
        # Remove common prefixes/suffixes
        prefixes = ["Here is the JSON:", "JSON:", "Output:", "Result:"]
        for prefix in prefixes:
            if cleaned.startswith(prefix):
                cleaned = cleaned[len(prefix) :].strip()

        # Remove trailing punctuation
        cleaned = cleaned.rstrip(".,;:")

        # Try parsing cleaned text
        try:
            return json.loads(cleaned)
        except json.JSONDecodeError:
            pass

        # If all strategies fail, raise an error with helpful context
        raise ValueError(f"Could not parse JSON from response. Text: {text[:500]}...")

    def image(self, prompt: str, file_name: str, width: int = 1024, height: int = 1024) -> str:
        """Generate image using Ollama (not supported by default)"""
        raise DazLlmError(
            ("Image generation not supported by Ollama. Use OpenAI or other providers for image generation.")
        )

    async def async_image(self, prompt: str, file_name: str, width: int = 1024, height: int = 1024) -> str:
        """Async image generation using Ollama (not supported by default)"""
        raise DazLlmError(
            ("Image generation not supported by Ollama. Use OpenAI or other providers for image generation.")
        )

    def get_context_length(self) -> int:
        """Get the context length for the current Ollama model"""
        try:
            # Try to get model info from Ollama API
            response = requests.get(
                f"{self.base_url}/api/show",
                json={"name": self.model},
                headers=self.headers,
            )

            if response.status_code == 200:
                model_info = response.json()
                # Check if context length is in the model info
                if "parameters" in model_info:
                    params = model_info["parameters"]
                    # Look for common parameter names for context length
                    for key in [
                        "num_ctx",
                        "context_length",
                        "max_context_length",
                        "context_size",
                    ]:
                        if key in params:
                            try:
                                return int(params[key])
                            except (ValueError, TypeError):
                                continue

                # Check modelfile for context length settings
                if "modelfile" in model_info:
                    modelfile = model_info["modelfile"]
                    # Look for PARAMETER num_ctx lines
                    import re

                    match = re.search(r"PARAMETER\s+num_ctx\s+(\d+)", modelfile, re.IGNORECASE)
                    if match:
                        return int(match.group(1))

        except Exception:
            pass  # Fall back to defaults

        # Fallback to known context lengths for common models
        context_lengths = {
            # Common Ollama models and their typical context lengths
            "mistral-small": 32768,  # Mistral models typically 32K
            "mistral": 32768,
            "mistral-large": 32768,
            "llama3": 8192,  # Llama 3 models
            "llama3:8b": 8192,
            "llama3:70b": 8192,
            "llama2": 4096,  # Llama 2 models
            "llama2:7b": 4096,
            "llama2:13b": 4096,
            "llama2:70b": 4096,
            "phi3": 4096,  # Phi-3 models
            "phi3:3.8b": 4096,
            "phi3:14b": 4096,
            "qwen3": 32768,  # Qwen models
            "qwen3:32b": 32768,
            "codestral": 32768,  # Code models
            "codeqwen": 32768,
            "gpt-oss:20b": 4096,  # GPT-OSS model
        }

        # Try exact match first
        if self.model in context_lengths:
            return context_lengths[self.model]

        # Try pattern matching for model families
        model_lower = self.model.lower()
        if "mistral" in model_lower:
            return 32768
        elif "llama3" in model_lower:
            return 8192
        elif "llama2" in model_lower:
            return 4096
        elif "phi3" in model_lower:
            return 4096
        elif "qwen" in model_lower:
            return 32768
        elif "code" in model_lower:
            return 32768

        # Conservative default for unknown models
        return 4096


import unittest  # noqa: E402


class TestLlmOllama(unittest.TestCase):
    def test_default_model(self):
        self.assertEqual(LlmOllama.default_model(), "mistral-small")
