import os
import torch
import json
import multiprocessing as mp

from glob import glob
try:
    from natsort import natsorted
except ModuleNotFoundError:
    # Model shards use zero-padded numeric names, so lexical order is equivalent.
    natsorted = sorted

from worker import worker
from utils import (
    load_quant_config,
    copy_files,
    get_ignore,
    process_and_save_info,
)


def main(
    input_dir,
    output_dir,
    model_type,
    mode,
    num_workers,
    custom_dequant_config,
    custom_quant_config,
    retry_files
):
    print(f"Make output_dir: {output_dir}")
    os.makedirs(output_dir, exist_ok=True)

    print(f"Copy files from {input_dir} to {output_dir}")
    copy_files(input_dir, output_dir)

    safetensor_files = natsorted(glob(os.path.join(input_dir, "*.safetensors")))
    if retry_files is not None:
        safetensor_files_ = []
        for sf in safetensor_files:
            for rf in retry_files:
                if sf.endswith(rf):
                    safetensor_files_.append(sf)
        safetensor_files = safetensor_files_
    assert len(safetensor_files) > 0

    num_gpus = torch.cuda.device_count()
    num_workers = min(num_workers, min(num_gpus, len(safetensor_files)))
    assert num_workers > 0, "No CUDA devices found or num_workers incorrect"

    print(f"{num_workers} GPUs will be used")
    print(f"Found {len(safetensor_files)} safetensor files")

    # ============================
    # 关键：统一使用 spawn context
    # ============================
    ctx = mp.get_context("spawn")

    task_queue = ctx.Queue()
    manager = ctx.Manager()
    global_lock = ctx.Lock()

    weight_map_shared = manager.dict()
    quant_map_shared = manager.dict()

    for sf in safetensor_files:
        # TODO: add a filter for specified files
        task_queue.put(sf)
    try:
        quant_config = load_quant_config(input_dir)
    except Exception:
        quant_config = None

    processes = []

    index_file = os.path.join(input_dir, "model.safetensors.index.json")
    tensor_index = None
    if os.path.exists(index_file):
        print("Loading tensor index")
        with open(index_file) as f:
            tensor_index = json.load(f)["weight_map"]

    for gpu_id in range(num_workers):
        p = ctx.Process(
            target=worker,
            args=(
                global_lock,
                task_queue,
                gpu_id,
                weight_map_shared,
                quant_map_shared,
                output_dir,
                mode,
                model_type,
                quant_config,
                custom_dequant_config,
                custom_quant_config,
                tensor_index
            ),
        )
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    output_model = output_dir.split(os.sep)[-1]
    qms_output_dir = os.path.join("online_quant_configs", output_model)
    os.makedirs(qms_output_dir, exist_ok=True)

    process_and_save_info(
        input_dir,
        output_dir,
        qms_output_dir,
        dict(weight_map_shared),
        mode in ("quant", "dequant+quant"),
        custom_quant_config.get("symmetric"),
        ignore=None if mode == "dequant" else get_ignore(custom_quant_config),
        group_size=custom_quant_config.get("group_size", -1),
    )

    if not("dequant" in mode and "+quant" not in mode):
        with open(os.path.join(qms_output_dir, "quant_map.json"), "w") as f:
            json.dump(dict(quant_map_shared), f, indent=2)


if __name__ == "__main__":
    import yaml
    from argparse import ArgumentParser

    mp.set_start_method("spawn", force=True)

    parser = ArgumentParser()
    parser.add_argument("--config", type=str)
    args = parser.parse_args()

    with open(args.config, "r") as f:
        cfg = yaml.safe_load(f)

    main(
        cfg['paths']['checkpoint_dir'],
        cfg['paths']['output_dir'],
        cfg['model']['model_type'],
        cfg['runtime']['mode'],
        cfg['runtime']['num_workers'],
        cfg['dequant'],
        cfg['quant'],
        cfg['retry_files']
    )