import subprocess
import os
import time
import logging


class CommandRunner:
    def __init__(self, machine_config):
        """
        :param machine_config: config/machine.yaml 中 'systems' -> 'current_system' 对应的内容
        """
        self.config = machine_config
        self.logger = logging.getLogger("NEP_Auto")
        self.mode = self.config.get("type", "local")  # local 或 slurm

    def run(self, tool_name, cwd=".", wait=True, extra_args=""):
        """
        核心运行方法
        :param tool_name: machine.yaml 中 tools 下的键名 (如 'gpumd', 'vasp')
        :param cwd: 执行命令的工作目录
        :param wait: 是否等待命令结束 (True: 阻塞, False: 后台运行)
        :param extra_args: 附加在命令后的参数
        """
        # 1. 获取工具配置
        tool_conf = self.config.get("tools", {}).get(tool_name)
        if not tool_conf:
            self.logger.error(f"❌ 找不到工具配置: {tool_name}")
            raise ValueError(f"Tool {tool_name} not defined in machine.yaml")

        cmd = tool_conf.get("command")
        env_setup = tool_conf.get("env_setup", "")

        # 2. 组装命令 (Local 模式)
        if self.mode == "local":
            full_cmd = f"{cmd} {extra_args}"

            # 如果有环境加载脚本，用 && 连接
            if env_setup:
                full_cmd = f"{env_setup} && {full_cmd}"

            self.logger.info(f"⚙️ [Local] Executing: {full_cmd}")
            self.logger.info(f"   📂 Workdir: {cwd}")

            try:
                # 使用 bash 执行以支持 source 命令
                process = subprocess.Popen(
                    full_cmd,
                    shell=True,
                    cwd=cwd,
                    executable="/bin/bash",
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True
                )

                if wait:
                    stdout, stderr = process.communicate()
                    if process.returncode != 0:
                        self.logger.error(f"❌ Execution failed (Code {process.returncode})")
                        self.logger.error(f"Stderr: {stderr}")
                        raise RuntimeError(f"Command failed: {full_cmd}")
                    return True
                else:
                    return process  # 返回进程对象供监控

            except Exception as e:
                self.logger.error(f"❌ Runner Error: {str(e)}")
                raise

        # 3. Slurm 模式 (预留接口，暂未实现具体逻辑)
        elif self.mode == "slurm":
            self.logger.warning("⚠️ Slurm mode not fully implemented yet.")
            # 这里未来会生成 sbatch 脚本并提交
            return False