import argparse
import subprocess
import os
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description="Spawn multi-window batch runner")
    parser.add_argument("--start", type=int, required=True, help="Start requirement (1-based, inclusive, e.g. 31)")
    parser.add_argument("--end", type=int, required=True, help="End requirement (1-based, inclusive, e.g. 50)")
    parser.add_argument("--workers", type=int, default=7, help="Number of physical CMD windows to spawn")
    parser.add_argument("--skip-research", action="store_true", help="Skip Phase 1 and use existing raw cases")
    parser.add_argument("--research-only", action="store_true", help="Only run research phases, skip Phase 2 and 3")
    parser.add_argument("--use-claude-sdk", action="store_true", help="Use pure Anthropic SDK for Phase 2/3")
    parser.add_argument("--platforms", type=str, default="", help="Comma-separated list of platforms to search")
    args = parser.parse_args()

    start_idx = args.start - 1
    end_idx = args.end - 1
    total_tasks = end_idx - start_idx + 1
    
    if total_tasks <= 0:
        print("Invalid range")
        return

    # Assign indices to workers using round-robin
    worker_tasks = {i: [] for i in range(args.workers)}
    for i, req_idx in enumerate(range(start_idx, end_idx + 1)):
        worker_id = i % args.workers
        worker_tasks[worker_id].append(req_idx)
        
    base_dir = Path(__file__).parent
    
    print(f"Spawning {args.workers} actual CMD windows to process {total_tasks} requirements...")

    for worker_id, indices in worker_tasks.items():
        if not indices:
            continue
            
        # Create a temporary batch file for this worker window
        bat_path = base_dir / f"worker_{worker_id}.bat"
        with open(bat_path, "w", encoding="utf-8") as f:
            f.write("@echo off\n")
            f.write("cd /d %~dp0\n")
            f.write("if exist ..\\..\\.venv\\Scripts\\activate.bat call ..\\..\\.venv\\Scripts\\activate.bat\n")
            f.write("echo Worker Group %d starting...\n" % worker_id)
            for idx in indices:
                f.write(f"echo =======================================\n")
                f.write(f"echo Running Requirement Index: {idx}\n")
                f.write(f"echo =======================================\n")
                
                cmd = f"python run_pipeline.py --index {idx}"
                if args.skip_research:
                    cmd += " --skip-research"
                if args.research_only:
                    cmd += " --research-only"
                if args.use_claude_sdk:
                    cmd += " --use-claude-sdk"
                if args.platforms:
                    cmd += f" --platforms {args.platforms}"
                
                # 第 1 遍执行
                f.write(f"echo [Attempt 1] {cmd}\n")
                f.write(f"{cmd}\n")
                
                # 第 2 遍执行（检查重试）
                if not args.research_only:
                    target_file = f"output\\{(idx+1):03d}\\strategy.json"
                    f.write(f"if not exist \"{target_file}\" (\n")
                    f.write(f"    echo =======================================\n")
                    f.write(f"    echo [Retry] target.json missing, retrying once...\n")
                    f.write(f"    echo =======================================\n")
                    f.write(f"    {cmd}\n")
                    f.write(f") else (\n")
                    f.write(f"    echo [Verified] {target_file} generated successfully.\n")
                    f.write(f")\n")
                else:
                    # 如果是 resaerch-only 则较难在 batch 中校验几个文件，默认让它安全重试一遍
                    f.write(f"echo [Retry For Research] Automatically trying once more to ensure completeness...\n")
                    f.write(f"{cmd}\n")
                
                f.write("timeout /t 2 > NUL\n")
            f.write("echo Worker Group %d finished!\n" % worker_id)
            f.write("pause\n")
            
        # Open a new CMD window executing this batch file
        subprocess.Popen(
            ["cmd.exe", "/c", "start", f"Worker_{worker_id}", "cmd.exe", "/c", str(bat_path.absolute())],
            cwd=str(base_dir)
        )
        print(f"Launched window for Worker {worker_id} -> Indices: {indices}")

if __name__ == "__main__":
    main()