import os
import shutil
import subprocess
import glob

class AgentToolbox:
    def __init__(self):
        self.root_dir = r"G:\IR_Experiment"
        self.output_base = os.path.join(self.root_dir, "Agent_Workspace")
        os.makedirs(self.output_base, exist_ok=True)

        # ================= 环境配置 (关键！) =================
        # 请根据你的实际情况修改这里的 python.exe 路径
        # 1. ir_final 环境 (用于 DarkIR, NAFNet, PromptIR)
        self.env_main = r"D:\conda\envs\ir_final\python.exe"
        # 2. swinir_env 环境 (用于 SwinIR)
        self.env_swinir = r"D:\conda\envs\swinir_env\python.exe" 
        
        # 如果找不到路径，尝试用系统默认的 'python' (前提是你激活了对应环境)
        if not os.path.exists(self.env_main): self.env_main = "python"
        if not os.path.exists(self.env_swinir): self.env_swinir = "python"

    def _run_cmd(self, cmd, cwd):
        """执行命令行的通用函数"""
        print(f"\n[Toolbox] 正在执行: {cmd} ...")
        try:
            subprocess.run(cmd, shell=True, check=True, cwd=cwd)
            return True
        except subprocess.CalledProcessError as e:
            print(f"❌ 执行失败: {e}")
            return False

    # ================= 工具 1: DarkIR (低光增强) =================
    def call_darkir(self, image_path):
        print(f"🌙 [调用 DarkIR] 正在处理低光: {os.path.basename(image_path)}")
        tool_dir = os.path.join(self.root_dir, "DarkIR")
        
        # 1. 适配输入：DarkIR 也是读文件夹的，我们把图复制到它的 input 目录
        input_dir = os.path.join(tool_dir, "test_input") # 对应 run_darkir.py 里的路径
        if os.path.exists(input_dir): shutil.rmtree(input_dir)
        os.makedirs(input_dir, exist_ok=True)
        shutil.copy(image_path, os.path.join(input_dir, os.path.basename(image_path)))

        # 2. 调用我们之前写好的脚本
        # 注意：这里调用的是 run_darkir.py，确保它里面的路径是对的
        cmd = f'"{self.env_main}" run_darkir.py'
        
        if self._run_cmd(cmd, cwd=tool_dir):
            # 3. 提取输出
            # DarkIR 输出在 results 文件夹
            result_dir = os.path.join(tool_dir, "results")
            # 找到生成的文件
            res_files = glob.glob(os.path.join(result_dir, "*.*"))
            if res_files:
                # 把结果移动到 Agent 工作区
                out_name = f"darkir_{os.path.basename(image_path)}"
                final_path = os.path.join(self.output_base, out_name)
                shutil.copy(res_files[0], final_path)
                return final_path
        return None

    # ================= 工具 2: SwinIR (超分放大) =================
    def call_swinir(self, image_path, scale=4):
        print(f"🔍 [调用 SwinIR] 正在放大 {scale}倍: {os.path.basename(image_path)}")
        tool_dir = os.path.join(self.root_dir, "SwinIR")
        
        # 1. 适配输入
        temp_input = os.path.join(tool_dir, "testsets", "agent_temp")
        if os.path.exists(temp_input): shutil.rmtree(temp_input)
        os.makedirs(temp_input, exist_ok=True)
        shutil.copy(image_path, os.path.join(temp_input, os.path.basename(image_path)))

        # 2. 构建命令
        # 使用 Real-World x4 模型
        model_path = r"model_zoo/swinir/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth"
        cmd = f'"{self.env_swinir}" main_test_swinir.py --task real_sr --scale {scale} --model_path {model_path} --folder_lq testsets/agent_temp --tile 400'

        if self._run_cmd(cmd, cwd=tool_dir):
            # 3. 提取输出
            # SwinIR 结果通常在 results/swinir_real_sr_x4 里面
            result_dir = os.path.join(tool_dir, "results", f"swinir_real_sr_x{scale}")
            # 找最新生成的文件
            res_files = glob.glob(os.path.join(result_dir, "*.*"))
            if res_files:
                # SwinIR 会给文件名加后缀，我们找包含原名的那个
                target_file = [f for f in res_files if os.path.basename(image_path).split('.')[0] in f][-1]
                
                out_name = f"swinir_{os.path.basename(image_path)}"
                final_path = os.path.join(self.output_base, out_name)
                shutil.copy(target_file, final_path)
                return final_path
        return None

# ================= 模拟 Agent 调度逻辑 =================
if __name__ == "__main__":
    toolbox = AgentToolbox()
    
    # 1. 准备一张测试图 (你可以换成任何存在的图片路径)
    # 假设我们用之前生成的“低光”测试图
    original_img = r"G:\datasets\realblur_dataset_test\075_blur_1.png" # 确保这张图存在！
    
    if not os.path.exists(original_img):
        print("❌ 测试图不存在，请修改 original_img 路径")
        exit()

    print(f"🏁 开始处理任务: {original_img}")

    # --- 步骤 1: 先提亮 (DarkIR) ---
    bright_img = toolbox.call_darkir(original_img)
    
    if bright_img:
        print(f"✅ 第一步完成: {bright_img}")
        
        # --- 步骤 2: 再放大 (SwinIR) ---
        # 把第一步的结果喂给第二步
        final_img = toolbox.call_swinir(bright_img)
        
        if final_img:
            print(f"🎉 任务全部完成！最终结果: {final_img}")
        else:
            print("❌ 第二步 SwinIR 失败")
    else:
        print("❌ 第一步 DarkIR 失败")