# Prompts: use modules under prompts/ where they exist, and root files where applicable
from prompts.physics_prompts import physics_reasoner_system_prompt, chat_system_prompt
from prompts.algebra_prompts import algebra_reasoner_system_prompt, algebra_system_prompt
from prompts.biology_prompts import biology_system_prompt
from prompts.english_prompts import english_system_prompt
from prompts.informatics_prompts import informatics_system_prompt, informatics_reasoner_system_prompt
from prompts.chemistry_prompts import chemistry_system_prompt, chemistry_reasoner_system_prompt
from prompts.geography_prompts import geography_system_prompt
from prompts.statistics_prompts import statistics_system_prompt, statistics_reasoner_system_prompt
from prompts.social_studies_prompts import social_studies_system_prompt
from prompts.geometry_prompts import geometry_system_prompt, geometry_reasonner_prompt
from prompts.history_prompts import history_system_prompt
from prompts.tg_suffix import TG_PLAIN_TEXT_SUFFIX
import logging
from db_config import connection_pool
import mysql.connector
import requests
import json
import base64
import sys
import os
import time

from llm_providers.openai_provider import call_openai_stream
from llm_providers.claude_provider import call_claude_stream
from llm_providers.gemini_provider import call_gemini_stream
from llm_providers.openrouter_provider import call_openrouter_stream


# Настраиваем логирование
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# Кэш в памяти для информации о моделях, чтобы избежать повторных запросов к БД
_model_info_cache = {}


def get_model_info(subject_id, mode):
    """
    Получает информацию о модели из базы данных на основе предмета и режима чата.
    Результаты кэшируются в памяти для ускорения повторных вызовов.
    
    Args:
        subject_id (int): ID предмета (из таблицы subject)
        mode (int): Режим чата (1 - решатель задач, 0 - обычный чат)
    
    Returns:
        tuple: (api_url, api_key, model_name, provider_name) - или None в случае ошибки
    """
    cache_key = (subject_id, mode)
    if cache_key in _model_info_cache:
        logger.info(f"Информация о модели для {cache_key} найдена в кэше.")
        return _model_info_cache[cache_key]

    conn = None
    cursor = None
    try:
        conn = connection_pool.get_connection()
        cursor = conn.cursor(dictionary=True)
        
        # Запрос к базе данных для получения информации о модели
        query = """
        SELECT p.api_url, p.api_key, m.name, p.name as provider
        FROM model m 
        INNER JOIN subject_model sm ON m.id = sm.model_id 
        INNER JOIN provider p ON m.provider_id = p.id
        WHERE sm.subject_id = %s AND sm.reasoner_mode = %s
        """
        
        cursor.execute(query, (subject_id, mode))
        model_info = cursor.fetchone()

        fallback_used = False
        if not model_info and mode == 1:
            logger.warning(
                "Модель для предмета %s в reasoner-режиме отсутствует. "
                "Переключаюсь на чат-модель как fallback.",
                subject_id,
            )
            cursor.execute(query, (subject_id, 0))
            model_info = cursor.fetchone()
            fallback_used = bool(model_info)
        
        result = None
        if model_info:
            logger.info(f"Модель для режима {mode}: {model_info['name']} от провайдера {model_info['provider']} (из БД)")
            result = (model_info['api_url'], model_info['api_key'], model_info['name'], model_info['provider'])
            if fallback_used:
                logger.info(
                    "Для предмета %s reasoner-режим будет использовать чат-модель %s.",
                    subject_id,
                    model_info['name'],
                )
        else:
            logger.error(f"Модель для предмета {subject_id} и режима {mode} не найдена")

        _model_info_cache[cache_key] = result  # Сохраняем результат (даже None) в кэш
        return result
            
    except mysql.connector.Error as err:
        logger.error(f"Ошибка базы данных: {err}")
        return None
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def get_llm_response_stream(chat_history, mode, subject, trace_id=None, plain_text_mode=False):
    """
    Отправляет сообщения в API и возвращает ответ в виде потока
    
    Args:
        chat_history (list): Список сообщений
        mode (str): Модель для использования ("chat-model" или "reasoner-model")
        subject (int): Предмет (1 - физика, 2 - алгебра)
        trace_id (str, optional): ID для отслеживания запроса
        plain_text_mode (bool, optional): Если True, добавляет суффикс для простого текста (без LaTeX/Markdown)
    """
    try:
        # Логируем полученные сообщения
        logging.debug(f"Полученные сообщения для API: {chat_history}")

        # Преобразуем режим в числовой формат для функции get_model_info
        reasoner_mode = 1 if mode == "reasoner-model" else 0

        # Latency timing
        t_entry = time.perf_counter()
        t_model_start = time.perf_counter()
        # Получаем информацию о модели из базы данных
        model_info = get_model_info(subject, reasoner_mode)
        if not model_info:
            logger.error(f"Не удалось получить информацию о модели для предмета {subject} и режима {mode}")
            yield "Ошибка: не удалось получить информацию о модели"
            return
            
        api_url, api_key, model_name, provider = model_info
        t_model_end = time.perf_counter()
        logger.info(
            f"[latency][{trace_id or '-'}] Model lookup done: subject={subject}, mode={mode}, provider={provider}, "
            f"model={model_name}, model_lookup_ms={int((t_model_end - t_model_start)*1000)}"
        )
        logger.info(f"Использую модель: {model_name} от провайдера: {provider}")
        
       
        # Подготовка сообщений
        t_prep_start = time.perf_counter()
        for message in chat_history:
            if not isinstance(message.get("content"), (str, list)):
                message["content"] = str(message.get("content", ""))

        # Выбор system_prompt в зависимости от предмета и режима
        system_prompt = None
        if subject == 1:  # Физика
            if reasoner_mode == 1:
                system_prompt = physics_reasoner_system_prompt
            else:
                system_prompt = chat_system_prompt
        elif subject == 2:  # Алгебра
            if reasoner_mode == 1:
                system_prompt = algebra_reasoner_system_prompt
            else:
                system_prompt = algebra_system_prompt
        elif subject == 3:  # Геометрия
            if reasoner_mode == 1:
                system_prompt = geometry_reasonner_prompt
            else:
                system_prompt = geometry_system_prompt
        elif subject == 4:  # Химия
            # Для химии используем отдельные промпты для чата и решателя
            if reasoner_mode == 1:
                system_prompt = chemistry_reasoner_system_prompt
            else:
                system_prompt = chemistry_system_prompt
        elif subject == 5:  # Биология
            # Для биологии всегда используем только обычный режим чата
            system_prompt = biology_system_prompt
        elif subject == 8:  # Обществознание
            system_prompt = social_studies_system_prompt
        elif subject == 10:  # География
            system_prompt = geography_system_prompt
        elif subject == 13:  # Статистика и вероятность
            if reasoner_mode == 1:
                system_prompt = statistics_reasoner_system_prompt
            else:
                system_prompt = statistics_system_prompt
        elif subject == 11:  # english
            system_prompt = english_system_prompt
        elif subject == 12:  # Информатика
            # Use subject-specific prompts based on mode
            if reasoner_mode == 1:
                system_prompt = informatics_reasoner_system_prompt
            else:
                system_prompt = informatics_system_prompt
        elif subject == 6:  # История
            system_prompt = history_system_prompt


        if not system_prompt:
            logger.warning(f"System prompt не определен для предмета {subject} и режима {reasoner_mode}")
            system_prompt = chat_system_prompt  # По умолчанию используем общий промпт

        # Apply Telegram plain text suffix if requested
        if plain_text_mode and TG_PLAIN_TEXT_SUFFIX not in (system_prompt or ""):
            system_prompt = (system_prompt or "") + TG_PLAIN_TEXT_SUFFIX

        # Добавляем системное сообщение
        chat_history_with_system = chat_history.copy()
        chat_history_with_system.insert(0, {"role": "system", "content": system_prompt})
        t_prep_end = time.perf_counter()
        logger.info(
            f"[latency][{trace_id or '-'}] Messages prepared: prep_ms={int((t_prep_end - t_prep_start)*1000)}, "
            f"total_since_entry_ms={int((t_prep_end - t_entry)*1000)}"
        )

        logger.info(f"Prepared messages for dispatch: {chat_history_with_system}")
        
        try:
            # Time to first chunk across providers
            t_provider_start = time.perf_counter()
            first_chunk_logged = False

            def yield_with_first_chunk(provider_iter):
                nonlocal first_chunk_logged
                for chunk in provider_iter:
                    if not first_chunk_logged:
                        first_chunk_logged = True
                        t_first = time.perf_counter()
                        logger.info(
                            f"[latency][{trace_id or '-'}] First provider chunk: provider={provider}, model={model_name}, "
                            f"time_to_first_chunk_ms={int((t_first - t_provider_start)*1000)}, "
                            f"total_since_entry_ms={int((t_first - t_entry)*1000)}"
                        )
                    yield chunk
            if provider == "claude_client":
                for chunk in yield_with_first_chunk(
                    call_claude_stream(chat_history_with_system, model_name, api_key, system_prompt)
                ):
                    yield chunk
            elif provider == "gemini_client":
                system_instruction = system_prompt
                reasoner_mode_int = 1 if mode == "reasoner-model" else 0
                for chunk in yield_with_first_chunk(
                    call_gemini_stream(chat_history, model_name, api_key, system_instruction, subject, reasoner_mode_int)
                ):
                    yield chunk
            elif provider == "openrouter_client":
                for chunk in yield_with_first_chunk(
                    call_openrouter_stream(chat_history_with_system, model_name, api_key, api_url)
                ):
                    yield chunk
            elif provider == "openai_client" or provider == "deepseek_client":
                for chunk in yield_with_first_chunk(
                    call_openai_stream(chat_history_with_system, model_name, api_key, api_url)
                ):
                    yield chunk
            else:
                logger.error(f"Unknown provider: {provider}")
                yield f"Error: unknown provider {provider}"

        except TimeoutError:
            logger.error("Превышено время ожидания ответа от API")
            yield "Извините, время ожидания ответа истекло. Пожалуйста, попробуйте еще раз."
        except Exception as api_error:
            logger.error(f"Ошибка при вызове API: {api_error}", exc_info=True)
            yield f"Произошла ошибка: {str(api_error)}"

    except Exception as e:
        logging.error(f"Общая ошибка при обработке запроса: {e}", exc_info=True)
        yield f"Произошла общая ошибка: {str(e)}"


# Пример использования:
if __name__ == "__main__":
    import sys