import argparse
import os
import sys
import shutil
import subprocess

def run_inference(input_path, output_path):
    # 1. 获取绝对路径
    cwd = os.path.dirname(os.path.abspath(__file__)) # 当前脚本所在目录
    python_exe = sys.executable # 当前环境的 python
    inference_script = os.path.join(cwd, "inference_codeformer.py")
    
    print(f"🚀 [CodeFormer] 启动中...")

    # 2. 准备临时输入目录 (CodeFormer 喜欢读文件夹)
    # 它的逻辑是：输入文件夹 -> 结果文件夹
    temp_in = os.path.join(cwd, "temp_worker_input")
    if os.path.exists(temp_in): shutil.rmtree(temp_in)
    os.makedirs(temp_in, exist_ok=True)
    
    # 复制图片进去
    img_name = os.path.basename(input_path)
    shutil.copy(input_path, os.path.join(temp_in, img_name))
    
    # 3. 构造命令
    # -w 0.5: 平衡画质和保真度
    # --face_upsample: 这一步会把人脸贴回原图，这是我们需要的
    cmd = [
        python_exe, 
        inference_script,
        "-w", "0.5", 
        "--input_path", temp_in,
        "--face_upsample" 
    ]
    
    # 4. 执行命令
    try:
        # cwd=cwd 保证它能找到 weights 文件夹
        subprocess.run(cmd, cwd=cwd, check=True)
        
        # 5. 提取结果
        # 默认结果路径: results/temp_worker_input_0.5/final_results/图片名.png
        res_dir = os.path.join(cwd, "results", "temp_worker_input_0.5", "final_results")
        res_file = os.path.join(res_dir, img_name)
        
        if os.path.exists(res_file):
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            shutil.copy(res_file, output_path)
            print(f"✅ CodeFormer 处理完成: {output_path}")
        else:
            print(f"❌ 未找到结果文件: {res_file}")
            print("   (可能是因为图片里没有人脸，或者模型下载失败)")
            
    except Exception as e:
        print(f"❌ 执行出错: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', required=True)
    parser.add_argument('-o', '--output', required=True)
    parser.add_argument('-m', '--model', help="Ignored")
    args = parser.parse_args()
    
    run_inference(args.input, args.output)