"""
Local provider implementation for InferenceKit.
Supports local models and inference engines.
"""

import asyncio
import os
import subprocess
from typing import Dict, Any, Optional, List
from datetime import datetime
from uuid import UUID, uuid4
from pathlib import Path

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


class LocalProvider:
    """Local provider implementation for running models locally."""

    def __init__(self, config: ProviderConfig):
        self.config = config
        self.model_path = config.base_url or "."
        self.engine = config.api_key or "auto"
        self.timeout = config.timeout
        self.max_retries = config.max_retries
        self.available_engines = self._detect_available_engines()

    def _detect_available_engines(self) -> List[str]:
        """Detect available local inference engines."""
        engines = []

        # Check for Ollama
        if self._check_command("ollama"):
            engines.append("ollama")

        # Check for LM Studio
        if self._check_command("lm-studio"):
            engines.append("lm-studio")

        # Check for other engines...

        return engines

    def _check_command(self, command: str) -> bool:
        """Check if a command is available."""
        try:
            subprocess.run([command, "--version"],
                         capture_output=True,
                         timeout=5,
                         check=True)
            return True
        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
            return False

    async def initialize(self) -> None:
        """Initialize the local provider."""
        if not self.available_engines:
            raise ConnectionError("No local inference engines found")

    async def shutdown(self) -> None:
        """Shutdown the local provider."""
        pass

    async def make_request(self, request: Request) -> Response:
        """Make a request to a local model."""
        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 a local model."""
        if self.engine == "ollama":
            return await self._ollama_generate(request)
        elif self.engine == "lm-studio":
            return await self._lm_studio_generate(request)
        else:
            return await self._generic_generate(request)

    async def _ollama_generate(self, request: Request) -> Response:
        """Generate text using Ollama."""
        try:
            # Check if model is available
            model_available = await self._check_ollama_model(request.model)
            if not model_available:
                raise ModelNotFoundError(f"Model {request.model} not found in Ollama")

            # Build command
            cmd = ["ollama", "run", request.model]

            # Add parameters
            params = []
            if request.parameters.get("temperature") is not None:
                params.append(f"-temperature {request.parameters['temperature']}")
            if request.parameters.get("max_tokens") is not None:
                params.append(f"-max-tokens {request.parameters['max_tokens']}")
            if request.parameters.get("top_p") is not None:
                params.append(f"-top-p {request.parameters['top_p']}")
            if request.parameters.get("stop"):
                params.append(f"-stop {request.parameters['stop']}")

            if params:
                cmd.extend(["--" + param for param in params])

            # Run command
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )

            # Send prompt
            stdout, stderr = await process.communicate(input=request.prompt.encode())

            if process.returncode != 0:
                raise ProviderError(f"Ollama failed: {stderr.decode()}")

            return Response(
                request_id=request.id,
                content=stdout.decode().strip(),
                metadata={
                    "model": request.model,
                    "engine": "ollama",
                    "created_at": datetime.utcnow().isoformat()
                }
            )

        except Exception as e:
            raise ProviderError(f"Failed to generate with Ollama: {e}")

    async def _lm_studio_generate(self, request: Request) -> Response:
        """Generate text using LM Studio."""
        # Implementation for LM Studio would go here
        raise NotImplementedError("LM Studio integration not yet implemented")

    async def _generic_generate(self, request: Request) -> Response:
        """Generate text using a generic approach."""
        # This could be a custom model or other inference engine
        raise NotImplementedError("Generic local model generation not implemented")

    async def _chat(self, request: Request) -> Response:
        """Chat using a local model."""
        if self.engine == "ollama":
            return await self._ollama_chat(request)
        elif self.engine == "lm-studio":
            return await self._lm_studio_chat(request)
        else:
            return await self._generic_chat(request)

    async def _ollama_chat(self, request: Request) -> Response:
        """Chat using Ollama."""
        try:
            # Check if model is available
            model_available = await self._check_ollama_model(request.model)
            if not model_available:
                raise ModelNotFoundError(f"Model {request.model} not found in Ollama")

            # Build command
            cmd = ["ollama", "run", request.model]

            # Add parameters
            params = []
            if request.parameters.get("temperature") is not None:
                params.append(f"-temperature {request.parameters['temperature']}")
            if request.parameters.get("max_tokens") is not None:
                params.append(f"-max-tokens {request.parameters['max_tokens']}")
            if request.parameters.get("top_p") is not None:
                params.append(f"-top-p {request.parameters['top_p']}")
            if request.parameters.get("stop"):
                params.append(f"-stop {request.parameters['stop']}")

            if params:
                cmd.extend(["--" + param for param in params])

            # Build chat messages
            messages = request.parameters.get("messages", [{
                "role": "user",
                "content": request.prompt
            }])

            chat_prompt = "\n".join([
                f"{msg.get('role', 'user')}: {msg.get('content', '')}"
                for msg in messages
            ])

            # Run command
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )

            # Send prompt
            stdout, stderr = await process.communicate(input=chat_prompt.encode())

            if process.returncode != 0:
                raise ProviderError(f"Ollama failed: {stderr.decode()}")

            return Response(
                request_id=request.id,
                content=stdout.decode().strip(),
                metadata={
                    "model": request.model,
                    "engine": "ollama",
                    "created_at": datetime.utcnow().isoformat()
                }
            )

        except Exception as e:
            raise ProviderError(f"Failed to chat with Ollama: {e}")

    async def _lm_studio_chat(self, request: Request) -> Response:
        """Chat using LM Studio."""
        # Implementation for LM Studio would go here
        raise NotImplementedError("LM Studio integration not yet implemented")

    async def _generic_chat(self, request: Request) -> Response:
        """Chat using a generic approach."""
        # This could be a custom model or other inference engine
        raise NotImplementedError("Generic local model chat not implemented")

    async def _generate_embedding(self, request: Request) -> Response:
        """Generate embeddings using a local model."""
        # Local embedding generation would need specific implementation
        raise NotImplementedError("Local embeddings not implemented")

    async def _check_ollama_model(self, model_name: str) -> bool:
        """Check if a model is available in Ollama."""
        try:
            process = await asyncio.create_subprocess_exec(
                "ollama", "list",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await process.communicate()

            if process.returncode != 0:
                return False

            available_models = stdout.decode().splitlines()
            return model_name in available_models
        except Exception:
            return False

    async def get_model_info(self) -> Dict[str, Any]:
        """Get information about available local models."""
        if self.engine == "ollama":
            return await self._get_ollama_model_info()
        else:
            return {"error": "Model info not available for this engine"}

    async def _get_ollama_model_info(self) -> Dict[str, Any]:
        """Get information about Ollama models."""
        try:
            process = await asyncio.create_subprocess_exec(
                "ollama", "show",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await process.communicate()

            if process.returncode != 0:
                return {"error": "Failed to get model info"}

            return {
                "engine": "ollama",
                "models": stdout.decode().strip(),
                "available": True
            }
        except Exception as e:
            return {"error": f"Failed to get model info: {e}"}

    async def list_models(self) -> List[str]:
        """List available local models."""
        if self.engine == "ollama":
            return await self._list_ollama_models()
        else:
            return []

    async def _list_ollama_models(self) -> List[str]:
        """List available Ollama models."""
        try:
            process = await asyncio.create_subprocess_exec(
                "ollama", "list",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await process.communicate()

            if process.returncode != 0:
                return []

            return [line.strip() for line in stdout.decode().splitlines() if line.strip()]
        except Exception:
            return []

    async def check_health(self) -> bool:
        """Check if the local provider is healthy."""
        if self.engine == "ollama":
            return await self._check_ollama_health()
        else:
            return False

    async def _check_ollama_health(self) -> bool:
        """Check if Ollama is running."""
        try:
            process = await asyncio.create_subprocess_exec(
                "ollama", "list",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                timeout=5
            )
            stdout, stderr = await process.communicate()
            return process.returncode == 0
        except Exception:
            return False

    async def get_system_info(self) -> Dict[str, Any]:
        """Get system information for local inference."""
        return {
            "engine": self.engine,
            "model_path": str(self.model_path),
            "available_engines": self.available_engines,
            "system_resources": {
                "cpu_count": os.cpu_count(),
                "memory": self._get_memory_info()
            }
        }

    def _get_memory_info(self) -> Dict[str, Any]:
        """Get memory information."""
        try:
            import psutil
            memory = psutil.virtual_memory()
            return {
                "total": memory.total,
                "available": memory.available,
                "percent": memory.percent
            }
        except ImportError:
            return {"error": "psutil not available"}