import argparse
import json
import subprocess
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import pandas as pd


def load_factory_codes(year: int) -> list[str]:
    csv_path = Path("data") / str(year) / "factory_data.csv"
    encodings = ("utf-8", "utf-8-sig", "gbk")
    last_error = None
    df = None
    for enc in encodings:
        try:
            df = pd.read_csv(csv_path, encoding=enc)
            break
        except UnicodeDecodeError as exc:
            last_error = exc
            continue
    if df is None:
        raise last_error if last_error else FileNotFoundError(f"Missing {csv_path}")
    col_map = {c.strip(): c for c in df.columns}
    if "工厂英文名" not in col_map:
        raise ValueError(f"{csv_path} 缺少字段: 工厂英文名")
    return df[col_map["工厂英文名"]].astype(str).str.strip().tolist()


def run_single(factory_code: str, python_exe: str, script_path: Path) -> tuple[str, int, str]:
    cmd = [python_exe, str(script_path), "--factory", factory_code]
    result = subprocess.run(cmd, capture_output=True, text=True)
    output = (result.stdout or "") + (result.stderr or "")
    return factory_code, result.returncode, output


def main():
    parser = argparse.ArgumentParser(description="Parallel GA calibration for all factories.")
    parser.add_argument("--max-workers", type=int, default=6, help="并行进程数上限（默认 6）")
    args = parser.parse_args()

    year = json.load(open("year.json", "r", encoding="utf-8"))["year"]
    codes = load_factory_codes(year)
    script_path = Path("ga_calibration_by_one_factory.py")
    if not script_path.exists():
        raise FileNotFoundError(f"找不到 {script_path}")

    python_exe = sys.executable
    results = []
    with ProcessPoolExecutor(max_workers=args.max_workers) as executor:
        future_map = {
            executor.submit(run_single, code, python_exe, script_path): code for code in codes
        }
        for future in as_completed(future_map):
            code = future_map[future]
            try:
                factory, rc, out = future.result()
            except Exception as exc:
                results.append((code, 1, f"Exception: {exc}"))
                continue
            results.append((factory, rc, out))
            status = "OK" if rc == 0 else f"FAIL({rc})"
            print(f"[{status}] {factory}")

    print("\n详细输出：")
    for factory, rc, out in results:
        status = "OK" if rc == 0 else f"FAIL({rc})"
        print(f"--- {factory} [{status}] ---")
        print(out.strip())
        print("---------------")


if __name__ == "__main__":
    main()