/** @file
    Terminal control utility functions.

    Copyright (C) 2018 Christian Zuckschwerdt

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.
*/

#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#ifndef _WIN32
#include <unistd.h>
#include <sys/ioctl.h>
#endif

#include "term_ctl.h"

#ifdef _WIN32
#include <stdlib.h>
#include <io.h>
#include <limits.h>
#include <windows.h>
#include <lm.h>

#ifndef STDOUT_FILENO
#define STDOUT_FILENO   1
#endif

#ifndef STDERR_FILENO
#define STDERR_FILENO   2
#endif

#ifndef UNUSED
#define UNUSED(x) (void)(x)
#endif

typedef struct console {
    CONSOLE_SCREEN_BUFFER_INFO info;
    BOOL                       redirected;
    BOOL                       ansi;
    HANDLE                     hnd;
    FILE                      *file;
    WORD                       fg, bg;
} console_t;

static WORD _term_get_win_color(BOOL fore, term_color_t color)
{
    // applies intensity only on foreground,
    // i.e. background will fallback to dim
    switch (color) {
       case TERM_COLOR_RESET: /* Cannot occur; just to suppress a warning */
            break;
       case TERM_COLOR_BLACK:
            return (0);
       case TERM_COLOR_BLUE:
            return (1);
       case TERM_COLOR_GREEN:
            return (2);
       case TERM_COLOR_CYAN:
            return (3);
       case TERM_COLOR_RED:
            return (4);
       case TERM_COLOR_MAGENTA:
            return (5);
       case TERM_COLOR_YELLOW:
            return (6);
       case TERM_COLOR_WHITE:
            return (7);
       case TERM_COLOR_BRIGHT_BLACK:
            return (fore ? 0 + FOREGROUND_INTENSITY : 0);
       case TERM_COLOR_BRIGHT_BLUE:
            return (fore ? 1 + FOREGROUND_INTENSITY : 1);
       case TERM_COLOR_BRIGHT_GREEN:
            return (fore ? 2 + FOREGROUND_INTENSITY : 2);
       case TERM_COLOR_BRIGHT_CYAN:
            return (fore ? 3 + FOREGROUND_INTENSITY : 3);
       case TERM_COLOR_BRIGHT_RED:
            return (fore ? 4 + FOREGROUND_INTENSITY : 4);
       case TERM_COLOR_BRIGHT_MAGENTA:
            return (fore ? 5 + FOREGROUND_INTENSITY : 5);
       case TERM_COLOR_BRIGHT_YELLOW:
            return (fore ? 6 + FOREGROUND_INTENSITY : 6);
       case TERM_COLOR_BRIGHT_WHITE:
            return (fore ? 7 + FOREGROUND_INTENSITY : 7);
    }
    fprintf(stderr,"FATAL: No mapping for TERM_COLOR_x=%d (fore: %d)\n", color, fore);
    return (0);
}

static void _term_set_color(console_t *console, BOOL fore, term_color_t color)
{
    WORD win_color;

    if (!console->file)
        return;

    if (color == TERM_COLOR_RESET) {
        console->fg = (console->info.wAttributes & 7);
        console->bg = (console->info.wAttributes & ~7);
    }
    else if (fore) {
        console->fg = _term_get_win_color(TRUE, color);
    }
    else if (color <= TERM_COLOR_WHITE ||
            (color >= TERM_COLOR_BRIGHT_BLACK && color <= TERM_COLOR_BRIGHT_WHITE)) {
        console->bg = 16 * _term_get_win_color(FALSE, color);
    }
    else
        return;

    win_color = console->bg + console->fg;

    /* Hack: as WinCon does not have color-themes (as Linux have) and no 'TERM_COLOR_BRIGHT_x'
     * codes are used, always use a high-intensity foreground color. This look best in
     * CMD with the default black background color.
     *
     * Note: do not do this for "BLACK" as that would turn it into "GREY".
     */
    if (fore && color != TERM_COLOR_RESET && color != TERM_COLOR_BLACK)
        win_color |= FOREGROUND_INTENSITY;

    fflush(console->file);
    SetConsoleTextAttribute(console->hnd, win_color);
}

/*
 * Cleanup in case we got a SIGINT signal in the middle of a
 * non-default colour output.
 */
static void _term_free(console_t *console)
{
    if (console->hnd && console->hnd != INVALID_HANDLE_VALUE) {
        fflush(console->file);
        SetConsoleTextAttribute(console->hnd, console->info.wAttributes);
    }
    free(console);
}

static BOOL _term_has_color(console_t *console)
{
    return console->hnd && !console->redirected;
}

static void *_term_init(FILE *fp)
{
    console_t *console = calloc(1, sizeof(*console));
    if (!console) {
        fprintf(stderr, "term_init failed\n");
        return NULL; // NOTE: return NULL on alloc failure.
    }

    int fd = fileno(fp);
    if (fd == STDOUT_FILENO) {
        console->hnd = GetStdHandle(STD_OUTPUT_HANDLE);
        console->file = fp;
    }
    else if (fd == STDERR_FILENO) {
        console->hnd = GetStdHandle(STD_ERROR_HANDLE);
        console->file = fp;
    }
    console->redirected = (console->hnd == INVALID_HANDLE_VALUE) ||
                         (!GetConsoleScreenBufferInfo(console->hnd, &console->info)) ||
                         (GetFileType(console->hnd) != FILE_TYPE_CHAR);

    // Test for Windows 10 to enable ANSI output, needs netapi32.dll
    LPWKSTA_INFO_100 pBuf = NULL;
    NET_API_STATUS nStatus;
    nStatus = NetWkstaGetInfo(NULL, 100, (LPBYTE *)&pBuf);
    if (nStatus == NERR_Success) {
        console->ansi = (pBuf->wki100_platform_id == 500) && (pBuf->wki100_ver_major >= 10);
    }
    if (pBuf != NULL) {
        NetApiBufferFree(pBuf);
    }

    // Windows 10 version 1511 added ANSI filters to cmd and terminal.
    // To use ANSI colors in Windows versions 1511 to 1903 requires setting VirtualTerminalLevel
    // by calling the SetConsoleMode API with the ENABLE_VIRTUAL_TERMINAL_PROCESSING flag.
    // ANSI colors are available by default in Windows version 1909 or newer
    if (console->ansi) {
        DWORD dwMode = 0;
        GetConsoleMode(console->hnd, &dwMode);
        dwMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
        SetConsoleMode(console->hnd, dwMode);
        // Check if it worked, it will fail for Legacy console-mode
        GetConsoleMode (console->hnd, &dwMode);
        if (!(dwMode & ENABLE_VIRTUAL_TERMINAL_PROCESSING)) {
            console->ansi = 0;
        }
    }

    _term_set_color(console, FALSE, TERM_COLOR_RESET); /* Set 'console->fg' and 'console->bg' */

    return console;
}
#endif /* _WIN32 */

int term_get_columns(void *ctx)
{
#ifdef _WIN32
    console_t *console = (console_t *)ctx;
    /*
     * Call this again as the screen dimensions could have changes since
     * we called '_term_init()'.
     */
    CONSOLE_SCREEN_BUFFER_INFO c_info;

    if (!console->hnd || console->hnd == INVALID_HANDLE_VALUE)
       return (80);

    if (!GetConsoleScreenBufferInfo(console->hnd, &c_info))
       return (80);
    return (c_info.srWindow.Right - c_info.srWindow.Left + 1);
#else
    FILE *fp = (FILE *)ctx;
    struct winsize w;
    ioctl(fileno(fp), TIOCGWINSZ, &w);
    return w.ws_col;
#endif
}

int term_has_color(void *ctx)
{
#ifdef _WIN32
    return _term_has_color(ctx);
#else
    FILE *fp = (FILE *)ctx;
    return isatty(fileno(fp)); // || get_env("force_color")
#endif
}

void *term_init(FILE *fp)
{
#ifdef _WIN32
    return _term_init(fp);
#else
    return fp;
#endif
}

void term_free(void *ctx)
{
    if (!ctx)
        return;
#ifdef _WIN32
    _term_free(ctx);
#else
    FILE *fp = (FILE *)ctx;
    if (term_has_color(ctx))
        fprintf(fp, "\033[0m");
#endif
}

void term_ring_bell(void *ctx)
{
#ifdef _WIN32
    Beep(800, 20);
    UNUSED(ctx);
#else
    FILE *fp = (FILE *)ctx;
    fprintf(fp, "\a");
#endif
}

void term_set_fg(void *ctx, term_color_t color)
{
#ifdef _WIN32
    console_t *console = (console_t *)ctx;
    if (!console->ansi) {
        _term_set_color(ctx, TRUE, color);
        return;
    }
    FILE *fp = console->file;
#else
    FILE *fp = (FILE *)ctx;
#endif
    if (color == TERM_COLOR_RESET)
        fprintf(fp, "\033[0m");
    else
        fprintf(fp, "\033[%d;1m", color);
}

void term_set_bg(void *ctx, term_color_t bg, term_color_t fg)
{
    if (bg < TERM_COLOR_BLACK
            || (bg > TERM_COLOR_WHITE && bg < TERM_COLOR_BRIGHT_BLACK)
            || bg > TERM_COLOR_BRIGHT_WHITE) {
        bg = 0;
    }
    if (fg < TERM_COLOR_BLACK
            || (fg > TERM_COLOR_WHITE && fg < TERM_COLOR_BRIGHT_BLACK)
            || fg > TERM_COLOR_BRIGHT_WHITE) {
        fg = 0;
    }

#ifdef _WIN32
    console_t *console = (console_t *)ctx;
    if (!console->ansi) {
        if (bg)
            _term_set_color(ctx, FALSE, bg);
        if (fg)
            _term_set_color(ctx, TRUE, fg);
        return;
    }
    FILE *fp = console->file;
#else
    FILE *fp = (FILE *)ctx;
#endif
    if (bg && fg)
        fprintf(fp, "\033[%d;%dm", bg + 10, fg);
    else if (bg)
        fprintf(fp, "\033[%dm", bg + 10);
    else if (fg)
        fprintf(fp, "\033[%dm", fg);
}

#define DIM(array) (int) (sizeof(array) / sizeof(array[0]))

static term_color_t color_map[] = {
                    TERM_COLOR_RESET,     /* "~0" */
                    TERM_COLOR_GREEN,
                    TERM_COLOR_WHITE,     /* "~2" */
                    TERM_COLOR_BLUE,
                    TERM_COLOR_CYAN,      /* "~4" */
                    TERM_COLOR_MAGENTA,
                    TERM_COLOR_YELLOW,    /* "~6" */
                    TERM_COLOR_BLACK,
                    TERM_COLOR_RED,       /* "~8" */
                  };

int term_set_color_map(int ascii_idx, term_color_t color)
{
    ascii_idx -= '0';
    if (ascii_idx < 0 || ascii_idx >= DIM(color_map))
        return -1;
    color_map[ascii_idx] = color;
    return ascii_idx;
}

int term_get_color_map(int ascii_idx)
{
    int i;

    ascii_idx -= '0';
    for (i = 0; ascii_idx >= 0 && i < DIM(color_map); i++)
        if (i == ascii_idx)
           return (int)color_map[i];
    return -1;
}

int term_puts(void *ctx, char const *buf)
{
    char const *p = buf;
    int i, len, buf_len, color;
    FILE *fp;

    if (!ctx)
        return fprintf(stderr, "%s", buf);

#ifdef _WIN32
    console_t *console = (console_t *)ctx;
    fp = console->file;
#else
    fp = (FILE *)ctx;
#endif

    if (!fp)
        fp = stderr;

    buf_len = (int)strlen(buf);
    for (i = len = 0; *p && i < buf_len; i++, p++) {
        if (*p == '~') {
            p++;
            color = ctx ? term_get_color_map(*p) : -1;
            if (color >= 0)
                term_set_fg(ctx, (term_color_t)color);
        }
        else {
            fputc(*p, fp);
            len++;
        }
    }
    return len;
}

int term_printf(void *ctx, _Printf_format_string_ char const *format, ...)
{
    int len;
    va_list args;
    char buf[4000];

    va_start(args, format);

    // Terminate first in case a buggy '_MSC_VER < 1900' is used.
    buf[sizeof(buf)-1] = '\0';
    vsnprintf(buf, sizeof(buf)-1, format, args);
    len = term_puts(ctx, buf);
    va_end (args);
    return len;
}

int term_help_fputs(void *ctx, char const *buf, FILE *fp)
{
    char const *p = buf;
    int i, len, buf_len, color, state = 0, set_color = -1, next_color = -1;
    if (!fp) {
        fp = stderr;
    }

    if (!ctx) {
        return fprintf(fp, "%s", buf);
    }

#ifdef _WIN32
    console_t *console = (console_t *)ctx;
    fp = console->file;
#else
    fp = (FILE *)ctx;
#endif

    buf_len = (int)strlen(buf);
    for (i = len = 0; *p && i < buf_len; i++, p++) {
        if (*p == '~') {
            p++;
            color = ctx ? term_get_color_map(*p) : -1;
            if (color >= 0)
                term_set_fg(ctx, (term_color_t)color);
            continue;
        }

        if (state == 0 && *p == '[') {
            state = 1;
            next_color = 5;
        }
        else if ((state == 1 || state == 2) && *p == ']' && ((p[1] == ' ' && p[2] != '|') || p[1] == '\n' || p[1] == '\0')) {
            state = 0;
            set_color = 0;
        }
        else if (state == 1 && *p == ' ') {
            state = 2;
            next_color = 6;
        }
        else if (state == 2 && *p == '|') {
            set_color = 0;
            next_color = 6;
        }

        else if (state == 0 && *p == '=' && p[1] == ' ') {
            state = 3;
            set_color = 1;
        }
        else if (state == 3 && *p == '=') {
            state = 0;
            next_color = 0;
        }

        else if (state == 0 && *p == '\'') {
            state = 4;
            next_color = 4;
        }
        else if (state == 4 && *p == '\'') {
            state = 0;
            set_color = 0;
        }

        else if (state == 0 && *p == '\"') {
            state = 5;
            set_color = 4;
        }
        else if (state == 5 && *p == '\"') {
            state = 0;
            next_color = 0;
        }

        if (set_color >= 0) {
            color = ctx ? (int)color_map[set_color] : -1;
            if (color >= 0)
                term_set_fg(ctx, (term_color_t)color);
        }
        set_color = next_color;
        next_color = -1;

        fputc(*p, fp);
        len++;
    }
    return len;
}

int term_help_fprintf(FILE *fp, _Printf_format_string_ char const *format, ...)
{
    int len;
    va_list args;
    char buf[4000];

    va_start(args, format);

    void *term = term_init(fp);
    if (!term_has_color(term)) {
        term_free(term);
        term = NULL;
    }

    // Terminate first in case a buggy '_MSC_VER < 1900' is used.
    buf[sizeof(buf) - 1] = '\0';
    vsnprintf(buf, sizeof(buf) - 1, format, args);
    len = term_help_fputs(term, buf, fp);

    term_free(term);

    va_end(args);
    return len;
}
