import argparse
import base64
import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

from openai import OpenAI


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
os.chdir(ROOT)

from config import api_key, base_url


DEFAULT_PROMPT = "improve this infographics to make it look more visual coherent and appealing."


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-root",
        default=None,
        help="chart_template_samples 下某次任务目录；默认取最新时间戳目录",
    )
    parser.add_argument(
        "--output-dir",
        default=None,
        help="输出目录；默认 <input-root>/gpt_image_2_improved",
    )
    parser.add_argument("--model", default="gpt-image-2")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--size", default="auto")
    parser.add_argument("--quality", default="auto")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def latest_sample_run(root: Path):
    runs = [
        p
        for p in root.iterdir()
        if p.is_dir() and (p / "run_config.json").is_file() and p.name[:8].isdigit()
    ]
    if not runs:
        raise SystemExit(f"no timestamp run found under {root}")
    return sorted(runs, key=lambda p: p.name)[-1]


def collect_pngs(input_root: Path, output_dir: Path, limit: int):
    pngs = [
        p
        for p in input_root.rglob("*.png")
        if output_dir not in p.parents
    ]
    pngs = sorted(pngs)
    if limit > 0:
        pngs = pngs[:limit]
    return pngs


def output_path_for(input_root: Path, output_dir: Path, image_path: Path):
    rel = image_path.relative_to(input_root)
    return output_dir / rel.parent / f"{image_path.stem}_gpt_image_2.png"


def write_jsonl(path: Path, record: dict):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


def improve_one(client, image_path: Path, out_path: Path, args):
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(image_path, "rb") as image_file:
        result = client.images.edit(
            model=args.model,
            image=image_file,
            prompt=args.prompt,
            size=args.size,
            quality=args.quality,
            n=1,
        )
    image_bytes = base64.b64decode(result.data[0].b64_json)
    with open(out_path, "wb") as f:
        f.write(image_bytes)
    return {
        "input": str(image_path),
        "output": str(out_path),
        "model": args.model,
        "prompt": args.prompt,
    }


def improve_job(job):
    args, image_path, out_path = job
    client = OpenAI(api_key=api_key, base_url=base_url)
    return improve_one(client, image_path, out_path, args)


def main():
    args = parse_args()
    input_root = Path(args.input_root) if args.input_root else latest_sample_run(
        ROOT / "output" / "chart_template_samples"
    )
    output_dir = Path(args.output_dir) if args.output_dir else input_root / "gpt_image_2_improved"
    output_dir.mkdir(parents=True, exist_ok=True)

    images = collect_pngs(input_root, output_dir, args.limit)
    manifest = output_dir / "manifest.jsonl"
    if manifest.exists() and args.overwrite:
        manifest.unlink()

    jobs = []
    for image_path in images:
        out_path = output_path_for(input_root, output_dir, image_path)
        if out_path.exists() and not args.overwrite:
            continue
        jobs.append((args, image_path, out_path))

    print(f"input_root={input_root}", flush=True)
    print(f"output_dir={output_dir}", flush=True)
    print(f"images={len(images)} pending={len(jobs)} workers={args.workers}", flush=True)

    if args.workers > 1:
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = [executor.submit(improve_job, job) for job in jobs]
            for index, future in enumerate(as_completed(futures), 1):
                record = future.result()
                write_jsonl(manifest, record)
                print(f"[{index}/{len(jobs)}] {record['output']}", flush=True)
    else:
        client = OpenAI(api_key=api_key, base_url=base_url)
        for index, (_args, image_path, out_path) in enumerate(jobs, 1):
            record = improve_one(client, image_path, out_path, args)
            write_jsonl(manifest, record)
            print(f"[{index}/{len(jobs)}] {record['output']}", flush=True)

    summary = {
        "input_root": str(input_root),
        "output_dir": str(output_dir),
        "model": args.model,
        "prompt": args.prompt,
        "total_images": len(images),
        "submitted_images": len(jobs),
    }
    with open(output_dir / "summary.json", "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)


if __name__ == "__main__":
    main()