#!/usr/bin/env python3
"""
启动脚本：同时运行 vLLM 服务器和 Streamlit 应用
保留原有的 Streamlit 入口，添加 vLLM 后台服务
"""

import subprocess
import time
import os
import signal
import sys
import threading
import logging
from pathlib import Path

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 添加项目根目录到路径
sys.path.append(str(Path(__file__).parent.parent))

class ServiceManager:
    def __init__(self):
        self.processes = []
        self.running = True
        
    def run_vllm(self):
        """启动 vLLM 服务器（后台服务）"""
        logger.info("🚀 正在启动 vLLM 推理服务器...")
        
        # vLLM 启动命令（优化显存使用）
        cmd = [
            "python", "-m", "vllm.entrypoints.openai.api_server",
            "--model", "Qwen/Qwen3.5-27B-GPTQ-Int4",
            "--port", "8000",
            "--host", "0.0.0.0",
            "--tensor-parallel-size", "1",
            "--quantization", "gptq",
            "--max-model-len", "32768",
            "--gpu-memory-utilization", "0.85",  # 保留 15% 显存给 Streamlit 等
            "--enforce-eager",
            "--max-num-batched-tokens", "8192",
            "--max-num-seqs", "4",
            "--disable-log-stats",  # 减少日志输出
        ]
        
        # 设置 CUDA 环境
        env = os.environ.copy()
        #env["CUDA_VISIBLE_DEVICES"] = "0"
        #env["VLLM_USE_TRITON"] = "1"  # 使用 Triton 加速
        
        # 启动进程
        process = subprocess.Popen(
            cmd,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1
        )
        
        self.processes.append(("vllm", process))
        
        # 在单独的线程中读取输出
        def log_output():
            for line in process.stdout:
                if line.strip():
                    logger.info(f"[vLLM] {line.strip()}")
        
        threading.Thread(target=log_output, daemon=True).start()
        
        # 等待 vLLM 就绪
        self.wait_for_vllm()
        
    def wait_for_vllm(self, max_retries=60):
        """等待 vLLM 服务器就绪"""
        import requests
        
        logger.info("⏳ 等待 vLLM 服务器初始化...")
        for i in range(max_retries):
            try:
                response = requests.get("http://localhost:8000/health")
                if response.status_code == 200:
                    logger.info("✅ vLLM 服务器已就绪！")
                    return True
            except:
                pass
            
            # 显示进度
            if i % 10 == 0:
                logger.info(f"⏳ 仍在等待 vLLM... ({i+1}/{max_retries})")
            time.sleep(2)
        
        logger.error("❌ vLLM 启动超时")
        return False
    
    def run_streamlit(self):
        """启动 Streamlit 应用（主服务）"""
        logger.info("🖥️ 正在启动 Streamlit 应用...")
        
        # 使用原有的入口点
        cmd = [
            "streamlit", "run",
            "streamlit_app.py",
            "--server.port=8501",
            "--server.address=0.0.0.0",
            "--server.enableCORS=false",
            "--server.enableXsrfProtection=false",
            "--server.maxUploadSize=10",  # 限制上传大小
        ]
        
        process = subprocess.Popen(cmd)
        self.processes.append(("streamlit", process))
        logger.info("✅ Streamlit 应用已启动")
        
    def monitor_services(self):
        """监控服务状态"""
        while self.running:
            time.sleep(10)
            for name, process in self.processes:
                if process.poll() is not None:
                    logger.error(f"❌ {name} 服务意外停止，退出码: {process.returncode}")
                    self.stop_all()
                    sys.exit(1)
    
    def stop_all(self, signum=None, frame=None):
        """停止所有服务"""
        logger.info("🛑 正在停止所有服务...")
        self.running = False
        for name, process in self.processes:
            logger.info(f"正在停止 {name}...")
            process.terminate()
            
        # 等待进程结束
        for name, process in self.processes:
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                logger.warning(f"强制终止 {name}")
                process.kill()
        
        logger.info("✅ 所有服务已停止")
        sys.exit(0)
    
    def run(self):
        """运行所有服务"""
        # 注册信号处理
        signal.signal(signal.SIGINT, self.stop_all)
        signal.signal(signal.SIGTERM, self.stop_all)
        
        # 启动 vLLM（后台）
        self.run_vllm()
        
        # 启动 Streamlit（前台）
        self.run_streamlit()
        
        # 开始监控
        self.monitor_services()

if __name__ == "__main__":
    manager = ServiceManager()
    manager.run()