"""
Ollama provider implementation for InferenceKit.
Supports models hosted on Ollama servers.
"""

import asyncio
import aiohttp
from typing import Dict, Any, Optional, List
from datetime import datetime
from uuid import UUID, uuid4

from ..types import (
    ModelType,
    ProviderType,
    Request,
    Response,
    ProviderConfig,
    ModelConfig,
)
from ..exceptions import (
    ProviderError,
    AuthenticationError,
    RateLimitError,
    ModelNotFoundError,
    ConnectionError,
    OperationTimeoutError,
)


class OllamaProvider:
    """Ollama provider implementation for remote Ollama servers."""

    def __init__(self, config: ProviderConfig):
        self.config = config
        self.base_url = config.base_url or "http://localhost:11434"
        self.session: Optional[aiohttp.ClientSession] = None
        self.timeout = config.timeout
        self.max_retries = config.max_retries
        self.headers = {
            "Content-Type": "application/json",
        }

    async def initialize(self) -> None:
        """Initialize the provider."""
        if self.session is None:
            self.session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=self.timeout)
            )

    async def shutdown(self) -> None:
        """Shutdown the provider."""
        if self.session:
            await self.session.close()
            self.session = None

    async def _make_request(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """Make a request to the Ollama API."""
        url = f"{self.base_url}/{endpoint}"

        for attempt in range(self.max_retries):
            try:
                async with self.session.post(
                    url,
                    json=data,
                    headers=self.headers
                ) as response:
                    if response.status == 401:
                        raise AuthenticationError("Invalid API key")
                    if response.status == 429:
                        raise RateLimitError("Rate limit exceeded")
                    if response.status == 404:
                        raise ModelNotFoundError(f"Model not found: {data.get('model')}")
                    if response.status >= 500:
                        raise ProviderError(f"Server error: {response.status}")

                    return await response.json()

            except (aiohttp.ClientConnectionError, asyncio.TimeoutError) as e:
                if attempt == self.max_retries - 1:
                    raise ConnectionError(f"Connection failed: {e}")
                await asyncio.sleep(2 ** attempt)

            except aiohttp.ClientResponseError as e:
                if e.status == 401:
                    raise AuthenticationError("Invalid API key")
                elif e.status == 429:
                    raise RateLimitError("Rate limit exceeded")
                elif e.status == 404:
                    raise ModelNotFoundError(f"Model not found: {data.get('model')}")
                else:
                    raise ProviderError(f"API error: {e}")

    async def make_request(self, request: Request) -> Response:
        """Make a request to the Ollama API."""
        if request.type == ModelType.TEXT_GENERATION:
            return await self._generate_text(request)
        elif request.type == ModelType.CHAT:
            return await self._chat(request)
        elif request.type == ModelType.EMBEDDING:
            return await self._generate_embedding(request)
        else:
            raise ValueError(f"Unsupported model type: {request.type}")

    async def _generate_text(self, request: Request) -> Response:
        """Generate text using Ollama API."""
        data = {
            "model": request.model,
            "prompt": request.prompt,
            "options": {
                "temperature": request.parameters.get("temperature", 0.7),
                "top_p": request.parameters.get("top_p", 1.0),
                "frequency_penalty": request.parameters.get("frequency_penalty", 0.0),
                "presence_penalty": request.parameters.get("presence_penalty", 0.0),
                "stop": request.parameters.get("stop"),
                "max_tokens": request.parameters.get("max_tokens"),
                "stream": request.parameters.get("stream", False),
            }
        }

        result = await self._make_request("generate", data)

        return Response(
            request_id=request.id,
            content=result.get("response", ""),
            metadata={
                "model": result.get("model"),
                "usage": result.get("usage", {}),
                "created_at": datetime.utcnow().isoformat()
            }
        )

    async def _chat(self, request: Request) -> Response:
        """Chat using Ollama API."""
        data = {
            "model": request.model,
            "messages": request.parameters.get("messages", [{
                "role": "user",
                "content": request.prompt
            }]),
            "options": {
                "temperature": request.parameters.get("temperature", 0.7),
                "top_p": request.parameters.get("top_p", 1.0),
                "frequency_penalty": request.parameters.get("frequency_penalty", 0.0),
                "presence_penalty": request.parameters.get("presence_penalty", 0.0),
                "stop": request.parameters.get("stop"),
                "max_tokens": request.parameters.get("max_tokens"),
                "stream": request.parameters.get("stream", False),
            }
        }

        result = await self._make_request("chat", data)

        return Response(
            request_id=request.id,
            content=result.get("message", {}).get("content", ""),
            metadata={
                "model": result.get("model"),
                "usage": result.get("usage", {}),
                "created_at": datetime.utcnow().isoformat()
            }
        )

    async def _generate_embedding(self, request: Request) -> Response:
        """Generate embeddings using Ollama API."""
        # Ollama doesn't support embeddings directly
        # This would need to use a different approach
        raise NotImplementedError("Embeddings are not supported by Ollama provider")

    async def get_model_info(self) -> Dict[str, Any]:
        """Get information about available models."""
        try:
            result = await self._make_request("models", {})
            return result
        except Exception:
            return {"error": "Failed to retrieve model information"}

    async def list_models(self) -> List[str]:
        """List available models."""
        try:
            result = await self._make_request("models", {})
            return [model["name"] for model in result.get("models", [])]
        except Exception:
            return []

    async def check_health(self) -> bool:
        """Check if the provider is healthy."""
        try:
            await self._make_request("health", {})
            return True
        except Exception:
            return False

    async def get_system_info(self) -> Dict[str, Any]:
        """Get system information for Ollama server."""
        try:
            result = await self._make_request("system", {})
            return result
        except Exception:
            return {"error": "Failed to get system info"}

    async def pull_model(self, model_name: str) -> bool:
        """Pull a model from the Ollama registry."""
        try:
            data = {"name": model_name}
            result = await self._make_request("pull", data)
            return result.get("status") == "success"
        except Exception:
            return False

    async def remove_model(self, model_name: str) -> bool:
        """Remove a model from the server."""
        try:
            data = {"name": model_name}
            result = await self._make_request("remove", data)
            return result.get("status") == "success"
        except Exception:
            return False

    async def list_running_models(self) -> List[str]:
        """List models currently running on the server."""
        try:
            result = await self._make_request("list", {})
            return [model["name"] for model in result.get("models", [])]
        except Exception:
            return []

    async def get_model_details(self, model_name: str) -> Dict[str, Any]:
        """Get details about a specific model."""
        try:
            data = {"name": model_name}
            result = await self._make_request("show", data)
            return result
        except Exception:
            return {"error": "Failed to get model details"}

    async def create_collection(self, collection_name: str, models: List[str]) -> bool:
        """Create a collection of models."""
        try:
            data = {
                "name": collection_name,
                "models": models
            }
            result = await self._make_request("collections/create", data)
            return result.get("status") == "success"
        except Exception:
            return False

    async def list_collections(self) -> List[str]:
        """List available collections."""
        try:
            result = await self._make_request("collections/list", {})
            return result.get("collections", [])
        except Exception:
            return []