#include <Python.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <linux/i2c.h>
#include <linux/i2c-dev.h>
#include <errno.h>
#include <time.h>
#include <pthread.h>
#include <stdlib.h> // For malloc and free

typedef struct SampleData {
    uint16_t value;
    struct timespec timestamp;
} SampleData;

typedef struct {
    PyObject_HEAD
    int file;
    uint8_t addr;
    pthread_t thread_id;          // 线程 ID
    int sample_interval_us;       // 采样间隔（微秒）
    int sample_count;             // 采样次数
    volatile int samples_ready;   // 样本准备就绪标志
    volatile int stop_sampling;   // 停止采样标志
    union {
        SampleData *data_buffer; // 动态分配的数据缓冲区，包含值和时间戳
    } buffer;
    int buffer_index;             // 缓冲区索引
} I2CHandle;

// 设置线程的CPU亲和性
static int set_cpu_affinity(pthread_t thread, int cpu) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu, &cpuset);

    if (pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset) != 0) {
        perror("pthread_setaffinity_np");
        return -1;
    }

    return 0;
}
// 读取16位寄存器的函数
static uint16_t read_reg(int file, uint8_t i2c_addr, uint8_t reg) {
    uint8_t buf[2];
    struct i2c_rdwr_ioctl_data packets;
    struct i2c_msg messages[2];

    messages[0].addr  = i2c_addr;  // INA226的I2C地址
    messages[0].flags = 0;
    messages[0].len   = 1;
    messages[0].buf   = &reg;

    messages[1].addr  = i2c_addr;
    messages[1].flags = I2C_M_RD;
    messages[1].len   = 2;
    messages[1].buf   = buf;

    packets.msgs  = messages;
    packets.nmsgs = 2;

    if (ioctl(file, I2C_RDWR, &packets) < 0) {
        perror("Failed to read from the i2c bus");
        return 0;
    }

    return (buf[0] << 8 | buf[1]);
}

// 写入16位寄存器的函数
static void write_reg(int file, uint8_t i2c_addr, uint8_t reg, uint16_t value) {
    uint8_t buf[3];
    struct i2c_rdwr_ioctl_data packets;
    struct i2c_msg messages[1];

    buf[0] = reg;
    buf[1] = (value >> 8) & 0xFF;
    buf[2] = value & 0xFF;

    messages[0].addr  = i2c_addr;  // INA226的I2C地址
    messages[0].flags = 0;
    messages[0].len   = 3;
    messages[0].buf   = buf;

    packets.msgs  = messages;
    packets.nmsgs = 1;

    if (ioctl(file, I2C_RDWR, &packets) < 0) {
        perror("Failed to write to the i2c bus");
    }
}

static void *sampling_thread(void *arg) {
    I2CHandle *self = (I2CHandle *)arg;

    struct timespec start_ts;
    struct timespec stop_ts;
    int count = 0;
    while (count < self->sample_count && !self->stop_sampling) {
        clock_gettime(CLOCK_REALTIME, &start_ts); // 记录当前时间
        uint16_t data = read_reg(self->file, self->addr, 0x01);  // 假设寄存器地址为 0x00
        self->buffer.data_buffer[self->buffer_index].value = data;
        self->buffer.data_buffer[self->buffer_index].timestamp = start_ts;
        self->buffer_index++;
        count++;
        clock_gettime(CLOCK_REALTIME, &stop_ts); // 
        double current_time = start_ts.tv_sec + start_ts.tv_nsec / 1e9;
        double complete_time = stop_ts.tv_sec + stop_ts.tv_nsec / 1e9;
        int sleep_us = self->sample_interval_us - (int)((complete_time - current_time) * 1e6);
        if(sleep_us > 0 && sleep_us < self->sample_interval_us)
        {
            usleep(sleep_us); // 等待0.1毫秒后再次检查时间
        } 
    }
    // 设置采样完成标志
    self->samples_ready = 1;
    self->stop_sampling = 1; // 停止采样
    pthread_exit(NULL);
}

static void ina226_close(void *self) {
    I2CHandle *handle = (I2CHandle *)self;
    if (handle->file > 0) {
        close(handle->file);
        handle->file = -1;
    }
    if (handle->thread_id != 0) {
        handle->stop_sampling = 1; // 设置停止采样标志
        pthread_join(handle->thread_id, NULL);
        handle->thread_id = 0;
    }
    free(handle->buffer.data_buffer); // 释放动态分配的内存
    handle->buffer.data_buffer = NULL;
}

static PyObject* ina226_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
    I2CHandle *self;
    self = (I2CHandle *)type->tp_alloc(type, 0);
    if (self != NULL) {
        self->file = -1;
        self->addr = 0;
        self->thread_id = 0;
        self->sample_interval_us = 0;
        self->sample_count = 0;
        self->samples_ready = 0;
        self->stop_sampling = 0;
        self->buffer_index = 0;
        self->buffer.data_buffer = NULL; // 初始化为 NULL
    }
    return (PyObject *)self;
}

static int ina226_init(I2CHandle *self, PyObject *args, PyObject *kwds) {
    static char *kwlist[] = {"address", NULL};
    uint8_t i2c_addr;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "B", kwlist, &i2c_addr)) {
        return -1;
    }

    self->file = open("/dev/i2c-3", O_RDWR);
    if (self->file < 0) {
        PyErr_SetFromErrno(PyExc_RuntimeError);
        return -1;
    }

    if (ioctl(self->file, I2C_SLAVE, i2c_addr) < 0) {
        close(self->file);
        PyErr_SetString(PyExc_RuntimeError, "Failed to acquire bus access and/or talk to slave.");
        return -1;
    }

    self->addr = i2c_addr;
    return 0;
}

static PyObject* ina226_write_reg(I2CHandle *self, PyObject *args) {
    uint8_t reg;
    uint16_t value;

    if (!PyArg_ParseTuple(args, "BH", &reg, &value)) {
        return NULL;
    }
    write_reg(self->file, self->addr, reg, value);
    Py_RETURN_NONE;
}

static PyObject* ina226_read_reg(I2CHandle *self, PyObject *args) {
    uint8_t reg;
    uint16_t data = 0;  // 初始化 data 为 0

    if (!PyArg_ParseTuple(args, "B", &reg)) {
        return NULL;
    }
    data = read_reg(self->file, self->addr, reg);

    return Py_BuildValue("h", data);
}

static PyObject* ina226_start_sampling(I2CHandle *self, PyObject *args) {
    int sample_interval_us;
    int sample_count;

    if (!PyArg_ParseTuple(args, "ii", &sample_interval_us, &sample_count)) {
        return NULL;
    }

    // 设置采样间隔和采样次数
    self->sample_interval_us = sample_interval_us;
    self->sample_count = sample_count;
    self->buffer_index = 0;  // 重置缓冲区索引

    // 分配数据缓冲区
    self->buffer.data_buffer = (SampleData*)malloc(sample_count * sizeof(SampleData));

    if (self->buffer.data_buffer == NULL) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to allocate memory for data buffer.");
        return NULL;
    }

    // 启动采样线程
    self->samples_ready = 0;
    self->stop_sampling = 0;
    if (pthread_create(&self->thread_id, NULL, sampling_thread, (void *)self) != 0) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to create thread.");
        free(self->buffer.data_buffer); // 释放内存
        self->buffer.data_buffer = NULL;
        return NULL;
    }

    if (set_cpu_affinity(self->thread_id, 3) != 0) {
        return NULL;
    } 

    Py_RETURN_NONE;
}

static PyObject* ina226_is_sampling_complete(I2CHandle *self, PyObject *unused) {
    // __attribute__((unused)) PyObject *dummy = unused;
    return PyBool_FromLong(self->samples_ready);
}

static PyObject* ina226_get_data(I2CHandle *self, PyObject *unused) {
    // __attribute__((unused)) PyObject *dummy = unused;
    PyObject *list = PyList_New(0);
    int index = 0;
    for (; index < self->buffer_index; index++) {
        PyObject *sub_tuple = PyTuple_New(2);
        double current_time = self->buffer.data_buffer[index].timestamp.tv_sec + self->buffer.data_buffer[index].timestamp.tv_nsec / 1e9;
        PyTuple_SetItem(sub_tuple, 0, Py_BuildValue("h", self->buffer.data_buffer[index].value));
        PyTuple_SetItem(sub_tuple, 1, Py_BuildValue("d", current_time));
        PyList_Append(list, sub_tuple);
    }

    return list;
}

static PyMethodDef I2CHandle_methods[] = {
    {"write_reg", (PyCFunction)ina226_write_reg, METH_VARARGS, "Write to an I2C register"},
    {"read_reg", (PyCFunction)ina226_read_reg, METH_VARARGS, "Read from an I2C register"},
    {"start_sampling", (PyCFunction)ina226_start_sampling, METH_VARARGS, "Start sampling data at specified intervals"},
    {"is_sampling_complete", (PyCFunction)ina226_is_sampling_complete, METH_NOARGS, "Check if sampling is complete"},
    {"get_data", (PyCFunction)ina226_get_data, METH_NOARGS, "Get sampled data with timestamps"},
    {NULL, NULL, 0, NULL}
};

static PyTypeObject I2CHandleType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "I2CHandle",                   /* tp_name */
    sizeof(I2CHandle),             /* tp_basicsize */
    0,                             /* tp_itemsize */
    (destructor)ina226_close,      /* tp_dealloc */
    0,                             /* tp_print */
    0,                             /* tp_getattr */
    0,                             /* tp_setattr */
    0,                             /* tp_reserved */
    0,                             /* tp_repr */
    0,                             /* tp_as_number */
    0,                             /* tp_as_sequence */
    0,                             /* tp_as_mapping */
    0,                             /* tp_hash */
    0,                             /* tp_call */
    0,                             /* tp_str */
    0,                             /* tp_getattro */
    0,                             /* tp_setattro */
    0,                             /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT,            /* tp_flags */
    "I2C Device Handle",           /* tp_doc */
    0,                             /* tp_traverse */
    0,                             /* tp_clear */
    0,                             /* tp_richcompare */
    0,                             /* tp_weaklistoffset */
    0,                             /* tp_iter */
    0,                             /* tp_iternext */
    I2CHandle_methods,             /* tp_methods */
    0,                             /* tp_members */
    0,                             /* tp_getset */
    0,                             /* tp_base */
    0,                             /* tp_dict */
    0,                             /* tp_descr_get */
    0,                             /* tp_descr_set */
    0,                             /* tp_dictoffset */
    (initproc)ina226_init,         /* tp_init */
    0,                             /* tp_alloc */
    ina226_new,                    /* tp_new */
};

// 定义 PyModuleDef 结构体
static struct PyModuleDef ina226module = {
    PyModuleDef_HEAD_INIT,
    "wtina226",                      /* m_name */
    "INA226 I2C communication module", /* m_doc */
    -1,                            /* m_size */
    I2CHandle_methods,             /* m_methods */
    NULL,                          /* m_slots */
    NULL,                          /* m_traverse */
    NULL,                          /* m_clear */
    NULL,                          /* m_free */
};

PyMODINIT_FUNC PyInit_wtina226(void) {
    PyObject *m;

    // 初始化类型
    if (PyType_Ready(&I2CHandleType) < 0)
        return NULL;

    // 创建模块
    m = PyModule_Create(&ina226module);
    if (m == NULL)
        return NULL;

    // 将类型添加到模块
    if (PyModule_AddObject(m, "I2CHandle", (PyObject *)&I2CHandleType) < 0) {
        Py_DECREF(&I2CHandleType);
        Py_DECREF(m);
        return NULL;
    }

    return m;
}
