NewQuant/strategy_manager/core/manager.py

import os
import sys
import time
import psutil
import subprocess
import logging
from pathlib import Path
from typing import Dict, Any, List
from datetime import datetime
import importlib.util
import json  # 确保导入json模块

# ==================== 动态路径配置 ====================
from core.path_utils import add_project_root_to_path
from core.whitelist_manager import WhitelistManager

# 添加项目根路径到sys.path
PROJECT_ROOT = add_project_root_to_path()


# ==================================================

class StrategyManager:
    def __init__(self, config_path: str = "config/main.json"):
        self.config = self._load_main_config(config_path)
        self.strategies_dir = Path("strategies")
        self.logs_dir = Path(self.config["logs_dir"])
        self.status_file = Path(self.config["status_file"])
        self.pid_dir = Path(self.config["pid_dir"])

        # 创建目录
        self.logs_dir.mkdir(exist_ok=True)
        self.pid_dir.mkdir(exist_ok=True)

        # 配置管理器日志
        self._setup_logger()

        # 初始化白名单管理器
        self.whitelist_manager = WhitelistManager()
        self.logger.info("📋 白名单管理器已初始化")

        self.strategies: Dict[str, Dict[str, Any]] = {}
        self.logger.info("🔄 正在加载策略配置...")
        self.load_strategies()
        self.logger.info("✅ 策略加载完成，共发现 %d 个策略", len(self.strategies))

    def _setup_logger(self):
        """配置管理器日志"""
        log_file = self.logs_dir / "manager.log"
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s [%(levelname)s] %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler(sys.stdout)
            ],
            force=True
        )
        self.logger = logging.getLogger("StrategyManager")

    def _load_main_config(self, config_path: str) -> Dict[str, Any]:
        path = Path(config_path)
        if not path.exists():
            self.logger.warning("⚠️ 主配置文件不存在，使用默认配置")
            return {
                "logs_dir": "logs",
                "status_file": "status.json",
                "pid_dir": "pids"
            }
        with open(path, 'r') as f:
            return json.load(f)

    def load_strategies(self):
        """递归扫描 strategies/ 目录，查找 .py 配置文件"""
        self.strategies = {}
        if not self.strategies_dir.exists():
            self.logger.error("❌ 策略配置目录不存在: %s", self.strategies_dir)
            return

        for config_file in self.strategies_dir.rglob("*.py"):
            try:
                spec = importlib.util.spec_from_file_location(
                    f"config_{config_file.stem}", config_file
                )
                config_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(config_module)

                if not hasattr(config_module, 'CONFIG'):
                    raise ValueError("配置文件缺少 CONFIG 变量")
                config = config_module.CONFIG

                required = ['name', 'strategy_class', 'enabled', 'engine_params', 'strategy_params']
                for field in required:
                    if field not in config:
                        raise ValueError("配置缺少必要字段: {}".format(field))

                relative_path = config_file.relative_to(self.strategies_dir)
                strategy_name = relative_path.parent.name
                symbol = config_file.stem
                strategy_key = "{}_{}".format(strategy_name, symbol)

                self.strategies[strategy_key] = {
                    "strategy_name": strategy_name,
                    "symbol": symbol,
                    "config_file": str(config_file),
                    "config": config,
                    "status": "stopped",
                    "pid": None,
                    "started_at": None,
                    "uptime": None
                }
                self.logger.info("📄 加载配置: %s", config_file)
            except Exception as e:
                self.logger.error("❌ 加载配置失败 %s: %s", config_file, e, exc_info=True)

    def get_status(self) -> Dict[str, Any]:
        """获取完整状态（包含白名单信息）"""
        self._refresh_status()

        # 构建状态数据
        status = {
            "timestamp": datetime.now().isoformat(),
            "total": len(self.strategies),
            "running": sum(1 for s in self.strategies.values() if s["status"] == "running"),
            "strategies": self.strategies
        }

        # 添加白名单信息到每个策略
        for name, info in status["strategies"].items():
            info["in_whitelist"] = self.whitelist_manager.is_in_whitelist(name)
            info["whitelist_enabled"] = self.whitelist_manager.is_enabled_in_whitelist(name)

        # 添加自动启动状态
        auto_start_status = self.whitelist_manager.get_auto_start_status()
        status["whitelist_auto_start_today"] = auto_start_status["should_auto_start"]
        status["whitelist_last_date"] = auto_start_status["last_auto_start_date"]
        status["whitelist_total"] = auto_start_status["whitelist_count"]
        status["whitelist_enabled"] = auto_start_status["enabled_count"]

        return status

    def _refresh_status(self):
        """刷新进程状态 - 双重验证"""
        for name, info in self.strategies.items():
            pid_file = self.pid_dir / "{}.pid".format(name)

            if pid_file.exists():
                try:
                    with open(pid_file, 'r') as f:
                        pid = int(f.read().strip())

                    if psutil.pid_exists(pid):
                        try:
                            proc = psutil.Process(pid)
                            if "python" in proc.name().lower():
                                info["status"] = "running"
                                info["pid"] = pid
                                if info["started_at"]:
                                    started = datetime.fromisoformat(info["started_at"])
                                    uptime = datetime.now() - started
                                    info["uptime"] = str(uptime).split('.')[0]
                                continue
                        except (psutil.NoSuchProcess, psutil.AccessDenied):
                            pass

                    self._cleanup_stopped_strategy(name, pid_file)
                except Exception as e:
                    self.logger.warning("⚠️ 刷新状态失败 %s: %s", name, e)
                    self._cleanup_stopped_strategy(name, pid_file)
            else:
                info["status"] = "stopped"
                info["pid"] = None
                info["started_at"] = None
                info["uptime"] = None

    def _is_running(self, name: str) -> bool:
        """检查策略是否运行中"""
        self._refresh_status()
        info = self.strategies[name]
        if not info["pid"]:
            return False

        pid_file = self.pid_dir / "{}.pid".format(name)
        if not pid_file.exists():
            return False

        try:
            with open(pid_file, 'r') as f:
                pid = int(f.read().strip())

            if psutil.pid_exists(pid):
                try:
                    proc = psutil.Process(pid)
                    return "python" in proc.name().lower()
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    return False
            return False
        except:
            return False

    def start_strategy(self, name: str) -> bool:
        """启动单个策略（守护进程模式）"""
        if name not in self.strategies:
            self.logger.error("❌ 策略不存在: %s", name)
            return False

        if self._is_running(name):
            self.logger.warning("⚠️ 策略已在运行: %s", name)
            return False

        info = self.strategies[name]
        config_file = Path(info["config_file"])

        # 创建策略专属日志目录
        strategy_log_dir = self.logs_dir / info["strategy_name"]
        strategy_log_dir.mkdir(exist_ok=True, parents=True)
        log_file = strategy_log_dir / "{}.log".format(info["symbol"])

        self.logger.info("\n" + "=" * 80)
        self.logger.info("🚀 启动策略: %s", name)
        self.logger.info("📄 配置文件: %s", config_file)
        self.logger.info("📝 日志文件: %s", log_file)

        try:
            # 启动子进程 - 关键修改：脱离终端会话
            with open(log_file, 'a') as f:
                # 使用 start_new_session=True 创建新会话，防止终端关闭影响
                # stdin 重定向到 /dev/null，完全脱离终端
                with open(os.devnull, 'r') as devnull:
                    process = subprocess.Popen(
                        [sys.executable, "launcher.py", "--config", str(config_file)],
                        stdout=f,
                        stderr=subprocess.STDOUT,
                        stdin=devnull,  # 关键：从/dev/null读取输入
                        # start_new_session=True,  # 关键：创建新会话，脱离终端
                        cwd=Path.cwd()
                    )

            # 更新状态
            info["pid"] = process.pid
            info["status"] = "running"
            info["started_at"] = datetime.now().isoformat()
            info["uptime"] = "00:00:00"

            # 保存PID文件
            pid_file = self.pid_dir / "{}.pid".format(name)
            with open(pid_file, 'w') as f:
                f.write(str(process.pid))

            self._save_status()
            self.logger.info("✅ 启动成功! PID: %d", process.pid)
            self.logger.info("ℹ️  该进程已脱离终端会话，关闭窗口不会停止策略")
            self.logger.info("=" * 80)
            return True

        except Exception as e:
            self.logger.error("❌ 启动失败: %s", e, exc_info=True)
            self._cleanup_stopped_strategy(name, self.pid_dir / "{}.pid".format(name))
            return False

    def stop_strategy(self, name: str, timeout: int = 30) -> bool:
        """停止单个策略"""
        if name not in self.strategies:
            self.logger.error("❌ 策略不存在: %s", name)
            return False

        if not self._is_running(name):
            self.logger.warning("⚠️ 策略未运行: %s", name)
            return False

        info = self.strategies[name]

        try:
            pid = info["pid"]
            process = psutil.Process(pid)

            self.logger.info("\n" + "=" * 80)
            self.logger.info("⏹️  正在停止策略: %s (PID: %d)", name, pid)

            # 优雅终止
            process.terminate()

            try:
                process.wait(timeout=timeout)
                self.logger.info("✅ 已优雅停止: %s", name)
            except psutil.TimeoutExpired:
                self.logger.warning("⏱️  超时，强制终止: %s", name)
                process.kill()
                process.wait()

            # 清理状态
            self._cleanup_stopped_strategy(name, self.pid_dir / "{}.pid".format(name))
            self._save_status()
            self.logger.info("=" * 80)
            return True

        except Exception as e:
            self.logger.error("❌ 停止失败 %s: %s", name, e, exc_info=True)
            return False

    def restart_strategy(self, name: str) -> bool:
        """重启策略"""
        self.logger.info("\n" + "=" * 80)
        self.logger.info("🔄 正在重启: %s", name)
        self.stop_strategy(name)
        time.sleep(2)
        return self.start_strategy(name)

    def start_all(self):
        """启动所有启用的策略"""
        self.logger.info("\n" + "=" * 100)
        self.logger.info("🚀 正在启动所有启用的策略...")
        self.logger.info("=" * 100)

        started = []
        for name, info in self.strategies.items():
            if info["config"]["enabled"] and not self._is_running(name):
                if self.start_strategy(name):
                    started.append(name)

        self.logger.info("\n✅ 成功启动 %d 个策略", len(started))
        if started:
            self.logger.info("📋 策略: %s", ", ".join(started))

    def stop_all(self):
        """停止所有运行的策略"""
        self.logger.info("\n" + "=" * 100)
        self.logger.info("⏹️  正在停止所有运行的策略...")
        self.logger.info("=" * 100)

        stopped = []
        for name in self.strategies.keys():
            if self._is_running(name):
                if self.stop_strategy(name):
                    stopped.append(name)

        self.logger.info("\n✅ 成功停止 %d 个策略", len(stopped))
        if stopped:
            self.logger.info("📋 策略: %s", ", ".join(stopped))

    def _cleanup_stopped_strategy(self, name: str, pid_file: Path):
        """清理已停止的策略状态"""
        pid_file.unlink(missing_ok=True)
        info = self.strategies[name]
        info["status"] = "stopped"
        info["pid"] = None
        info["started_at"] = None
        info["uptime"] = None

    def _save_status(self):
        """状态持久化（修复：排除不可序列化的config字段）"""
        try:
            status = self.get_status()

            # 创建JSON安全的版本（排除config字段，因为它可能包含timedelta等不可序列化对象）
            status_for_json = status.copy()
            status_for_json["strategies"] = {}

            for name, info in status["strategies"].items():
                # 复制所有字段除了config
                strategy_info = {k: v for k, v in info.items() if k != "config"}
                status_for_json["strategies"][name] = strategy_info

            with open(self.status_file, 'w') as f:
                json.dump(status_for_json, f, indent=2, ensure_ascii=False)

            self.logger.debug("💾 状态已保存到 %s", self.status_file)
        except Exception as e:
            self.logger.error("❌ 保存状态失败: %s", e, exc_info=True)

    # ==================== 白名单管理方法 ====================

    def add_to_whitelist(self, name: str) -> bool:
        """
        添加策略到白名单

        Args:
            name: 策略标识符

        Returns:
            是否添加成功
        """
        if name not in self.strategies:
            self.logger.error("❌ 策略不存在: %s", name)
            return False

        if self.whitelist_manager.add(name, enabled=True):
            self.logger.info("✅ 添加到白名单: %s", name)
            self._save_status()
            return True

        return False

    def remove_from_whitelist(self, name: str) -> bool:
        """
        从白名单移除策略

        Args:
            name: 策略标识符

        Returns:
            是否移除成功
        """
        if self.whitelist_manager.remove(name):
            self.logger.info("✅ 从白名单移除: %s", name)
            self._save_status()
            return True

        return False

    def set_whitelist_enabled(self, name: str, enabled: bool) -> bool:
        """
        设置策略在白名单中的启用状态

        Args:
            name: 策略标识符
            enabled: 是否启用

        Returns:
            是否设置成功
        """
        if self.whitelist_manager.set_enabled(name, enabled):
            self.logger.info("✅ 设置白名单状态: %s -> %s", name, enabled)
            self._save_status()
            return True

        return False

    def auto_start_whitelist_strategies(self) -> Dict[str, bool]:
        """
        自动启动白名单中所有未运行的策略
        一天只执行一次

        Returns:
            Dict[str, bool]: 每个策略的启动结果
        """
        if not self.whitelist_manager.should_auto_start_today():
            self.logger.info("⏰ 今天已经执行过自动启动，跳过")
            return {}

        self.logger.info("🚀 开始执行白名单自动启动...")

        results = {}
        whitelist = self.whitelist_manager.get_all()

        for name, config in whitelist.items():
            if not config.get("enabled", True):
                self.logger.info("⏭️  跳过禁用策略: %s", name)
                continue

            if name not in self.strategies:
                self.logger.warning("⚠️  策略不在系统中: %s", name)
                continue

            # 检查是否已在运行
            if self._is_running(name):
                self.logger.info("✅ 策略已在运行: %s", name)
                results[name] = True
                continue

            # 尝试启动
            self.logger.info("🚀 启动白名单策略: %s", name)
            success = self.start_strategy(name)

            # 记录启动结果
            results[name] = success

            if success:
                self.logger.info("✅ 白名单策略启动成功: %s", name)
            else:
                self.logger.error("❌ 白名单策略启动失败: %s", name)

        # 更新日期
        self.whitelist_manager.update_last_auto_start_date(
            datetime.now().date().isoformat()
        )

        # 统计结果
        success_count = sum(1 for v in results.values() if v)
        fail_count = len(results) - success_count

        self.logger.info("📊 白名单自动启动完成: 成功 %d, 失败 %d", success_count, fail_count)

        return results


def print_status_table(status: Dict[str, Any]):
    """格式化打印状态表格"""
    print("\n" + "=" * 130)
    print("策略状态总览 (更新时间: {})".format(status['timestamp']))
    print("总计: {} | 运行中: {} | 已停止: {}".format(
        status['total'], status['running'], status['total'] - status['running']
    ))
    print("=" * 130)

    if not status["strategies"]:
        print("未找到任何策略")
        return

    print(
        "配置标识                          策略名称                              状态      PID       运行时长        启动时间")
    print("-" * 130)

    for name, info in status["strategies"].items():
        status_text = "RUNNING" if info["status"] == "running" else "STOPPED"
        pid_text = str(info["pid"]) if info["pid"] else "-"
        uptime_text = info["uptime"] if info["uptime"] else "-"
        started_text = info["started_at"][:19] if info["started_at"] else "-"

        print("{:<35} {:<40} {:<10} {:<10} {:<15} {:<25}".format(
            name, info['config']['name'], status_text, pid_text, uptime_text, started_text
        ))

    print("=" * 130)


if __name__ == "__main__":
    manager = StrategyManager()

    if len(sys.argv) > 1:
        command = sys.argv[1]
        if command == "status":
            print_status_table(manager.get_status())
        elif command == "start-all":
            manager.start_all()
        elif command == "stop-all":
            manager.stop_all()
        elif command.startswith("start:"):
            name = command.split(":", 1)[1]
            manager.start_strategy(name)
        elif command.startswith("stop:"):
            name = command.split(":", 1)[1]
            manager.stop_strategy(name)
        elif command.startswith("restart:"):
            name = command.split(":", 1)[1]
            manager.restart_strategy(name)
        else:
            print("未知命令:", command)
            print("用法: python manager.py [status|start-all|stop-all|start:NAME|stop:NAME|restart:NAME]")
    else:
        print("用法: python manager.py [status|start-all|stop-all|start:NAME|stop:NAME|restart:NAME]")