#define _GNU_SOURCE
#include <sched.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <ctype.h>
#include <sys/types.h>

#include "pluma-cpu-affinity.h"
#include "pluma-debug.h"

static cpu_set_t original_affinity;
static gboolean affinity_saved = FALSE;

#define CPUINFO_PATH "/proc/cpuinfo"
#define VENDOR_LINE "vendor_id"
#define VENDOR1 "centaurhauls"
#define VENDOR2 "shanghai"

gboolean
pluma_cpu_bind_to_core (gint core_id)
{
    cpu_set_t cpu_set;
    gint num_cores;
    pid_t pid = getpid();  // 获取当前进程ID，而不是线程

    // 获取系统CPU核心数
    num_cores = sysconf(_SC_NPROCESSORS_ONLN);

    // 检查core_id是否有效
    if (core_id < 0 || core_id >= num_cores)
    {
        g_warning("Invalid core_id %d. System has %d cores (0-%d)",
                  core_id, num_cores, num_cores - 1);
        return FALSE;
    }

    // 保存当前CPU亲和性
    if (sched_getaffinity(pid, sizeof(original_affinity), &original_affinity) == 0)
    {
        affinity_saved = TRUE;
    }
    else
    {
        g_warning("Failed to get current CPU affinity: %s", g_strerror(errno));
    }

    // 设置CPU亲和性到指定核心
    CPU_ZERO(&cpu_set);
    CPU_SET(core_id, &cpu_set);

    if (sched_setaffinity(pid, sizeof(cpu_set), &cpu_set) == 0)
    {
        g_debug("Successfully bound process %d to CPU core %d", pid, core_id);
        return TRUE;
    }
    else
    {
        g_warning("Failed to bind process %d to CPU core %d: %s",
                  pid, core_id, g_strerror(errno));
        return FALSE;
    }
}

gboolean
pluma_cpu_unbind (void)
{
    pid_t pid = getpid();  // 获取当前进程ID，而不是线程

    if (affinity_saved)
    {
        if (sched_setaffinity(pid, sizeof(original_affinity), &original_affinity) == 0)
        {
            g_debug("Successfully restored original CPU affinity for process %d", pid);
            affinity_saved = FALSE;
            return TRUE;
        }
        else
        {
            g_warning("Failed to restore original CPU affinity for process %d: %s",
                      pid, g_strerror(errno));
            return FALSE;
        }
    }

    return TRUE;
}

gboolean
pluma_is_zhaoxin_cpu(void)
{
	if (check_zx() == 1)
	{
		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

int check_zx()
{
	FILE *fp = fopen(CPUINFO_PATH, "r");
	if (!fp)
	{
		return 0;
	}

	char buffer[128];
	while (fgets(buffer, sizeof(buffer), fp))
	{
		if (strncmp(buffer, VENDOR_LINE, 9) == 0)
		{
			for (char *p = buffer; *p; ++p)
			{
				*p = tolower(*p);
			}
			if (strstr(buffer, VENDOR1) || strstr(buffer, VENDOR2))
			{
				fclose(fp);
				return 1;
			}
			else
			{
				break;
			}
		}
	}
	fclose(fp);
	return 0;
} 