import logging
from typing import Optional, Dict, Any
from ollama import Client

class OllamaClient:
    """Клиент для работы с Ollama API."""

    def __init__(self, base_url: str):
        """
        Инициализация клиента.

        Args:
            base_url: Базовый URL для подключения к Ollama серверу
        """
        self.base_url = base_url
        self.client = Client(host=base_url)
        self.logger = logging.getLogger(__name__)

    def generate(self, model: str, prompt: str, num_ctx: int = 32000, **kwargs) -> str:
        """
        Генерация ответа от модели.

        Args:
            model: Название модели
            prompt: Входной промпт
            num_ctx: Размер контекста (по умолчанию 32000)
            **kwargs: Дополнительные параметры для запроса

        Returns:
            Сгенерированный ответ

        Raises:
            Exception: Если произошла ошибка при генерации
        """
        try:
            self.logger.info(f"Generating response for model {model}")
            self.logger.debug(f"Prompt: {prompt[:200]}...")  # Логируем начало промпта
            # Объединяем options из kwargs с num_ctx
            options = {'num_ctx': num_ctx}
            if 'options' in kwargs:
                options.update(kwargs.pop('options'))
            response = self.client.generate(
                model=model,
                prompt=prompt,
                options=options,
                **kwargs
            )
            # Проверяем структуру ответа
            self.logger.debug(f"Response structure: {response}")
            self.logger.debug(f"Response type: {type(response)}")
            
            # Если это объект GenerateResponse (как в ollama 0.3+)
            if hasattr(response, 'response'):
                return response.response
            elif hasattr(response, 'text'):
                return response.text
            elif isinstance(response, dict):
                if 'response' in response:
                    return response['response']
                elif 'text' in response:
                    return response['text']
                else:
                    # Попробуем извлечь любое строковое значение
                    for key, value in response.items():
                        if isinstance(value, str):
                            return value
                    raise ValueError(f"Unexpected response format - no text or response field: {response}")
            elif isinstance(response, str):
                return response
            else:
                raise ValueError(f"Unexpected response format: {response}")
        except Exception as e:
            error_msg = f"Error generating response for model {model}: {e}"
            self.logger.error(error_msg)
            self.logger.debug(f"Failed prompt: {prompt[:500]}")
            raise

    def chat(self, model: str, messages: list, num_ctx: int = 32000, **kwargs) -> str:
        """
        Диалог с моделью.

        Args:
            model: Название модели
            messages: Список сообщений в формате [{'role': 'user', 'content': '...'}, ...]
            num_ctx: Размер контекста (по умолчанию 32000)
            **kwargs: Дополнительные параметры для запроса

        Returns:
            Ответ от модели

        Raises:
            Exception: Если произошла ошибка при чате
        """
        try:
            self.logger.info(f"Chatting with model {model}")
            # Объединяем options из kwargs с num_ctx
            options = {'num_ctx': num_ctx}
            if 'options' in kwargs:
                options.update(kwargs.pop('options'))
            response = self.client.chat(
                model=model,
                messages=messages,
                options=options,
                **kwargs
            )
            return response['message']['content']
        except Exception as e:
            self.logger.error(f"Error in chat: {e}")
            raise

    def list_models(self) -> list:
        """
        Получение списка доступных моделей.

        Returns:
            Список моделей
        """
        try:
            response = self.client.list()
            return [model['name'] for model in response['models']]
        except Exception as e:
            self.logger.error(f"Error listing models: {e}")
            raise
