/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright (c) 2012 DEY Storage Systems, Inc.  All rights reserved.
 * Copyright 2012 Nexenta Systems, Inc.  All rights reserved.
 * Copyright 2019 Joyent, Inc.
 * Copyright 2022 Oxide Computer Co.
 */

/*
 * This implements psrinfo(8), a utility to report various information
 * about processors, cores, and threads (virtual cpus).  This is mostly
 * intended for human consumption - this utility doesn't do much more than
 * simply process kstats for human readability.
 *
 * All the relevant kstats are in the cpu_info kstat module.
 */

#include <sys/sysmacros.h>

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <kstat.h>
#include <libintl.h>
#include <locale.h>
#include <libgen.h>
#include <ctype.h>
#include <errno.h>
#include <err.h>

#include <libdevinfo.h>

#define _(x)    gettext(x)
#if XGETTEXT
        /* These CPU states are here for benefit of xgettext */
        _("on-line")
        _("off-line")
        _("faulted")
        _("powered-off")
        _("no-intr")
        _("spare")
        _("unknown")
        _("disabled")
#endif

/*
 * We deal with sorted linked lists, where the sort key is usually the
 * cpu id, core id, or chip id.  We generalize this with simple node.
 */
struct link {
        long            l_id;
        struct link     *l_next;
        void            *l_ptr;
};

/*
 * A physical chip.  A chip can contain multiple cores and virtual cpus.
 */
struct pchip {
        struct link     p_link;
        int             p_ncore;
        int             p_nvcpu;
        struct link     *p_cores;
        struct link     *p_vcpus;
        int             p_doit;
};

struct core {
        struct link     c_link;
        struct link     c_link_pchip;

        int             c_nvcpu;
        int             c_doit;

        struct pchip    *c_pchip;
        struct link     *c_vcpus;
};

struct vcpu {
        struct link     v_link;

        struct link     v_link_core;
        struct link     v_link_pchip;

        int             v_doit;

        struct pchip    *v_pchip;
        struct core     *v_core;

        char            *v_state;
        long            v_state_begin;
        char            *v_cpu_type;
        char            *v_fpu_type;
        long            v_clock_mhz;
        long            v_pchip_id;     /* 1 per socket */
        char            *v_impl;
        char            *v_brand;
        char            *v_socket;
        long            v_core_id;      /* n per chip_id */
};

static struct link *pchips = NULL;
static struct link *cores = NULL;
static struct link *vcpus = NULL;

static uint_t nr_cpus;
static uint_t nr_cores;
static uint_t nr_chips;

static const char *cmdname;

static void
usage(char *msg)
{
        if (msg != NULL)
                (void) fprintf(stderr, "%s: %s\n", cmdname, msg);
        (void) fprintf(stderr, _("usage: \n"
            "\t%s -r propname\n"
            "\t%s [-v] [-p] [processor_id ...]\n"
            "\t%s -s [-p] processor_id\n"
            "\t%s -t [-S <state> | -c | -p]\n"),
            cmdname, cmdname, cmdname, cmdname);
        exit(2);
}

/* like perror, but includes the command name */
static void
die(const char *msg)
{
        (void) fprintf(stderr, "%s: %s: %s\n", cmdname, msg, strerror(errno));
        exit(2);
}

static char *
mystrdup(const char *src)
{
        char *dst;

        if ((dst = strdup(src)) == NULL)
                die(_("strdup() failed"));
        return (dst);
}

static void *
zalloc(size_t size)
{
        void *ptr;

        if ((ptr = calloc(1, size)) == NULL)
                die(_("calloc() failed"));
        return (ptr);
}

/*
 * Insert a new node on a list, at the insertion point given.
 */
static void
ins_link(struct link **ins, struct link *item)
{
        item->l_next = *ins;
        *ins = item;
}

/*
 * Find an id on a sorted list.  If the requested id is not found,
 * then the insertpt will be set (if not null) to the location where
 * a new node should be inserted with ins_link (see above).
 */
static void *
find_link(void *list, int id, struct link ***insertpt)
{
        struct link **ins = list;
        struct link *l;

        while ((l = *ins) != NULL) {
                if (l->l_id == id)
                        return (l->l_ptr);
                if (l->l_id > id)
                        break;
                ins = &l->l_next;
        }
        if (insertpt != NULL)
                *insertpt = ins;
        return (NULL);
}

/*
 * Print the linked list of ids in parens, taking care to collapse
 * ranges, so instead of (0 1 2 3) it should print (0-3).
 */
static void
print_links(struct link *l)
{
        int     start = -1;
        int     end = 0;

        (void) printf(" (");
        while (l != NULL) {
                if (start < 0) {
                        start = l->l_id;
                }
                end = l->l_id;
                if ((l->l_next == NULL) ||
                    (l->l_next->l_id > (l->l_id + 1))) {
                        /* end of the contiguous group */
                        if (start == end) {
                                (void) printf("%d", start);
                        } else {
                                (void) printf("%d-%d", start, end);
                        }
                        if (l->l_next)
                                (void) printf(" ");
                        start = -1;
                }
                l = l->l_next;
        }
        (void) printf(")");
}

static const char *
timestr(long t)
{
        static char buffer[256];
        (void) strftime(buffer, sizeof (buffer), _("%m/%d/%Y %T"),
            localtime(&t));
        return (buffer);
}

static void
print_vp(int nspec)
{
        struct pchip *chip;
        struct core *core;
        struct vcpu *vcpu;
        struct link *l1, *l2;
        int len;
        for (l1 = pchips; l1; l1 = l1->l_next) {

                chip = l1->l_ptr;

                if ((nspec != 0) && (chip->p_doit == 0))
                        continue;

                vcpu = chip->p_vcpus->l_ptr;

                /*
                 * Note that some of the way these strings are broken up are
                 * to accommodate the legacy translations so that we won't
                 * have to retranslate for this utility.
                 */
                if ((chip->p_ncore == 1) || (chip->p_ncore == chip->p_nvcpu)) {
                        (void) printf(_("%s has %d virtual %s"),
                            _("The physical processor"),
                            chip->p_nvcpu,
                            chip->p_nvcpu > 1 ?
                            _("processors") :
                            _("processor"));
                } else {
                        (void) printf(_("%s has %d %s and %d virtual %s"),
                            _("The physical processor"),
                            chip->p_ncore, _("cores"),
                            chip->p_nvcpu,
                            chip->p_nvcpu > 1 ?
                            _("processors") : _("processor"));
                }

                print_links(chip->p_vcpus);
                (void) putchar('\n');

                if ((chip->p_ncore == 1) || (chip->p_ncore == chip->p_nvcpu)) {
                        if (strlen(vcpu->v_impl)) {
                                (void) printf("  %s\n", vcpu->v_impl);
                        }
                        if (((len = strlen(vcpu->v_brand)) != 0) &&
                            (strncmp(vcpu->v_brand, vcpu->v_impl, len) != 0))
                                (void) printf("\t%s", vcpu->v_brand);
                } else {
                        for (l2 = chip->p_cores; l2; l2 = l2->l_next) {
                                core = l2->l_ptr;
                                (void) printf(_("  %s has %d virtual %s"),
                                    _("The core"),
                                    core->c_nvcpu,
                                    chip->p_nvcpu > 1 ?
                                    _("processors") : _("processor"));
                                print_links(core->c_vcpus);
                                (void) putchar('\n');
                        }
                        if (strlen(vcpu->v_impl)) {
                                (void) printf("    %s\n", vcpu->v_impl);
                        }
                        if (((len = strlen(vcpu->v_brand)) != 0) &&
                            (strncmp(vcpu->v_brand, vcpu->v_impl, len) != 0))
                                (void) printf("      %s", vcpu->v_brand);
                }
                if (strcmp(vcpu->v_socket, "Unknown") != 0) {
                        (void) printf("\t[ %s: %s ]", _("Socket"),
                            vcpu->v_socket);
                }
                (void) putchar('\n');
        }
}

static void
print_ps(void)
{
        int online = 1;
        struct pchip *p = NULL;
        struct vcpu *v;
        struct link *l;

        /*
         * Report "1" if all cpus colocated on the same chip are online.
         */
        for (l = pchips; l != NULL; l = l->l_next) {
                p = l->l_ptr;
                if (p->p_doit)
                        break;
        }
        if (p == NULL)
                return; /* should never happen! */
        for (l = p->p_vcpus; l != NULL; l = l->l_next) {
                v = l->l_ptr;
                if (strcmp(v->v_state, "on-line") != 0) {
                        online = 0;
                        break;
                }
        }

        (void) printf("%d\n", online);
}

static void
print_s(void)
{
        struct link *l;

        /*
         * Find the processor (there will be only one) that we selected,
         * and report whether or not it is online.
         */
        for (l = vcpus; l != NULL; l = l->l_next) {
                struct vcpu *v = l->l_ptr;
                if (v->v_doit) {
                        (void) printf("%d\n",
                            strcmp(v->v_state, "on-line") == 0 ? 1 : 0);
                        return;
                }
        }
}

static void
print_p(int nspec)
{
        struct          link *l1, *l2;
        int             online = 0;

        /*
         * Print the number of physical packages with at least one processor
         * online.
         */
        for (l1 = pchips; l1 != NULL; l1 = l1->l_next) {
                struct pchip *p = l1->l_ptr;
                if ((nspec == 0) || (p->p_doit)) {

                        for (l2 = p->p_vcpus; l2 != NULL; l2 = l2->l_next) {
                                struct vcpu *v = l2->l_ptr;
                                if (strcmp(v->v_state, "on-line") == 0) {
                                        online++;
                                        break;
                                }
                        }
                }
        }
        (void) printf("%d\n", online);
}

static void
print_v(int nspec)
{
        struct link     *l;

        for (l = vcpus; l != NULL; l = l->l_next) {
                struct vcpu *v = l->l_ptr;

                if ((nspec != 0) && (!v->v_doit))
                        continue;
                (void) printf(_("Status of virtual processor %d as of: "),
                    l->l_id);
                (void) printf("%s\n", timestr(time(NULL)));
                (void) printf(_("  %s since %s.\n"),
                    _(v->v_state), timestr(v->v_state_begin));
                if (v->v_clock_mhz) {
                        (void) printf(
                            _("  The %s processor operates at %llu MHz,\n"),
                            v->v_cpu_type, (unsigned long long)v->v_clock_mhz);
                } else {
                        (void) printf(
                            _("  The %s processor operates at " \
                            "an unknown frequency,\n"), v->v_cpu_type);
                }
                switch (*v->v_fpu_type) {
                case '\0':
                        (void) printf(
                            _("\tand has no floating point processor.\n"));
                        break;
                case 'a': case 'A':
                case 'e': case 'E':
                case 'i': case 'I':
                case 'o': case 'O':
                case 'u': case 'U':
                case 'y': case 'Y':
                        (void) printf(
                            _("\tand has an %s floating point processor.\n"),
                            v->v_fpu_type);
                        break;
                default:
                        (void) printf(
                            _("\tand has a %s floating point processor.\n"),
                            v->v_fpu_type);
                        break;
                }
        }
}

static void
print_normal(int nspec)
{
        struct link     *l;
        struct vcpu     *v;

        for (l = vcpus; l != NULL; l = l->l_next) {
                v = l->l_ptr;
                if ((nspec == 0) || (v->v_doit)) {
                        (void) printf(_("%d\t%-8s  since %s\n"),
                            l->l_id, _(v->v_state), timestr(v->v_state_begin));
                }
        }
}

static bool
valid_propname(const char *propname)
{
        size_t i;

        const char *props[] = {
                "smt_enabled",
        };

        for (i = 0; i < ARRAY_SIZE(props); i++) {
                if (strcmp(propname, props[i]) == 0)
                        break;
        }

        return (i != ARRAY_SIZE(props));
}

static void
read_property(const char *propname)
{
        di_prop_t prop = DI_PROP_NIL;
        di_node_t root_node;
        bool show_all = strcmp(propname, "all") == 0;

        if (!show_all && !valid_propname(propname))
                errx(EXIT_FAILURE, _("unknown CPU property %s"), propname);

        if ((root_node = di_init("/", DINFOPROP)) == NULL)
                err(EXIT_FAILURE, _("failed to read root node"));

        while ((prop = di_prop_sys_next(root_node, prop)) != DI_PROP_NIL) {
                const char *name = di_prop_name(prop);
                char *val;
                int nr_vals;

                if (!valid_propname(name))
                        continue;

                if (!show_all && strcmp(di_prop_name(prop), propname) != 0)
                        continue;

                if ((nr_vals = di_prop_strings(prop, &val)) < 1) {
                        err(EXIT_FAILURE,
                            _("error reading property %s"), name);
                } else if (nr_vals != 1) {
                        errx(EXIT_FAILURE, _("invalid property %s"), name);
                }

                printf("%s=%s\n", name, val);

                if (!show_all)
                        exit(EXIT_SUCCESS);
        }

        if (!show_all)
                errx(EXIT_FAILURE, _("property %s was not found"), propname);

        di_fini(root_node);
}

static void
print_total(int opt_c, int opt_p, const char *opt_S)
{
        uint_t count = 0;

        if (opt_c) {
                printf("%u\n", nr_cores);
                return;
        } else if (opt_p) {
                printf("%u\n", nr_chips);
                return;
        } else if (opt_S == NULL || strcmp(opt_S, "all") == 0) {
                printf("%u\n", nr_cpus);
                return;
        }


        for (struct link *l = vcpus; l != NULL; l = l->l_next) {
                struct vcpu *v = l->l_ptr;
                if (strcmp(opt_S, v->v_state) == 0)
                        count++;
        }

        printf("%u\n", count);
}

int
main(int argc, char **argv)
{
        kstat_ctl_t     *kc;
        kstat_t         *ksp;
        kstat_named_t   *knp;
        struct vcpu     *vc;
        struct core     *core;
        struct pchip    *chip;
        struct link     **ins;
        char            *s;
        int             nspec;
        int             optc;
        int             opt_c = 0;
        int             opt_p = 0;
        const char      *opt_r = NULL;
        const char      *opt_S = NULL;
        int             opt_s = 0;
        int             opt_t = 0;
        int             opt_v = 0;
        int             ex = 0;

        cmdname = basename(argv[0]);


        (void) setlocale(LC_ALL, "");
#if !defined(TEXT_DOMAIN)
#define TEXT_DOMAIN     "SYS_TEST"
#endif
        (void) textdomain(TEXT_DOMAIN);

        /* collect the kstats */
        if ((kc = kstat_open()) == NULL)
                die(_("kstat_open() failed"));

        if ((ksp = kstat_lookup(kc, "cpu_info", -1, NULL)) == NULL)
                die(_("kstat_lookup() failed"));

        for (ksp = kc->kc_chain; ksp; ksp = ksp->ks_next) {

                if (strcmp(ksp->ks_module, "cpu_info") != 0)
                        continue;
                if (kstat_read(kc, ksp, NULL) == -1)
                        die(_("kstat_read() failed"));

                vc = find_link(&vcpus, ksp->ks_instance, &ins);
                if (vc == NULL) {
                        vc = zalloc(sizeof (struct vcpu));
                        vc->v_link.l_id = ksp->ks_instance;
                        vc->v_link_core.l_id = ksp->ks_instance;
                        vc->v_link_pchip.l_id = ksp->ks_instance;
                        vc->v_link.l_ptr = vc;
                        vc->v_link_core.l_ptr = vc;
                        vc->v_link_pchip.l_ptr = vc;
                        ins_link(ins, &vc->v_link);
                        nr_cpus++;
                }

                if ((knp = kstat_data_lookup(ksp, "state")) != NULL) {
                        vc->v_state = mystrdup(knp->value.c);
                } else {
                        vc->v_state = "unknown";
                }

                if ((knp = kstat_data_lookup(ksp, "cpu_type")) != NULL) {
                        vc->v_cpu_type = mystrdup(knp->value.c);
                }
                if ((knp = kstat_data_lookup(ksp, "fpu_type")) != NULL) {
                        vc->v_fpu_type = mystrdup(knp->value.c);
                }

                if ((knp = kstat_data_lookup(ksp, "state_begin")) != NULL) {
                        vc->v_state_begin = knp->value.l;
                }

                if ((knp = kstat_data_lookup(ksp, "clock_MHz")) != NULL) {
                        vc->v_clock_mhz = knp->value.l;
                }

                if ((knp = kstat_data_lookup(ksp, "brand")) == NULL) {
                        vc->v_brand = _("(unknown)");
                } else {
                        vc->v_brand = mystrdup(knp->value.str.addr.ptr);
                }

                if ((knp = kstat_data_lookup(ksp, "socket_type")) == NULL) {
                        vc->v_socket = "Unknown";
                } else {
                        vc->v_socket = mystrdup(knp->value.str.addr.ptr);
                }

                if ((knp = kstat_data_lookup(ksp, "implementation")) == NULL) {
                        vc->v_impl = _("(unknown)");
                } else {
                        vc->v_impl = mystrdup(knp->value.str.addr.ptr);
                }
                /*
                 * Legacy code removed the chipid and cpuid fields... we
                 * do the same for compatibility.  Note that the original
                 * pattern is a bit strange, and we have to emulate this because
                 * on SPARC we *do* emit these.  The original pattern we are
                 * emulating is: $impl =~ s/(cpuid|chipid)\s*\w+\s+//;
                 */
                if ((s = strstr(vc->v_impl, "chipid")) != NULL) {
                        char *x = s + strlen("chipid");
                        while (isspace(*x))
                                x++;
                        if ((!isalnum(*x)) && (*x != '_'))
                                goto nochipid;
                        while (isalnum(*x) || (*x == '_'))
                                x++;
                        if (!isspace(*x))
                                goto nochipid;
                        while (isspace(*x))
                                x++;
                        (void) strcpy(s, x);
                }
nochipid:
                if ((s = strstr(vc->v_impl, "cpuid")) != NULL) {
                        char *x = s + strlen("cpuid");
                        while (isspace(*x))
                                x++;
                        if ((!isalnum(*x)) && (*x != '_'))
                                goto nocpuid;
                        while (isalnum(*x) || (*x == '_'))
                                x++;
                        if (!isspace(*x))
                                goto nocpuid;
                        while (isspace(*x))
                                x++;
                        (void) strcpy(s, x);
                }
nocpuid:

                if ((knp = kstat_data_lookup(ksp, "chip_id")) != NULL)
                        vc->v_pchip_id = knp->value.l;
                chip = find_link(&pchips, vc->v_pchip_id, &ins);
                if (chip == NULL) {
                        chip = zalloc(sizeof (struct pchip));
                        chip->p_link.l_id = vc->v_pchip_id;
                        chip->p_link.l_ptr = chip;
                        ins_link(ins, &chip->p_link);
                        nr_chips++;
                }
                vc->v_pchip = chip;

                if ((knp = kstat_data_lookup(ksp, "core_id")) != NULL)
                        vc->v_core_id = knp->value.l;
                core = find_link(&cores, vc->v_core_id, &ins);
                if (core == NULL) {
                        core = zalloc(sizeof (struct core));
                        core->c_link.l_id = vc->v_core_id;
                        core->c_link.l_ptr = core;
                        core->c_link_pchip.l_id = vc->v_core_id;
                        core->c_link_pchip.l_ptr = core;
                        core->c_pchip = chip;
                        ins_link(ins, &core->c_link);
                        chip->p_ncore++;
                        (void) find_link(&chip->p_cores, core->c_link.l_id,
                            &ins);
                        ins_link(ins, &core->c_link_pchip);
                        nr_cores++;
                }
                vc->v_core = core;



                /* now put other linkages in place */
                (void) find_link(&chip->p_vcpus, vc->v_link.l_id, &ins);
                ins_link(ins, &vc->v_link_pchip);
                chip->p_nvcpu++;

                (void) find_link(&core->c_vcpus, vc->v_link.l_id, &ins);
                ins_link(ins, &vc->v_link_core);
                core->c_nvcpu++;
        }

        (void) kstat_close(kc);

        nspec = 0;

        while ((optc = getopt(argc, argv, "cpr:S:stv")) != EOF) {
                switch (optc) {
                case 'c':
                        opt_c = 1;
                        break;
                case 'p':
                        opt_p = 1;
                        break;
                case 'r':
                        opt_r = optarg;
                        break;
                case 'S':
                        opt_S = optarg;
                        break;
                case 's':
                        opt_s = 1;
                        break;
                case 't':
                        opt_t = 1;
                        break;
                case 'v':
                        opt_v = 1;
                        break;
                default:
                        usage(NULL);
                }
        }

        if (opt_r != NULL) {
                if (optind != argc)
                        usage(_("cannot specify CPUs with -r"));
                if (opt_c || opt_p || opt_S != NULL || opt_s || opt_t || opt_v)
                        usage(_("cannot specify other arguments with -r"));

                read_property(opt_r);
                return (EXIT_SUCCESS);
        }

        if (opt_t != 0) {
                if (optind != argc)
                        usage(_("cannot specify CPUs with -t"));
                if (opt_s || opt_v)
                        usage(_("cannot specify -s or -v with -t"));
                if (opt_S != NULL && (opt_c || opt_p))
                        usage(_("cannot specify CPU state with -c or -p"));
                if (opt_c && opt_p)
                        usage(_("cannot specify -c and -p"));

                print_total(opt_c, opt_p, opt_S);
                return (EXIT_SUCCESS);
        }

        if (opt_S != NULL || opt_c)
                usage(_("cannot specify -S or -c without -t"));

        while (optind < argc) {
                long id;
                char *eptr;
                struct link *l;
                id = strtol(argv[optind], &eptr, 10);
                l = find_link(&vcpus, id, NULL);
                if ((*eptr != '\0') || (l == NULL)) {
                        (void) fprintf(stderr,
                            _("%s: processor %s: Invalid argument\n"),
                            cmdname, argv[optind]);
                        ex = 2;
                } else {
                        ((struct vcpu *)l->l_ptr)->v_doit = 1;
                        ((struct vcpu *)l->l_ptr)->v_pchip->p_doit = 1;
                        ((struct vcpu *)l->l_ptr)->v_core->c_doit = 1;
                }
                nspec++;
                optind++;
        }

        if (opt_s && opt_v) {
                usage(_("options -s and -v are mutually exclusive"));
        }
        if (opt_s && nspec != 1) {
                usage(_("must specify exactly one processor if -s used"));
        }
        if (opt_v && opt_p) {
                print_vp(nspec);
        } else if (opt_s && opt_p) {
                print_ps();
        } else if (opt_p) {
                print_p(nspec);
        } else if (opt_v) {
                print_v(nspec);
        } else if (opt_s) {
                print_s();
        } else {
                print_normal(nspec);
        }

        return (ex);
}