"""
Этот модуль предоставляет инструменты для взаимодействия с различными LLM,
включая OpenAI (Gemini) и MistralAI, обеспечивая унифицированный интерфейс.
"""

from typing import Dict, Any, List, Optional
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage

from openai import OpenAI

from mistralai import Mistral

# --- Конфигурация LLM ---
# Пустой словарь для локальных моделей (Ollama)
LOCAL_MODELS: Dict[str, Any] = {}

DEFAULT_TEMPERATURE = 1.0

# Конфигурация внешних моделей
MODELS: Dict[str, Dict[str, Any]] = {
    "gemini-2.0-flash-r": {
        "name": "gemini-2.0-flash-r",
        "provider": "openai",  # Изменено на "openai"
        "model_name": "gemini-2.0-flash",  # Добавлено имя модели для LangChain
        "apiBase":
        "https://render-service-gsu7.onrender.com/g/v1beta",  # "generativelanguage.googleapis.com", # "https://render-service-gsu7.onrender.com/g/v1beta",
        "apiKey": "AIzaSyDpueKFWVqknVKlQn6TdasLmJ2lvAUiBik",
        "stream": True,
        "capabilities": ["vision"],
    },
    "gemini-2.5-flash-lite-openrouter": {
        "name": "google/gemini-2.5-flash-lite",
        "provider": "openai",
        "model_name": "google/gemini-2.5-flash-lite",
        "apiBase": "https://openrouter.ai/api/v1",  # Добавлено /api/v1
        "apiKey":
        "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",
        "stream": True,
        "capabilities": ["vision", "reasoning"],
        "model_kwargs": {
            "include_reasoning": True
        }
    },
    "gemini-2.5-flash-r": {
        "name": "gemini-2.5-flash-r",
        "provider": "openai",  # Изменено на "openai"
        "model_name": "gemini-2.5-flash",  # Добавлено имя модели для LangChain
        "apiBase":
        "https://render-service-gsu7.onrender.com/g/v1beta",  # "generativelanguage.googleapis.com", # "https://render-service-gsu7.onrender.com/g/v1beta",
        "apiKey": "AIzaSyDpueKFWVqknVKlQn6TdasLmJ2lvAUiBik",
        "stream": True,
        "capabilities": ["vision"],
    },
    "gemini-2.5-flash-openrouter": {
        "name": "google/gemini-2.5-flash",
        "provider": "openai",
        "model_name": "google/gemini-2.5-flash",
        "apiBase": "https://openrouter.ai/api/v1",  # Добавлено /api/v1
        "apiKey":
        "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",
        "stream": True,
        "capabilities": ["vision", "reasoning"],
        "model_kwargs": {
            "include_reasoning": True
        }
    },
    "gemini-3.0-flash-openrouter": {
        "name": "google/gemini-3-flash-preview",
        "provider": "openai",
        "model_name": "google/gemini-3-flash-preview",
        "apiBase": "https://openrouter.ai/api/v1",  # Добавлено /api/v1
        "apiKey":
        "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",
        "stream": True,
        "capabilities": ["vision", "reasoning"],
        "model_kwargs": {
            "include_reasoning": True
        }
    },
    "gemini-3.1-pro-openrouter": {
        "name": "google/gemini-3.1-pro-preview",
        "provider": "openai",
        "model_name": "google/gemini-3.1-pro-preview",
        "apiBase": "https://openrouter.ai/api/v1",  # Добавлено /api/v1
        "apiKey":
        "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",
        "stream": True,
        "capabilities": ["vision", "reasoning"],
        "model_kwargs": {
            "include_reasoning": True
        }
    },
    "claude-4.5-sonnet-bh": {
        "name": "claude-4.5-sonnet-bh",
        "provider": "openai",  # Изменено на "openai"
        "model_name":
        "claude-sonnet-4.5",  # Добавлено имя модели для LangChain
        "apiBase": "https://bothub.chat/api/v2/openai/v1",
        "apiKey":
        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpZCI6ImQ3MTMwYTk0LTJiNmMtNDgzZi04Y2NiLWM5MjgwNzlmNDFlZSIsImlzRGV2ZWxvcGVyIjp0cnVlLCJpYXQiOjE3NTQ0NzYwMDcsImV4cCI6MjA3MDA1MjAwN30._64_DXNZatdd4-4VbXp8dOXdDhIlNWMQt9A_Nn6mu3g",
        "stream": True,
        "capabilities": ["vision"],
    },
    "mistral-small-latest": {
        "name": "mistral-small-latest",
        "provider": "mistralai",
        "model_name":
        "mistral-small-latest",  # Добавлено имя модели для LangChain
        "apiBase": "https://api.mistral.ai",
        "apiKey": "Q0m29fvxBY0Cfdj4sjHaKqccy1NjonLW",
        "stream": True,
        "capabilities": ["vision"],
    },
    "mistral-small-latest-r": {
        "name": "mistral-small-latest",
        "provider": "mistralai",
        "model_name":
        "mistral-small-latest",  # Добавлено имя модели для LangChain
        "apiBase": "https://render-service-gsu7.onrender.com/m",
        "apiKey": "Q0m29fvxBY0Cfdj4sjHaKqccy1NjonLW",
        "stream": True,
        "capabilities": ["vision"],
    },
    "mistral-small-latest-error": {
        "name": "mistral-small-latest",
        "provider": "mistralai",
        "model_name":
        "mistral-small-latest-error",  # Добавлено имя модели для LangChain
        "apiBase": "https://render-service-gsu7.onrender.com/m",
        "apiKey": "Q0m29fvxBY0Cfdj4sjHaKqccy1NjonLW",
        "stream": True,
        "capabilities": ["vision"],
    },
    "claude-sonnet-4.6-openrouter": {
        "name": "anthropic/claude-sonnet-4.6",
        "provider": "openai",
        "model_name": "anthropic/claude-sonnet-4.6",
        "apiBase": "https://openrouter.ai/api/v1",  # Добавлено /api/v1
        "apiKey":
        "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",
        "stream": True,
        "capabilities": ["vision", "reasoning"],
        "model_kwargs": {
            "include_reasoning": True
        }
    },
    "deepseek-v4-flash-openrouter": {
        "name": "deepseek/deepseek-v4-flash",
        "provider": "openai",
        "model_name": "deepseek/deepseek-v4-flash",
        "apiBase": "https://openrouter.ai/api/v1",  # Добавлено /api/v1
        "apiKey":
        "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",
        "stream": True,
        "capabilities": ["vision", "reasoning"],
        "model_kwargs": {
            "include_reasoning": True
        }
    }
}


class CustomLLM:
    """
    Класс-обертка для взаимодействия с различными LLM.
    """

    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.name = config["name"]
        self.provider = config["provider"]

        try:

            if self.provider == "openai":
                # Инициализация Gemini через OpenAI API
                self.client = OpenAI(api_key=config["apiKey"],
                                     base_url=config["apiBase"])
                self.model_name = config["model_name"]
                self.stream = config["stream"]

            elif self.provider == "mistralai":
                # Для Mistral используем прямые HTTP-запросы
                self.client = Mistral(api_key=config["apiKey"],
                                      server_url=config["apiBase"])
                self.model_name = config["model_name"]
                self.stream = config["stream"]
                # self.capabilities = config.get("capabilities", [])
            else:
                raise ValueError(f"Неизвестный провайдер LLM: {self.provider}")

        except Exception as e:
            print(f"Ошибка при вызове OpenAI API: {e}")
            raise

    def invoke(self, messages: List[Any], temperature: float = DEFAULT_TEMPERATURE, max_tokens: Optional[int] = None) -> str:
        """
        Отправляет запрос к LLM и возвращает ответ.
        Messages могут быть строкой или списком объектов Langchain Message.
        """
        if self.provider == "openai":
            # OpenAI ожидает список сообщений в определенном формате
            openai_messages = []
            if isinstance(messages, str):
                openai_messages.append({"role": "user", "content": messages})
            elif isinstance(messages, list):
                for msg in messages:
                    if isinstance(msg, HumanMessage):
                        openai_messages.append({
                            "role": "user",
                            "content": msg.content
                        })
                    elif isinstance(msg, SystemMessage):
                        openai_messages.append({
                            "role": "system",
                            "content": msg.content
                        })
                    elif isinstance(
                            msg, AIMessage):  # Добавляем обработку AIMessage
                        openai_messages.append({
                            "role": "assistant",
                            "content": msg.content
                        })
                    else:  # Предполагаем, что это другие типы сообщений Langchain или словари
                        openai_messages.append({
                            "role":
                            msg.type if hasattr(msg, 'type') else 'user',
                            "content":
                            msg.content
                        })

            try:
                kwargs = {"temperature": temperature}
                if max_tokens:
                    kwargs["max_tokens"] = int(max_tokens)
                    
                response = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=openai_messages,
                    stream=self.stream,
                    **kwargs
                )

                if self.stream:
                    # Handle streaming responses
                    collected_chunks = []
                    collected_messages = []
                    for chunk in response:
                        collected_chunks.append(chunk)
                        chunk_message = chunk.choices[0].delta.content
                        if chunk_message is not None:
                            collected_messages.append(chunk_message)
                    full_reply_content = ''.join(collected_messages)
                    return full_reply_content
                else:
                    return response.choices[0].message.content

            except Exception as e:
                print(f"Ошибка при вызове OpenAI API: {e}")
                raise
        elif self.provider == "mistralai":
            mistralai_messages = []
            if isinstance(messages, str):
                mistralai_messages.append({
                    "role": "user",
                    "content": messages
                })
            elif isinstance(messages, list):
                for msg in messages:
                    if isinstance(msg, HumanMessage):
                        mistralai_messages.append({
                            "role": "user",
                            "content": msg.content
                        })
                    elif isinstance(msg, SystemMessage):
                        mistralai_messages.append({
                            "role": "system",
                            "content": msg.content
                        })
                    elif isinstance(
                            msg, AIMessage):  # Добавляем обработку AIMessage
                        mistralai_messages.append({
                            "role": "assistant",
                            "content": msg.content
                        })
                    else:  # Предполагаем, что это другие типы сообщений Langchain или словари
                        mistralai_messages.append({
                            "role":
                            msg.type if hasattr(msg, 'type') else 'user',
                            "content":
                            msg.content
                        })

            try:
                kwargs = {"temperature": temperature}
                if max_tokens:
                    kwargs["max_tokens"] = int(max_tokens)
                    
                response = self.client.chat.stream(model=self.model_name,
                                                   messages=mistralai_messages,
                                                   **kwargs)

                if self.stream:
                    collected_messages = []
                    for chunk in response:
                        delta = chunk.data.choices[0].delta
                        chunk_message = getattr(
                            delta, "content",
                            None)  # или delta.content, если всегда есть

                        if chunk_message is not None:
                            collected_messages.append(chunk_message)

                        if getattr(delta, "finish_reason", None) is not None:
                            break

                    full_reply_content = "".join(collected_messages)
                    return full_reply_content
                else:
                    return response.choices[0].message.content

            except Exception as e:
                print(f"Ошибка при вызове OpenAI API: {e}")
                raise
        return ""

def get_llm(name: str) -> CustomLLM:
    """
    Возвращает экземпляр CustomLLM для заданной модели.
    """
    config = MODELS.get(name)
    if not config:
        raise ValueError(f"Модель '{name}' не сконфигурирована.")
    return CustomLLM(config)
