/*
 * Copyright (C) 2008 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <ctype.h>
#include <dirent.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#define MAX_LINE 512
#define MAX_FILENAME 64

const char *EXPECTED_VERSION = "Latency Top version : v0.1\n";
const char *SYSCTL_FILE = "/proc/sys/kernel/latencytop";
const char *GLOBAL_STATS_FILE = "/proc/latency_stats";
const char *THREAD_STATS_FILE_FORMAT = "/proc/%d/task/%d/latency";

struct latency_entry {
    struct latency_entry *next;
    unsigned long count;
    unsigned long max;
    unsigned long total;
    char reason[MAX_LINE];
};

static inline void check_latencytop() { }

static struct latency_entry *read_global_stats(struct latency_entry *list, int erase);
static struct latency_entry *read_process_stats(struct latency_entry *list, int erase, int pid);
static struct latency_entry *read_thread_stats(struct latency_entry *list, int erase, int pid, int tid, int fatal);

static struct latency_entry *alloc_latency_entry(void);
static void free_latency_entry(struct latency_entry *e);

static void set_latencytop(int on);
static struct latency_entry *read_latency_file(FILE *f, struct latency_entry *list);
static void erase_latency_file(FILE *f);

static struct latency_entry *find_latency_entry(struct latency_entry *e, char *reason);
static void print_latency_entries(struct latency_entry *head);

static void signal_handler(int sig);
static void disable_latencytop(void);

static int numcmp(const long long a, const long long b);
static int lat_cmp(const void *a, const void *b);

static void clear_screen(void);
static void usage(const char *cmd);

struct latency_entry *free_entries;

int main(int argc, char *argv[]) {
    struct latency_entry *e;
    int delay, iterations;
    int pid, tid;
    int count, erase;
    int i;

    delay = 1;
    iterations = 0;
    pid = tid = 0;
    
    for (i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "-d")) {
            if (i >= argc - 1) {
                fprintf(stderr, "Option -d expects an argument.\n");
                exit(EXIT_FAILURE);
            }
            delay = atoi(argv[++i]);
            continue;
        }
        if (!strcmp(argv[i], "-n")) {
            if (i >= argc - 1) {
                fprintf(stderr, "Option -n expects an argument.\n");
                exit(EXIT_FAILURE);
            }
            iterations = atoi(argv[++i]);
            continue;
        }
        if (!strcmp(argv[i], "-h")) {
            usage(argv[0]);
            exit(EXIT_SUCCESS);
        }
        if (!strcmp(argv[i], "-p")) {
            if (i >= argc - 1) {
                fprintf(stderr, "Option -p expects an argument.\n");
                exit(EXIT_FAILURE);
            }
            pid = atoi(argv[++i]);
            continue;
        }
        if (!strcmp(argv[i], "-t")) {
            if (i >= argc - 1) {
                fprintf(stderr, "Option -t expects an argument.\n");
                exit(EXIT_FAILURE);
            }
            tid = atoi(argv[++i]);
            continue;
        }
        fprintf(stderr, "Invalid argument \"%s\".\n", argv[i]);
        usage(argv[0]);
        exit(EXIT_FAILURE);
    }

    if (tid && !pid) {
        fprintf(stderr, "If you provide a thread ID with -t, you must provide a process ID with -p.\n");
        exit(EXIT_FAILURE);
    }

    check_latencytop();

    free_entries = NULL;

    signal(SIGINT, &signal_handler);
    signal(SIGTERM, &signal_handler);

    atexit(&disable_latencytop);

    set_latencytop(1);

    count = 0;
    erase = 1;

    while ((iterations == 0) || (count++ < iterations)) {

        sleep(delay);

        e = NULL;
        if (pid) {
            if (tid) {
                e = read_thread_stats(e, erase, pid, tid, 1);
            } else {
                e = read_process_stats(e, erase, pid);
            }
        } else {
            e = read_global_stats(e, erase);
        }
        erase = 0;

        clear_screen();
        if (pid) {
            if (tid) {
                printf("Latencies for thread %d in process %d:\n", tid, pid);
            } else {
                printf("Latencies for process %d:\n", pid);
            }
        } else {
            printf("Latencies across all processes:\n");
        }
        print_latency_entries(e);
    }

    set_latencytop(0);

    return 0;
}

static struct latency_entry *read_global_stats(struct latency_entry *list, int erase) {
    FILE *f;
    struct latency_entry *e;

    if (erase) {
        f = fopen(GLOBAL_STATS_FILE, "w");
        if (!f) {
            fprintf(stderr, "Could not open global latency stats file: %s\n", strerror(errno));
            exit(EXIT_FAILURE);
        }
        fprintf(f, "erase\n");
        fclose(f);
    }
    
    f = fopen(GLOBAL_STATS_FILE, "r");
    if (!f) {
        fprintf(stderr, "Could not open global latency stats file: %s\n", strerror(errno));
        exit(EXIT_FAILURE);
    }

    e = read_latency_file(f, list);

    fclose(f);

    return e;
}

static struct latency_entry *read_process_stats(struct latency_entry *list, int erase, int pid) {
    char dirname[MAX_FILENAME];
    DIR *dir;
    struct dirent *ent;
    struct latency_entry *e;
    int tid;

    sprintf(dirname, "/proc/%d/task", pid);
    dir = opendir(dirname);
    if (!dir) {
        fprintf(stderr, "Could not open task dir for process %d.\n", pid);
        fprintf(stderr, "Perhaps the process has terminated?\n");
        exit(EXIT_FAILURE);
    }

    e = list;
    while ((ent = readdir(dir))) {
        if (!isdigit(ent->d_name[0]))
            continue;

        tid = atoi(ent->d_name);

        e = read_thread_stats(e, erase, pid, tid, 0);
    }

    closedir(dir);

    return e;
}

static struct latency_entry *read_thread_stats(struct latency_entry *list, int erase, int pid, int tid, int fatal) {
    char filename[MAX_FILENAME];
    FILE *f;
    struct latency_entry *e;

    sprintf(filename, THREAD_STATS_FILE_FORMAT, pid, tid);

    if (erase) {
        f = fopen(filename, "w");
        if (!f) {
            if (fatal) {
                fprintf(stderr, "Could not open %s: %s\n", filename, strerror(errno));
                fprintf(stderr, "Perhaps the process or thread has terminated?\n");
                exit(EXIT_FAILURE);
            } else {
                return list;
            }
        }
        fprintf(f, "erase\n");
        fclose(f);
    }
    
    f = fopen(GLOBAL_STATS_FILE, "r");
    if (!f) {
        if (fatal) {
            fprintf(stderr, "Could not open %s: %s\n", filename, strerror(errno));
            fprintf(stderr, "Perhaps the process or thread has terminated?\n");
            exit(EXIT_FAILURE);
        } else {
            return list;
        }
    }

    e = read_latency_file(f, list);

    fclose(f);

    return e;
}

static struct latency_entry *alloc_latency_entry(void) {
    struct latency_entry *e;

    if (free_entries) {
        e = free_entries;
        free_entries = free_entries->next;
    } else {
        e = calloc(1, sizeof(struct latency_entry));
        if (!e) {
            fprintf(stderr, "Could not allocate latency entry: %s\n", strerror(errno));
            exit(EXIT_FAILURE);
        }
    }

    return e;
}

static void free_latency_entry(struct latency_entry *e) {
    e->next = free_entries;
    free_entries = e;
}

static struct latency_entry *find_latency_entry(struct latency_entry *head, char *reason) {
    struct latency_entry *e;

    e = head;

    while (e) {
        if (!strcmp(e->reason, reason))
            return e;
        e = e->next;
    }

    return NULL;
}

static void set_latencytop(int on) {
    FILE *f;

    f = fopen(SYSCTL_FILE, "w");
    if (!f) {
        fprintf(stderr, "Could not open %s: %s\n", SYSCTL_FILE, strerror(errno));
        exit(EXIT_FAILURE);
    }

    fprintf(f, "%d\n", on);

    fclose(f);
}

static void erase_latency_file(FILE *f) {
    fprintf(f, "erase\n");
}

static struct latency_entry *read_latency_file(FILE *f, struct latency_entry *list) {
    struct latency_entry *e, *head;
    char line[MAX_LINE];
    unsigned long count, max, total;
    char reason[MAX_LINE];

    head = list;

    if (!fgets(line, MAX_LINE, f)) {
        fprintf(stderr, "Could not read latency file version: %s\n", strerror(errno));
        exit(EXIT_FAILURE);
    }

    if (strcmp(line, EXPECTED_VERSION) != 0) {
        fprintf(stderr, "Expected version: %s\n", EXPECTED_VERSION);
        fprintf(stderr, "But got version: %s", line);
        exit(EXIT_FAILURE);
    }

    while (fgets(line, MAX_LINE, f)) {
        sscanf(line, "%ld %ld %ld %s", &count, &total, &max, reason);
        if (max > 0 || total > 0) {
            e = find_latency_entry(head, reason);
            if (e) {
                e->count += count;
                if (max > e->max)
                    e->max = max;
                e->total += total;
            } else {
                e = alloc_latency_entry();
                e->count = count;
                e->max = max;
                e->total = total;
                strcpy(e->reason, reason);
                e->next = head;
                head = e;
            }
        }
    }

    return head;
}

static void print_latency_entries(struct latency_entry *head) {
    struct latency_entry *e, **array;
    unsigned long average;
    int i, count;

    e = head;
    count = 0;
    while (e) {
        count++;
        e = e->next;
    }

    e = head;
    array = calloc(count, sizeof(struct latency_entry *));
    if (!array) {
        fprintf(stderr, "Error allocating array: %s\n", strerror(errno));
        exit(EXIT_FAILURE);
    }
    for (i = 0; i < count; i++) {
        array[i] = e;
        e = e->next;
    }

    qsort(array, count, sizeof(struct latency_entry *), &lat_cmp);

    printf("%10s  %10s  %7s  %s\n", "Maximum", "Average", "Count", "Reason");
    for (i = 0; i < count; i++) {
        e = array[i];
        average = e->total / e->count;
        printf("%4lu.%02lu ms  %4lu.%02lu ms  %7ld  %s\n",
            e->max / 1000, (e->max % 1000) / 10,
            average / 1000, (average % 1000) / 10,
            e->count,
            e->reason);
    }

    free(array);
}

static void signal_handler(int sig) {
    exit(EXIT_SUCCESS);
}

static void disable_latencytop(void) {
    set_latencytop(0);
}

static void clear_screen(void) {
    printf("\n\n");
}

static void usage(const char *cmd) {
    fprintf(stderr, "Usage: %s [ -d delay ] [ -n iterations ] [ -p pid [ -t tid ] ] [ -h ]\n"
                    "    -d delay       Time to sleep between updates.\n"
                    "    -n iterations  Number of updates to show (0 = infinite).\n"
                    "    -p pid         Process to monitor (default is all).\n"
                    "    -t tid         Thread (within specified process) to monitor (default is all).\n"
                    "    -h             Display this help screen.\n",
        cmd);
}

static int numcmp(const long long a, const long long b) {
    if (a < b) return -1;
    if (a > b) return 1;
    return 0;
}

static int lat_cmp(const void *a, const void *b) {
    const struct latency_entry *pa, *pb;

    pa = (*((struct latency_entry **)a));
    pb = (*((struct latency_entry **)b));

    return numcmp(pb->max, pa->max);
}