/*
 * Copyright (c) 2026 Imre Vadász <imre@vdsz.com>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

/*
 * Device driver for Intel's C-state residency counters.
 * Supports Nehalem and later Core CPUs.
 * Currently supports models up to CoffeeLake generation.
 */

#include <sys/param.h>
#include <sys/bus.h>
#include <sys/types.h>
#include <sys/malloc.h>
#include <sys/systm.h>
#include <sys/module.h>
#include <sys/conf.h>
#include <sys/cpu_topology.h>
#include <sys/kernel.h>
#include <sys/thread2.h>
#include <sys/sensors.h>
#include <sys/bitops.h>

#include <machine/clock.h>
#include <machine/specialreg.h>
#include <machine/cpufunc.h>
#include <machine/cputypes.h>
#include <machine/md_var.h>

#include "cpu_if.h"

#define ATOM_MSR_PKG_CSTATE_COUNTER_HZ  (1000 * 1000)
#define ATOM_MSR_PKG_C2_RESIDENCY       MSR_PKG_C3_RESIDENCY
#define ATOM_MSR_PKG_C4_RESIDENCY       MSR_PKG_C6_RESIDENCY
#define ATOM_MSR_PKG_C6_RESIDENCY       MSR_PKG_C7_RESIDENCY

#define ATOM_MSR_CORE_C1_RESIDENCY      0x660
#define ATOM_MSR_CORE_C6_RESIDENCY      MSR_CORE_C6_RESIDENCY

struct corecstat_sensor {
        uint64_t        tsc_count;
        uint64_t        cst_count;
        u_int           msr;
        int             bits;
        uint64_t        mask;
        struct ksensor  sensor;
};

struct corecstat_softc {
        device_t                sc_dev;

        int                     sc_is_atom;
        /*
         * sc_is_atom == 0
         *   0 ==> Nehalem/Westmere
         *   1 ==> Sandy Bridge / Ivy Bridge / Haswell / Broadwell / Skylake
         *   2 ==> Haswell/Broadwell - low-power
         *
         * sc_is_atom == 1
         *   0 ==> Original Atom
         *   1 ==> Airmont Atom
         */
        int                     sc_version;

        struct corecstat_sensor sc_core_sens[3];
        u_int                   sc_core_sens_cnt;

#define PC2_IDX  0
#define PC3_IDX  1
#define PC6_IDX  2
#define PC7_IDX  3
#define PC8_IDX  4
#define PC9_IDX  5
#define PC10_IDX 6
        struct corecstat_sensor *sc_pkg_sens;
        u_int                   sc_pkg_sens_cnt;

        struct sensor_task      *sc_senstask;
        struct ksensordev       *sc_cpu_sensordev;

        struct ksensordev       *sc_sensordev;
};

/*
 * Device methods.
 */
static void     corecstat_identify(driver_t *driver, device_t parent);
static int      corecstat_probe(device_t dev);
static int      corecstat_attach(device_t dev);
static int      corecstat_detach(device_t dev);
static void     corecstat_init_atom_pkg(struct corecstat_softc *sc, int cpu);
static void     corecstat_init_atom_cpu(struct corecstat_softc *sc, int cpu);
static void     corecstat_init_core_pkg(struct corecstat_softc *sc, int cpu);
static void     corecstat_init_core_cpu(struct corecstat_softc *sc, int cpu);
static void     corecstat_refresh(void *arg);
static void     corecstat_pkg_sens_init(struct corecstat_softc *sc,
                                        int cpu, char *desc, u_int msr,
                                        int bits);
static void     corecstat_core_sens_init(struct corecstat_softc *sc,
                                        int cpu, char *desc, u_int msr,
                                        int bits);
static void     corecstat_sens_init(struct corecstat_sensor *sens,
                                    char *desc, u_int msr, int bits);
static void     corecstat_sens_update(struct corecstat_softc *sc,
                                      struct corecstat_sensor *sens);
static int      corecstat_try(u_int msr, char *name);

static device_method_t corecstat_methods[] = {
        /* Device interface */
        DEVMETHOD(device_identify,      corecstat_identify),
        DEVMETHOD(device_probe,         corecstat_probe),
        DEVMETHOD(device_attach,        corecstat_attach),
        DEVMETHOD(device_detach,        corecstat_detach),

        DEVMETHOD_END
};

static driver_t corecstat_driver = {
        "corecstat",
        corecstat_methods,
        sizeof(struct corecstat_softc),
};

static devclass_t corecstat_devclass;
DRIVER_MODULE(corecstat, cpu, corecstat_driver, corecstat_devclass, NULL, NULL);
MODULE_VERSION(corecstat, 1);

static void
corecstat_identify(driver_t *driver, device_t parent)
{
        device_t child;

        /* Make sure we're not being doubly invoked. */
        if (device_find_child(parent, "corecstat", -1) != NULL)
                return;

        /* Check that the vendor is Intel. */
        if (cpu_vendor_id != CPU_VENDOR_INTEL)
                return;

        /*
         * We add a child for each CPU since settings must be performed
         * on each CPU in the SMP case.
         */
        child = device_add_child(parent, "corecstat", -1);
        if (child == NULL)
                device_printf(parent, "add corecstat child failed\n");
}

static int
corecstat_probe(device_t dev)
{
        int cpu_family, cpu_model;

        if (resource_disabled("corecstat", 0))
                return (ENXIO);

        cpu_model = CPUID_TO_MODEL(cpu_id);
        cpu_family = CPUID_TO_FAMILY(cpu_id);

        if (cpu_family == 0x06) {
                switch (cpu_model) {
                /* ATOM */
                case 0x1c:
                case 0x26:
                case 0x27:
                case 0x35:
                case 0x36:
                        break;
                /* Silvermont Atom */
                case 0x37:
                case 0x4a:
                case 0x4c: /* Airmont Atom */
                case 0x4d:
                case 0x5a:
                case 0x5d:
                        break;
                /* Nehalem */
                case 0x1a:
                case 0x1e:
                case 0x1f:
                case 0x25: /* Westmere */
                case 0x2c: /* Westmere */
                case 0x2e:
                case 0x2f: /* Westmere */
                        break;
                /* Sandy Bridge */
                case 0x2a:
                case 0x2d:
                case 0x3a: /* Ivy Bridge */
                case 0x3e: /* Ivy Bridge-E */
                        break;
                /* Haswell */
                case 0x3c:
                case 0x3f: /* Haswell-E */
                case 0x45:
                case 0x46:
                        break;
                /* Broadwell */
                case 0x3d:
                case 0x47:
                case 0x4f:
                case 0x56:
                        break;
                /* Skylake */
                case 0x4e:
                case 0x5e:
                /* Kabylake/Coffeelake */
                case 0x8e:
                        break;
                default:
                        return (ENXIO);
                }
        }

        device_set_desc(dev, "CPU Package C-State residency counters");

        return (BUS_PROBE_GENERIC);
}

static int
corecstat_attach(device_t dev)
{
        struct corecstat_softc *sc = device_get_softc(dev);
        const struct cpu_node *node;
        int cpu = device_get_unit(device_get_parent(dev));
        int cpu_family, cpu_model;

        sc->sc_dev = dev;
        sc->sc_is_atom = 0;
        sc->sc_core_sens_cnt = 0;
        sc->sc_pkg_sens_cnt = 0;

        cpu_model = CPUID_TO_MODEL(cpu_id);
        cpu_family = CPUID_TO_FAMILY(cpu_id);

        if (cpu_family == 0x06) {
                switch (cpu_model) {
                /* ATOM */
                case 0x1c:
                case 0x26:
                case 0x27:
                case 0x35:
                case 0x36:
                        /* Only Package C-States available, no Core C-states */
                        sc->sc_is_atom = 1;
                        sc->sc_version = 0;
                        break;
                /* Silvermont Atom */
                case 0x37:
                case 0x4a:
                case 0x4c: /* Airmont Atom */
                case 0x4d:
                case 0x5a:
                case 0x5d:
                        sc->sc_is_atom = 1;
                        sc->sc_version = 1;
                        break;
                /* Nehalem */
                case 0x1a:
                case 0x1e:
                case 0x1f:
                case 0x25: /* Westmere */
                case 0x2c: /* Westmere */
                case 0x2e:
                case 0x2f: /* Westmere */
                        /* PKG_C3, PKG_C6, PKG_C7 */
                        /* CORE_C3, CORE_C6 */
                        sc->sc_version = 0;
                        break;
                /* Sandy Bridge */
                case 0x2a:
                case 0x2d:
                /* Ivy Bridge */
                case 0x3a:
                case 0x3e: /* Ivy Bridge-E */
                /* Haswell */
                case 0x3c:
                case 0x3f: /* Haswell-E */
                case 0x46:
                /* Broadwell */
                case 0x47:
                case 0x4f:
                case 0x56:
                /* Skylake */
                case 0x4e:
                case 0x5e:
                        /* PKG_C2, PKG_C3, PKG_C6, PKG_C7 */
                        /* CORE_C3, CORE_C6, CORE_C7 */
                        sc->sc_version = 1;
                        break;
                case 0x45:      /* Haswell low-power CPUs */
                case 0x3d:      /* Broadwell low-power CPUs */
                case 0x8e:      /* Kabylake/Coffeelake CPUs */
                        /* PKG_C2, PKG_C3, PKG_C6, PKG_C7, PKG_C8, PKG_C9, PKG_C10 */
                        /* CORE_C3, CORE_C6, CORE_C7 */
                        sc->sc_version = 2;
                        break;
                default:
                        return (ENXIO);
                }
        }

        sc->sc_cpu_sensordev = CPU_GET_SENSDEV(devclass_find_unit("cpu", cpu));
        if (sc->sc_cpu_sensordev == NULL)
                return (1);

        node = get_cpu_node_by_cpuid(cpu);
        while (node != NULL) {
                if (node->type == CORE_LEVEL &&
                    !(cpu == BSRCPUMASK(node->members)))
                        return (ENXIO);
                if (node->type == CHIP_LEVEL) {
                        if (node->child_no == 0)
                                node = NULL;
                        break;
                }
                node = node->parent_node;
        }
        if (node != NULL && cpu == BSRCPUMASK(node->members)) {
                sc->sc_sensordev = kmalloc(sizeof(struct ksensordev),
                    M_DEVBUF, M_WAITOK | M_ZERO);

                ksnprintf(sc->sc_sensordev->xname, sizeof(
                    sc->sc_sensordev->xname), "cpu_node%d", get_chip_ID(cpu));

                if (sc->sc_is_atom)
                        corecstat_init_atom_pkg(sc, cpu);
                else
                        corecstat_init_core_pkg(sc, cpu);

                sensordev_install(sc->sc_sensordev);
        } else if (sc->sc_is_atom && sc->sc_version == 0) {
                return (ENXIO);
        }

        if (sc->sc_is_atom)
                corecstat_init_atom_cpu(sc, cpu);
        else
                corecstat_init_core_cpu(sc, cpu);

        sc->sc_senstask = sensor_task_register2(sc, corecstat_refresh, 1,
            device_get_unit(dev));

        return (0);
}

static int
corecstat_detach(device_t dev)
{
        struct corecstat_softc *sc = device_get_softc(dev);

        if (sc->sc_senstask != NULL) {
                sensor_task_unregister2(sc->sc_senstask);
                sc->sc_senstask = NULL;
        }
        if (sc->sc_cpu_sensordev != NULL) {
                int i;

                for (i = 0; i < sc->sc_core_sens_cnt; i++) {
                        sensor_detach(sc->sc_cpu_sensordev,
                            &sc->sc_core_sens[i].sensor);
                }
        }
        if (sc->sc_sensordev != NULL) {
                sensordev_deinstall(sc->sc_sensordev);
                kfree(sc->sc_sensordev, M_DEVBUF);
                sc->sc_sensordev = NULL;
        }
        if (sc->sc_pkg_sens != NULL) {
                kfree(sc->sc_pkg_sens, M_DEVBUF);
                sc->sc_pkg_sens = NULL;
        }

        return (0);
}

static void
corecstat_init_atom_pkg(struct corecstat_softc *sc, int cpu)
{
        if (sc->sc_version == 0) {
                sc->sc_pkg_sens = kmalloc(3 * sizeof(struct corecstat_sensor),
                    M_DEVBUF, M_WAITOK | M_ZERO);
                corecstat_pkg_sens_init(sc, cpu, "node%d PC2 residency",
                    ATOM_MSR_PKG_C2_RESIDENCY, 64);
                corecstat_pkg_sens_init(sc, cpu, "node%d PC4 residency",
                    ATOM_MSR_PKG_C4_RESIDENCY, 64);
                corecstat_pkg_sens_init(sc, cpu, "node%d PC6 residency",
                    ATOM_MSR_PKG_C6_RESIDENCY, 64);
        } else {
                sc->sc_pkg_sens = kmalloc(1 * sizeof(struct corecstat_sensor),
                    M_DEVBUF, M_WAITOK | M_ZERO);
                corecstat_pkg_sens_init(sc, cpu, "node%d PC6 residency",
                    ATOM_MSR_PKG_C6_RESIDENCY, 64);
        }
}

static void
corecstat_init_atom_cpu(struct corecstat_softc *sc, int cpu)
{
        /* No Core C-state residency counters on old Atoms */
        KKASSERT(sc->sc_version != 0);

        corecstat_core_sens_init(sc, cpu, "node%d core%d C1 res.",
            ATOM_MSR_CORE_C1_RESIDENCY, 64);
        corecstat_core_sens_init(sc, cpu, "node%d core%d C6 res.",
            ATOM_MSR_CORE_C6_RESIDENCY, 64);
}

static void
corecstat_init_core_pkg(struct corecstat_softc *sc, int cpu)
{
        sc->sc_pkg_sens = kmalloc(7 * sizeof(struct corecstat_sensor),
            M_DEVBUF, M_WAITOK | M_ZERO);
        if (sc->sc_version == 1 || sc->sc_version == 2) {
                corecstat_pkg_sens_init(sc, cpu, "node%d PC2 residency",
                    MSR_PKG_C2_RESIDENCY, 64);
        }

        corecstat_pkg_sens_init(sc, cpu, "node%d PC3 residency",
            MSR_PKG_C3_RESIDENCY, 64);
        corecstat_pkg_sens_init(sc, cpu, "node%d PC6 residency",
            MSR_PKG_C6_RESIDENCY, 64);
        corecstat_pkg_sens_init(sc, cpu, "node%d PC7 residency",
            MSR_PKG_C7_RESIDENCY, 64);

        if (sc->sc_version == 2) {
                corecstat_pkg_sens_init(sc, cpu, "node%d PC8 residency",
                    MSR_PKG_C8_RESIDENCY, 60);
                corecstat_pkg_sens_init(sc, cpu, "node%d PC9 residency",
                    MSR_PKG_C9_RESIDENCY, 60);
                corecstat_pkg_sens_init(sc, cpu, "node%d PC10 residency",
                    MSR_PKG_C10_RESIDENCY, 60);
        }
}

static void
corecstat_init_core_cpu(struct corecstat_softc *sc, int cpu)
{
        corecstat_core_sens_init(sc, cpu, "node%d core%d C3 res.",
            MSR_CORE_C3_RESIDENCY, 64);
        corecstat_core_sens_init(sc, cpu, "node%d core%d C6 res.",
            MSR_CORE_C6_RESIDENCY, 64);

        if (sc->sc_version == 1 || sc->sc_version == 2) {
                corecstat_core_sens_init(sc, cpu, "node%d core%d C7 res.",
                    MSR_CORE_C7_RESIDENCY, 64);
        }
}

static void
corecstat_refresh(void *arg)
{
        struct corecstat_softc *sc = (struct corecstat_softc *)arg;
        int i;

        for (i = 0; i < sc->sc_core_sens_cnt; i++)
                corecstat_sens_update(sc, &sc->sc_core_sens[i]);
        for (i = 0; i < sc->sc_pkg_sens_cnt; i++)
                corecstat_sens_update(sc, &sc->sc_pkg_sens[i]);
}

static void
corecstat_core_sens_init(struct corecstat_softc *sc,
    int cpu, char *desc, u_int msr, int bits)
{
        struct corecstat_sensor *sens = &sc->sc_core_sens[sc->sc_core_sens_cnt];
        char buf[sizeof(sens->sensor.desc)];

        ksnprintf(buf, sizeof(buf), desc, get_chip_ID(cpu),
            get_core_number_within_chip(cpu));
        if (corecstat_try(msr, buf)) {
                sc->sc_core_sens_cnt++;
                corecstat_sens_init(sens, buf, msr, bits);
                sensor_attach(sc->sc_cpu_sensordev, &sens->sensor);
        }
}

static void
corecstat_pkg_sens_init(struct corecstat_softc *sc,
    int cpu, char *desc, u_int msr, int bits)
{
        struct corecstat_sensor *sens = &sc->sc_pkg_sens[sc->sc_pkg_sens_cnt];
        char buf[sizeof(sens->sensor.desc)];

        ksnprintf(buf, sizeof(buf), desc, get_chip_ID(cpu));
        if (corecstat_try(msr, buf)) {
                sc->sc_pkg_sens_cnt++;
                corecstat_sens_init(sens, buf, msr, bits);
                sensor_attach(sc->sc_sensordev, &sens->sensor);
        }
}

static void
corecstat_sens_init(struct corecstat_sensor *sens, char *desc, u_int msr,
    int bits)
{
        uint64_t a, b;

        KKASSERT(bits > 0 && bits <= 64);
        sens->sensor.type = SENSOR_PERCENT;
        sens->sensor.flags |= SENSOR_FUNKNOWN;
        sens->msr = msr;
        sens->bits = bits;
        sens->mask = __BITS64(0, bits-1);

        crit_enter();
        a = rdtsc();
        b = rdmsr(msr);
        crit_exit();
        sens->tsc_count = a;
        sens->cst_count = b & sens->mask;

        strlcpy(sens->sensor.desc, desc, sizeof(sens->sensor.desc));
}

static void
corecstat_sens_update(struct corecstat_softc *sc,
    struct corecstat_sensor *sens)
{
        uint64_t a, b, msr;
        uint64_t tscdiff, cstdiff;

        msr = sens->msr;
        crit_enter();
        a = rdtsc();
        b = rdmsr(msr);
        crit_exit();

        b &= sens->mask;

        if (a < sens->tsc_count)
                tscdiff = a + (~sens->tsc_count);
        else
                tscdiff = a - sens->tsc_count;

        if (b < sens->cst_count)
                cstdiff = b + ((~sens->cst_count) & sens->mask);
        else
                cstdiff = b - sens->cst_count;

        sens->tsc_count = a;
        sens->cst_count = b;

        /*
         * On old Atom CPUs (before Silvermont) the Package-cstate counters
         * run at 1 MHz.
         */
        if (sc->sc_is_atom && sc->sc_version == 0) {
                if (cstdiff < (1ULL << 44)) {
                        cstdiff *= ATOM_MSR_PKG_CSTATE_COUNTER_HZ;
                        cstdiff /= tsc_frequency;
                } else {
                        cstdiff /= tsc_frequency;
                        cstdiff *= ATOM_MSR_PKG_CSTATE_COUNTER_HZ;
                }
        }

        if (tscdiff > 0) {
                /* Make sure we don't calculate total garbage here */
                if (cstdiff < (1ULL << 44)) {
                        sens->sensor.value = (cstdiff * 100000) / tscdiff;
                        sens->sensor.flags &=
                            ~(SENSOR_FUNKNOWN | SENSOR_FINVALID);
                } else {
                        sens->sensor.flags |= SENSOR_FINVALID;
                }
        }
}

static int
corecstat_try(u_int msr, char *name)
{
        uint64_t val;

        if (rdmsr_safe(msr, &val) != 0) {
                kprintf("msr %s (0x%08x) not available\n", name, msr);
                return 0;
        }
        return 1;
}