import os
import json
import torch
import multiprocessing as mp
import traceback

from safetensors.torch import save_file
import random
import time
import fcntl
import shutil


def save_weights(save_path, output):
    tmp_path = save_path + ".tmp"
    save_file(output, tmp_path)

    with open(tmp_path, "rb") as f:
        os.fsync(f.fileno())
    os.replace(tmp_path, save_path)


# TODO: 分开加载，不覆盖

def worker(
    global_lock,
    task_queue: mp.Queue,
    gpu_id: int,
    weight_map_shared,
    quant_map_shared,
    output_dir: str,
    mode: str,  # dequant, quant, dequant+quant
    model_type: str | None = None,
    quant_config: dict | None = None,  # from original config.json
    custom_dequant_config: dict | None = None,  # from .yml
    custom_quant_config: dict | None = None,  # from .yml
    tensor_index: dict | None = None,
    input_dir: str | None = None
):
    try:
        torch.cuda.set_device(gpu_id)
        device = torch.device(f"cuda:{gpu_id}")
        print(f"[Worker {gpu_id}] Started on {device}")

        from functions.dequant import dequant
        from functions.quant import quant

        while True:
            try:
                sf = task_queue.get(timeout=5)
            except Exception:
                print(f"[Worker {gpu_id}] Task queue empty, exiting")
                break

            try:
                print(f"[Worker {gpu_id}] Processing {sf}")

                output = None

                if "dequant" in mode:
                    output = dequant(
                        file_name=sf,
                        quant_config=quant_config,
                        custom_dequant_config=custom_dequant_config,
                        device=device,
                        model_type=model_type,
                        tensor_index=tensor_index,
                        input_dir=os.path.dirname(sf)
                    )

                    if "+quant" in mode:
                        output, quant_map = quant(
                            state_dict=output,
                            custom_quant_config=custom_quant_config,
                            device=device,
                            model_type=model_type
                        )
                else:
                    if "quant" in mode:
                        output, quant_map = quant(
                            file_name=sf,
                            custom_quant_config=custom_quant_config,
                            device=device,
                            model_type=model_type
                        )
                    else:
                        raise ValueError(f"only support quant, dequant, and dequant+quant, get {mode}")

                file_name = os.path.basename(sf)

                for wname, _ in output.items():
                    weight_map_shared[wname] = file_name
                
                try:
                    for wname, dtype in quant_map.items():
                        quant_map_shared[wname] = dtype
                except Exception as e:
                    pass

                save_path = os.path.join(output_dir, file_name)
                sleep_time = random.uniform(0, 2)
                time.sleep(sleep_time)

                with global_lock:
                    print(f"[Worker {gpu_id}] Saving {file_name}")
                    save_weights(save_path, output)

                # 主动释放显存
                del output
                torch.cuda.empty_cache()

            except Exception as e:
                print(f"[Worker {gpu_id}] Error processing {sf}")
                traceback.print_exc()

    except Exception:
        traceback.print_exc()