import os
import aiohttp
import numpy as np
from openai import AsyncOpenAI
import time
from lightrag.utils import setup_logger, EmbeddingFunc
import asyncio
from typing import List, Dict
import logging
import yaml
from src.project_rag import ProjectRAGManager

setup_logger("lightrag", level="INFO")
logger = logging.getLogger("lightrag")

def load_config():
    with open('config.yaml', 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

config = load_config()
MAX_CONCURRENT_LLM_REQUESTS = config['performance']['max_concurrent_llm_requests']
MAX_RETRIES = 3
RETRY_DELAY = 1.0

openai_client = None

def get_openai_client():
    """Получить или создать клиент OpenAI"""
    global openai_client
    if openai_client is None:
        openai_client = AsyncOpenAI(
            base_url=os.getenv('OPENAI_BASE_URL'),
            api_key=os.getenv('OPENAI_API_KEY')
        )
    return openai_client

async def process_single_llm_request(client: AsyncOpenAI, messages: List[Dict], **kwargs) -> str:
    """Обработка одного запроса к LLM с повторными попытками"""
    for attempt in range(MAX_RETRIES):
        try:
            start_time = time.time()
            
            completion = await client.chat.completions.create(
                model=config['models']['llm_model'],
                messages=messages,
                temperature=kwargs.get('temperature', 0.3),
                max_tokens=kwargs.get('max_tokens', 5000),
                stop=kwargs.get('stop', None)
            )
            
            processing_time = time.time() - start_time
            logger.debug(f"LLM запрос обработан за {processing_time:.2f}s")
            
            return completion.choices[0].message.content
            
        except Exception as e:
            if attempt < MAX_RETRIES - 1:
                logger.warning(f"Ошибка при запросе к LLM (попытка {attempt + 1}/{MAX_RETRIES}): {e}")
                await asyncio.sleep(RETRY_DELAY * (2 ** attempt))
            else:
                logger.error(f"Не удалось выполнить запрос к LLM после {MAX_RETRIES} попыток: {e}")
                raise e

async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
    """
    Оригинальная функция для работы с OpenAI-совместимым API
    (сохранена для обратной совместимости)
    """
    client = get_openai_client()
    
    messages = []
    
    if system_prompt:
        messages.append({
            'role': 'system', 
            'content': system_prompt
        })
    
    if history_messages:
        for msg in history_messages:
            if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                messages.append(msg)
            elif isinstance(msg, str):
                messages.append({
                    'role': 'user',
                    'content': msg
                })
    
    messages.append({
        'role': 'user', 
        'content': prompt
    })
    
    logger.info(f"Отправляем запрос к OpenAI API: {len(messages)} сообщений")
    return await process_single_llm_request(client, messages, **kwargs)

async def embedding_func(texts: List[str]) -> np.ndarray:
    """
    ОПТИМИЗИРОВАННАЯ функция для VLLM - без внутреннего батчинга!
    
    LightRAG уже передает оптимально размеченные батчи.
    Мы просто отправляем их в VLLM как есть.
    
    Args:
        texts: Список текстов
    
    Returns:
        numpy массив с эмбеддингами
    """
    if not texts:
        return np.array([])
    
    if isinstance(texts, str):
        texts = [texts]
    
    logger.info(f" Отправляем батч из {len(texts)} текстов в VLLM")
    
    for attempt in range(MAX_RETRIES):
        try:
            start_time = time.time()
            
            payload = {
                "input": texts,
                "model": config['models']['embedding_model']
            }
            await asyncio.sleep(0.01)  # 100ms пауза
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    os.getenv('EMBEDDING_API_URL'), 
                    json=payload, 
                    ssl=False,
                    timeout=aiohttp.ClientTimeout(total=120)
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        raise Exception(f"VLLM API Error: {response.status} - {error_text}")
                    
                    result = await response.json()
                    
                    if "data" not in result:
                        raise Exception("Invalid response structure from VLLM API")
                    
                    embeddings = result["data"]
                    batch_embeddings = [e["embedding"] for e in embeddings]
                    
                    processing_time = time.time() - start_time
                    
                    return np.array(batch_embeddings, dtype=np.float32)
                    
        except Exception as e:
            if attempt < MAX_RETRIES - 1:
                wait_time = RETRY_DELAY * (2 ** attempt)
                logger.warning(f"Ошибка при обращении к VLLM (попытка {attempt + 1}/{MAX_RETRIES}): {e}")
                await asyncio.sleep(wait_time)
            else:
                logger.error(f"Не удалось получить эмбеддинги от VLLM после {MAX_RETRIES} попыток: {e}")
                raise


embedding_func_obj = EmbeddingFunc(
    embedding_dim=1024,
    max_token_size=7192,
    func=embedding_func
)

rag_manager = ProjectRAGManager(
    llm_func=llm_model_func,
    embedding_func=embedding_func_obj
)
