import os, sys
import time
import subprocess
import platform
import json

os.environ["BLAS_NUM_THREADS"] = "1"
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "build"))
from kt_kernel import kt_kernel_ext
from kt_kernel_ext.kvcache import ggml_type
import torch
from torch import inf, nn
from torch.nn import init

from tqdm import tqdm

qlen = 4096
kvlen = 0
page_table = list(range(20))
page_size = 256
pages_count = 200


hidden_size = 7168
num_heads = 128
kv_lora_rank = 512
q_lora_rank = 512
nope_size = 128
rope_size = 64
page_size = 512
layer_num = 10


rope_theta = 10000
max_qlen = qlen + kvlen
max_kvlen = 4096
max_position_embeddings = 163840

rope_scaling = {
    "beta_fast": 32,
    "beta_slow": 1,
    "factor": 40,
    "mscale": 1.0,
    "mscale_all_dim": 1.0,
    "original_max_position_embeddings": 4096,
    "type": "yarn",
}

CPUINFER_PARAM = 304
# 初始化 CPUInfer（此处使用原始构造函数，可根据需要调整配置参数）
CPUInfer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM)


warm_up_iter = 20
test_iter = 100


# 获取脚本相关信息，用于生成结果保存文件名
script_path = os.path.abspath(__file__)
script_dir = os.path.dirname(script_path)
script_name = os.path.splitext(os.path.basename(script_path))[0]
json_path = os.path.join(script_dir, "bench_results " + ".jsonl")


def get_git_commit():
    """
    获取当前 git 提交记录（commit hash 和提交信息），
    并检查是否存在未提交的更改（dirty）
    """
    result = {}
    try:
        commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
        commit_msg = subprocess.check_output(["git", "log", "-1", "--pretty=%B"]).decode("utf-8").strip()
        result["commit"] = commit
        result["commit_message"] = commit_msg

        # 检查是否存在未提交的更改
        dirty_output = subprocess.check_output(["git", "status", "--porcelain"]).decode("utf-8").strip()
        if dirty_output:
            result["dirty"] = True
            result["dirty_files"] = dirty_output.splitlines()
        else:
            result["dirty"] = False
    except Exception as e:
        result["commit"] = None
        result["commit_message"] = None
        result["dirty"] = None
        result["error"] = str(e)
    return result


def get_system_info():
    """
    获取系统信息，包括系统名称、CPU 型号、内存大小（GB）、CPU 核数及 socket 数量
    """
    info = {}
    uname = platform.uname()
    info["system_name"] = uname.system
    info["node_name"] = uname.node

    # 获取 CPU 型号（仅 Linux 支持）
    cpu_model = None
    if os.path.exists("/proc/cpuinfo"):
        try:
            with open("/proc/cpuinfo", "r") as f:
                for line in f:
                    if "model name" in line:
                        cpu_model = line.split(":", 1)[1].strip()
                        break
        except Exception as e:
            cpu_model = f"Error: {e}"
    info["cpu_model"] = cpu_model

    # 获取内存大小（单位：GB），仅 Linux 支持
    mem_total_gb = None
    if os.path.exists("/proc/meminfo"):
        try:
            with open("/proc/meminfo", "r") as f:
                for line in f:
                    if "MemTotal" in line:
                        mem_kb = float(line.split(":", 1)[1].split()[0])
                        mem_total_gb = round(mem_kb / (1024 * 1024), 2)
                        break
        except Exception as e:
            mem_total_gb = f"Error: {e}"
    info["memory_size_GB"] = mem_total_gb

    info["cpu_core_count"] = os.cpu_count()

    # 解析 /proc/cpuinfo 获取 socket 数量
    sockets = set()
    if os.path.exists("/proc/cpuinfo"):
        try:
            with open("/proc/cpuinfo", "r") as f:
                for line in f:
                    if "physical id" in line:
                        sockets.add(line.split(":", 1)[1].strip())
        except Exception as e:
            sockets = set()
    info["cpu_socket_count"] = len(sockets) if len(sockets) > 0 else 1

    return info


def record_results(result, filename=json_path):
    """
    将结果以 JSON 格式追加到文件中
    """
    with open(filename, "a") as f:
        f.write(json.dumps(result) + "\n")


def bench_mla(quant_mode: str):
    """
    测试 MLA 模型的性能
    """
    with torch.inference_mode():
        # 这里可以添加 MLA 模型的具体实现和测试代码
        hidden_type = 1  # ggml_type::GGML_TYPE_FP16（固定）
        if quant_mode == "fp32":
            q_a_proj_type = 0  # ggml_type::GGML_TYPE_F32
            q_b_proj_type = 0
            kv_a_proj_with_mqa_type = 0
            kv_b_proj_type = 0
            w_o_type = 0
            bytes_per_elem = 4.000000
        elif quant_mode == "fp16":
            q_a_proj_type = 1  # ggml_type::GGML_TYPE_F32
            q_b_proj_type = 1
            kv_a_proj_with_mqa_type = 1
            kv_b_proj_type = 1
            w_o_type = 1
            bytes_per_elem = 2.000000
        elif quant_mode == "q4_k_m":
            q_a_proj_type = 12  # ggml_type::GGML_TYPE_Q4_K
            q_b_proj_type = 12
            kv_a_proj_with_mqa_type = 12  # ggml_type::GGML_TYPE_Q6_K
            kv_b_proj_type = 12
            w_o_type = 12
            bytes_per_elem = 0.5625
        else:
            raise ValueError("不支持的量化模式")

        # 构建各层 MLA 模型的输入数据
        mlas = []
        for i in tqdm(range(layer_num)):
            q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False, dtype=torch.float16)
            q_b_proj = nn.Linear(q_lora_rank, num_heads * (nope_size + rope_size), bias=False, dtype=torch.float16)
            kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank + rope_size, bias=False, dtype=torch.float16)
            kv_b_proj = nn.Linear(num_heads * (nope_size + nope_size), kv_lora_rank, bias=False, dtype=torch.float16)
            o_proj = nn.Linear(num_heads * nope_size, hidden_size, bias=False, dtype=torch.float16)

            init.normal_(q_a_proj.weight, mean=0.0, std=0.02)
            init.normal_(q_b_proj.weight, mean=0.0, std=0.02)
            init.normal_(kv_a_proj_with_mqa.weight, mean=0.0, std=0.02)
            init.normal_(kv_b_proj.weight, mean=0.0, std=0.02)
            init.normal_(o_proj.weight, mean=0.0, std=0.02)
            q_a_proj_weight = q_a_proj.weight.to(torch.float16).to("cpu").contiguous()
            q_b_proj_weight = q_b_proj.weight.to(torch.float16).to("cpu").contiguous()
            kv_a_proj_with_mqa_weight = kv_a_proj_with_mqa.weight.to("cpu").to(torch.float16).contiguous()
            kv_b_proj_weight = kv_b_proj.weight.to(torch.float16).to("cpu").contiguous()
            o_proj_weight = o_proj.weight.to(torch.float16).to("cpu").contiguous()

            config = kt_kernel_ext.mla.MLAConfig(
                hidden_size,
                q_lora_rank,
                kv_lora_rank,
                num_heads,
                nope_size,
                rope_size,
            )
            config.max_qlen = max_qlen
            config.max_kvlen = max_kvlen
            config.max_position_embeddings = max_position_embeddings
            config.rope_scaling_factor = rope_scaling["factor"]
            config.rope_theta = rope_theta
            config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
            config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
            config.rope_scaling_mscale = rope_scaling["mscale"]
            config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
            config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]

            config.q_a_proj = q_a_proj_weight.data_ptr()
            config.q_b_proj = q_b_proj_weight.data_ptr()
            config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
            config.kv_b_proj = kv_b_proj_weight.data_ptr()
            config.o_proj = o_proj_weight.data_ptr()

            config.q_a_proj_type = ggml_type.FP16
            config.q_b_proj_type = ggml_type.FP16
            config.kv_a_proj_with_mqa_type = ggml_type.FP16
            config.kv_b_proj_type = ggml_type.FP16
            config.w_o_type = ggml_type.FP16

            config.pool = CPUInfer.backend_

            mla = kt_kernel_ext.mla.MLA(config)
            mla.load_weights()
            mla.set_local_pages(pages_count)
            mlas.append(mla)

        print("Generating data...")
        input_tensor = (
            torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cpu").to("cpu").contiguous()
        )
        output_tensor = (
            torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cpu").to("cpu").contiguous()
        )

        print("Warming up...")

        for i in tqdm(range(warm_up_iter)):
            mlas[i % layer_num].forward(
                [qlen],
                [page_table],
                [kvlen],
                input_tensor[i % layer_num].data_ptr(),
                output_tensor[i % layer_num].data_ptr(),
            )

        print("Start testing...")

        start = time.perf_counter()
        for i in tqdm(range(test_iter)):
            mlas[i % layer_num].forward(
                [qlen],
                [page_table],
                [kvlen],
                input_tensor[i % layer_num].data_ptr(),
                output_tensor[i % layer_num].data_ptr(),
            )

        end = time.perf_counter()
        total_time = end - start

        time_per_iter_us = (total_time * 1e6) / test_iter
        bandwidth = (
            bytes_per_elem
            * (
                q_lora_rank * hidden_size
                + (kv_lora_rank + rope_size) * hidden_size
                + (nope_size + rope_size) * q_lora_rank * num_heads
                + (nope_size + nope_size) * kv_lora_rank * num_heads
                + hidden_size * nope_size * num_heads
                + hidden_size * qlen
            )
            * test_iter
            / (total_time * 1e9)
        )
        flops = (
            2
            * (
                q_lora_rank * hidden_size * qlen
                + kv_lora_rank * hidden_size * qlen
                + num_heads * (nope_size + rope_size) * q_lora_rank * qlen
                + num_heads * qlen * nope_size * kv_lora_rank
                + num_heads * (kvlen + qlen) * kv_lora_rank * qlen
                + num_heads * rope_size * qlen * (qlen + kvlen)
                + num_heads * kv_lora_rank * (qlen + kvlen) * qlen
                + num_heads * nope_size * kv_lora_rank * qlen
                + hidden_size * num_heads * nope_size * qlen
            )
            * test_iter
            / (total_time * 1e12)
        )

        print("Quant mode:", quant_mode)
        print("Time(s):", total_time)
        print("Iteration:", test_iter)
        print("Time(us) per iteration:", time_per_iter_us)
        print("Bandwidth:", bandwidth, "GB/s")
        print("TFLOPS:", flops)
        print("")

        # 整理测试结果
        result = {
            "test_name": os.path.basename(__file__),
            "quant_mode": quant_mode,
            "total_time_seconds": total_time,
            "iterations": test_iter,
            "time_per_iteration_us": time_per_iter_us,
            "bandwidth_GBs": bandwidth,
            "flops_TFLOPS": flops,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
            "test_parameters": {
                "qlen": qlen,
                "kvlen": kvlen,
                "page_table": page_table,
                "page_size": page_size,
                "pages_count": pages_count,
                "hidden_size": hidden_size,
                "num_heads": num_heads,
                "kv_lora_rank": kv_lora_rank,
                "q_lora_rank": q_lora_rank,
                "nope_size": nope_size,
                "rope_size": rope_size,
                "layer_num": layer_num,
                "rope_theta": rope_theta,
                "max_qlen": max_qlen,
                "max_kvlen": max_kvlen,
                "max_position_embeddings": max_position_embeddings,
                "rope_scaling": rope_scaling,
                "warm_up_iter": warm_up_iter,
                "test_iter": test_iter,
                "CPUInfer_parameter": CPUINFER_PARAM,
            },
        }
        # 添加 git 与系统信息
        result.update(get_git_commit())
        result.update(get_system_info())
        # 将结果记录到 JSON 文件中
        print(result)
        record_results(result)


bench_mla("fp16")
