# scripts/launch_llm_server.py

"""
Скрипт для запуска локального API-сервера vLLM.
Поднимает процесс в фоне, мониторит логи и health-эндпоинт до 3 минут.
Все пути и параметры загружаются из конфигурации Hydra.
"""
import logging
import os
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import List, Optional
from urllib.parse import ParseResult, urlparse

import hydra
import requests
from omegaconf import DictConfig

# Настройка логгера
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger: logging.Logger = logging.getLogger(__name__)

def tail_log(path: str, n: int = 20) -> List[str]:
    """
    Читает последние N строк из лог-файла, фильтруя важные сообщения.
    
    Args:
        path (str): Абсолютный путь к файлу логов.
        n (int): Количество строк для чтения.
        
    Returns:
        List[str]: Список отфильтрованных строк лога.
    """
    if not os.path.exists(path):
        return ["(лог еще не создан)"]
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            lines: List[str] = f.readlines()[-n:]
        
        filt: List[str] = [ln for ln in lines if re.search(r"\b(INFO|WARNING|ERROR)\b", ln)]
        return filt if filt else lines
    except Exception as e:
        return [f"(не удалось прочитать лог: {e})"]

@hydra.main(version_base=None, config_path="../configs", config_name="main_config")
def main(cfg: DictConfig) -> None:
    """
    Основная функция лаунчера. Формирует команду запуска из конфигурации,
    стартует фоновый процесс и ожидает его готовности.
    
    Args:
        cfg (DictConfig): Конфигурация проекта (Hydra).
    """
    # 1. Определяем абсолютный путь к корню проекта
    project_root: Path = Path(__file__).resolve().parents[1]
    
    # 2. Извлекаем параметры из конфигурации
    llm_cfg: DictConfig = cfg.get("llm", {})
    model: str = llm_cfg.get("model", "Qwen/Qwen3-VL-8B-Instruct-FP8")
    base_url: str = llm_cfg.get("api_url", "http://127.0.0.1:8000/v1")
    
    max_model_len: str = str(llm_cfg.get("max_model_len", 8192))
    gpu_util: str = str(llm_cfg.get("gpu_memory_utilization", 0.85))
    dtype: str = llm_cfg.get("dtype", "bfloat16")
    served_name: str = llm_cfg.get("served_model_name", model.split("/")[-1].lower())
    
    log_file_rel: str = llm_cfg.get("log_file", "logs/vllm_server.log")
    log_path: str = str(project_root / log_file_rel)
    
    os.makedirs(os.path.dirname(log_path), exist_ok=True)
    
    download_dir: Optional[str] = llm_cfg.get("download_dir")

    if not model or not base_url:
        logger.error("В конфигурации отсутствуют обязательные параметры (model или api_url).")
        sys.exit(1)

    parsed: ParseResult = urlparse(base_url)
    host: str = parsed.hostname or "127.0.0.1"
    port: int = parsed.port or 8000
    scheme: str = parsed.scheme or "http"
    health_url: str = f"{scheme}://{host}:{port}/health"

    # 3. Собираем базовую команду
    cmd: List[str] = [
        "nohup", sys.executable, "-m", "vllm.entrypoints.openai.api_server",
        "--model", model,
        "--host", host,
        "--port", str(port),
        "--max-model-len", max_model_len,
        "--dtype", dtype,
        "--gpu-memory-utilization", gpu_util,
        "--served-model-name", served_name,
        "--trust-remote-code"
    ]
    
    if download_dir:
        cmd.extend(["--download-dir", str(project_root / download_dir)])

    logger.info(f"⚙️ Стартую vLLM с командой: {' '.join(cmd)}")
    logger.info(f"📄 Лог-файл: {log_path}")
    logger.info(f"🌐 Health URL: {health_url}")

    try:
        # 4. Запускаем сервер в фоне
        log_fh = open(log_path, "w", encoding="utf-8", errors="ignore")
        proc: subprocess.Popen = subprocess.Popen(
            cmd, 
            stdout=log_fh, 
            stderr=subprocess.STDOUT, 
            stdin=subprocess.DEVNULL,
            start_new_session=True,
            env=os.environ.copy()
        )
        
        # 5. Мониторинг запуска — до 3 минут
        logger.info("⏳ Ожидаем загрузки модели (может занять до 3 минут)...")
        is_ready: bool = False
        max_checks: int = 60
        interval_sec: int = 15
        
        for i in range(max_checks):
            if i % 3 == 0:
                logger.info(f"📊 Текущие логи (проверка {i // 3 + 1}):")
                for ln in tail_log(log_path, n=10):
                    logger.info(f"   > {ln.rstrip()}")

            ret: Optional[int] = proc.poll()
            if ret is not None:
                logger.error(f"❌ Процесс vLLM неожиданно завершился (код {ret}).")
                logger.error("Последние строки лога:")
                for ln in tail_log(log_path, n=30):
                    logger.error(f"   > {ln.rstrip()}")
                sys.exit(1)

            try:
                resp: requests.Response = requests.get(health_url, timeout=5)
                if resp.status_code == 200:
                    logger.info(f"✅ Сервер готов! Время загрузки: ~{i * interval_sec} секунд.")
                    is_ready = True
                    break
            except Exception:
                pass

            logger.info(f"🔄 [{i + 1}/{max_checks}] Инициализация... ждем {interval_sec} сек.")
            time.sleep(interval_sec)

        if not is_ready:
            logger.error("❌ Сервер не стал готов за 3 минуты. Завершаем ожидание.")
            sys.exit(1)
            
        # Успешный запуск - выводим все нужные параметры для клиента
        logger.info("🚀 ИИ-бэкенд успешно запущен и работает в фоне!")
        logger.info(f"🔗 Используйте в клиенте base_url: {base_url}")
        logger.info(f"🧠 Имя для обращения по API (served_model_name): {served_name}")
        logger.info(f"📦 Оригинальное название модели: {model}")

    except Exception as e:
        logger.error(f"❌ Критическая ошибка при запуске: {e}", exc_info=True)
        sys.exit(1)

if __name__ == "__main__":
    main()