#include <azur/gint/render.h>
#include <gint/rtc.h>
#include <num/num.h>
#include <num/vec.h>
#include <random>
#include <string.h>
#include <stdlib.h>
#include "settings.h"
#include "render.h"
#include "util.h"
using namespace libnum;

uint8_t BOSONX_SHADER_PTRACE = -1;
static azrp_shader_t bosonx_shader_ptrace;

__attribute__((constructor))
static void register_shader(void)
{
    BOSONX_SHADER_PTRACE = azrp_register_shader(bosonx_shader_ptrace, NULL);
}

//---

struct command {
    uint8_t shader_id;
    uint8_t y;
    int16_t x;
    int16_t pty, height;
    uint16_t bgcolor;
    num t;
};

/* Nested circles */
#define CLIGHT_R1  10
#define CLIGHT_R2  15
#define CDARK_R1   36
#define CBEAMS1_R  41
#define CDARK_R2   45
#define CMAIN_R    48
#define CBEAMS2_R  52
#define CMAX_R     55

/* Ticks */
#define TICK_SIZE   7
#define TICK_COUNT 17 /* > 1 */

/* Traces */
#define TRACES_R             CDARK_R1
#define TRACES_NUM_THREADS   24 /* < 255*/
#define TRACES_PALETTE_SIZE  (TRACES_NUM_THREADS+1)
#define TRACES_TURN_MIN      num(0.02)
#define TRACES_TURN_RATIO    5
#define TRACES_TURN_MAX      (TRACES_TURN_RATIO * TRACES_TURN_MIN)

/* Precomputed and preallocated data */

/* Concentric circles pattern */
static uint8_t circles[CMAIN_R+1][CMAIN_R+1];
/* Trace image palette */
static uint16_t palette[TRACES_PALETTE_SIZE] = { 0 };
/* Trace turn vectors */
static vec2 trace_turns[TRACES_TURN_RATIO * TRACES_R];
/* Ticks tileset */
static image_t *img_ticks = NULL;
/* Trace image (regenerated every frame */
static image_t *img_trace = NULL;
/* Random seed for the entire animation */
static int seed = 0;

static void add_circle(int r, int value)
{
    int x = -r;
    int y = 0;
    int err = 2 - 2 * r;

    do {
        for(int i = x + CMAIN_R; i <= CMAIN_R; i++)
            circles[CMAIN_R - y][i] = value;
        r = err;
        if(r <= y)
            err += ++y * 2 + 1;
        if(r > x || err > y)
            err += ++x * 2 + 1;
    }
    while(x <= 0);
}

__attribute__((constructor))
static void precompute_circles(void)
{
    memset(circles, 0, sizeof circles);
    add_circle(CMAIN_R, 1);
    add_circle(CDARK_R2, 2);
    add_circle(CDARK_R1, 1);
    add_circle(CLIGHT_R2, 3);
    add_circle(CLIGHT_R1, 1);
}

__attribute__((constructor))
static void precompute_ticks(void)
{
    img_ticks = image_alloc(TICK_SIZE*TICK_COUNT, TICK_SIZE, IMAGE_P8_RGB565A);
    if(!img_ticks)
        return;

    static uint16_t img_ticks_palette[2] = { 0x5555, C_WHITE };
    image_set_palette(img_ticks, img_ticks_palette, 2, false);

    num angle_step = num(1.57080) / (TICK_COUNT - 1);

    for(int i = 0; i < TICK_COUNT; i++) {
        num alpha = i * angle_step;
        vec2 D(num_cos(alpha), num_sin(alpha));

        /* num_cos / num_sin rounding... */
        if(i == 0)
            D = vec2(1, 0);
        else if(i == TICK_COUNT - 1)
            D = vec2(0, 1);

        vec2 norm_D(-D.y, D.x);

        /* Render the segment that goes through the center with slope alpha,
           rendering with width 1 and length TICK_SIZE / 2 on both sides. */

        for(int y = 0; y < TICK_SIZE; y++)
        for(int x = 0; x < TICK_SIZE; x++) {
            vec2 p(+num(x) - num(TICK_SIZE - 1) / 2,
                   -num(y) + num(TICK_SIZE - 1) / 2);

            if(D.dot(p).abs() <= TICK_SIZE / 2 && norm_D.dot(p).abs() <= 1)
                image_set_pixel(img_ticks, i * TICK_SIZE + x, y, 0x81);
            else
                image_set_pixel(img_ticks, i * TICK_SIZE + x, y, 0x80);
        }
    }
}

static void render_tick(int cx, int cy, int radius, num angle)
{
    num c = num_cos(angle);
    num s = num_sin(angle);

    bool flip = false;

    /* Normalize to 0..π/2 */
    if(angle < 0) {
        angle = -angle;
        flip = !flip;
    }
    angle = angle % num(3.14159);
    if(angle > num(1.57080)) {
        angle = num(3.14159) - angle;
        flip = !flip;
    }

    int n = (angle * TICK_COUNT / num(1.57080)).ifloor();
    if(n < 0)
        n = 0;
    if(n >= TICK_COUNT)
        n = TICK_COUNT - 1;
    azrp_subimage(
        cx + int(c * radius) - TICK_SIZE / 2,
        cy - int(s * radius) - TICK_SIZE / 2,
        img_ticks, n * TICK_SIZE, 0, TICK_SIZE, TICK_SIZE,
        flip ? IMAGE_HFLIP : 0);
}

static void bosonx_shader_ptrace(void *, void *cmd0, void *frag0)
{
    struct command *cmd = (struct command *)cmd0;
    uint16_t *frag = (uint16_t *)frag0;
    frag += azrp_width * cmd->y + cmd->x;

    int pty = cmd->pty;
    int h = azrp_frag_height - cmd->y;
    if(h > cmd->height)
        h = cmd->height;

    cmd->height -= h;
    cmd->y = 0;
    cmd->pty += h;

    uint16_t palette[4] = { 0,
        render_blend(cmd->bgcolor, C_BLACK, 10),
        render_blend(cmd->bgcolor, C_BLACK, 16),
        render_blend(cmd->bgcolor, C_BLACK, 4),
    };

    do {
        int ly = pty - CMAX_R;
        uint16_t *fragc = frag + CMAX_R;

        if(ly >= -CMAIN_R && ly <= CMAIN_R) {
            int data_y = (ly <= 0) ? ly + CMAIN_R : CMAIN_R - ly;

            for(int lx = -CMAIN_R; lx <= 0; lx++) {
                int8_t c = circles[data_y][lx + CMAIN_R];
                if(c)
                    fragc[lx] = palette[c];
            }
            for(int lx = 1; lx <= CMAIN_R; lx++) {
                int8_t c = circles[data_y][CMAIN_R - lx];
                if(c)
                    fragc[lx] = palette[c];
            }
        }

        frag += azrp_width;
        pty++;
    }
    while(--h > 0);
}

/* Cannot use a constructor here because compiler schedules the global object
   constructor (which sets everything to 0) to run after global constructors */
static void precompute_turn_vectors(void)
{
    static bool computed = false;
    if(computed)
        return;

    for(int i = 0; i < TRACES_TURN_RATIO * TRACES_R; i++) {
        trace_turns[i].x = num_cos(i * TRACES_TURN_MIN);
        trace_turns[i].y = num_sin(i * TRACES_TURN_MIN);
    }
    computed = true;
}

void generate_traces(uint16_t color, bool successful, num t)
{
    int const N = 2 * TRACES_R;
    precompute_turn_vectors();

    if(!img_trace) {
        img_trace = image_alloc(N, N, IMAGE_P8_RGB565A);
        if(!img_trace)
            return;
        image_set_palette(img_trace, palette, TRACES_PALETTE_SIZE, false);
    }

    /* Clear out the entire image */
    uint8_t *px = (uint8_t *)img_trace->data;
    memset(px, 0x80, N * N);

    /* Generate the palette */
    int R5 = (color >> 11);
    int G6 = (color >> 5) & 0x3f;
    int B5 = (color & 0x1f);

    for(int i = 1; i < TRACES_PALETTE_SIZE; i++) {
        /* alpha=31: color, alpha=0: white */
        int alpha = (i >= TRACES_NUM_THREADS / 4) ? 0 :
                    31 - (i * 31) / (TRACES_NUM_THREADS / 4);
        int R = ((alpha * R5) + (31 - alpha) * 0x1f) / 31;
        int G = ((alpha * G6) + (31 - alpha) * 0x3f) / 31;
        int B = ((alpha * B5) + (31 - alpha) * 0x1f) / 31;
        palette[i] = (R << 11) | (G << 5) | B;
    }

    std::minstd_rand r(seed);
    std::minstd_rand r2(seed+1);

    int steps = (t * TRACES_R / TIME_PTRACE_ANIMATION).ifloor();
    if(steps > TRACES_R)
        steps = TRACES_R;

    for(int i = 0; i < TRACES_NUM_THREADS; i++) {
        vec2 dir = num_rand_vec2_std(r);
        int turn_factor = 1 + (r2() % 4);

        for(int s = 0; s < steps; s++) {
            vec2 straight = dir * s;
            vec2 rotated = straight;

            if(successful) {
                vec2 const &turn = trace_turns[s * turn_factor];
                rotated.x = straight.x * turn.x - straight.y * turn.y;
                rotated.y = straight.x * turn.y + straight.y * turn.x;
            }

            int x = (N / 2 + rotated.x.ifloor());
            int y = (N / 2 - rotated.y.ifloor());

            /* "Increase" the color of the pixel. The palette is a gradient so
               this will imitate an overlay effect */
            if((unsigned)y < N) {
                if((unsigned)x < N)
                    px[N * y + x]++;
                if((unsigned)x+1 < N)
                    px[N * y + x+1]++;
            }
            if((unsigned)y+1 < N) {
                if((unsigned)x < N)
                    px[N * (y+1) + x]++;
                if((unsigned)x+1 < N)
                    px[N * (y+1) + x+1]++;
            }
        }
    }

    /* Make sure every value is within the palette range */
    for(int i = 0; i < N * N; i++) {
        if((uint8_t)(px[i] - 0x80) > TRACES_PALETTE_SIZE)
            px[i] = 0x80 + TRACES_PALETTE_SIZE - 1;
    }
}

void shader_ptrace_reset(void)
{
    seed = rtc_ticks();
}

void shader_ptrace(int x, int y, uint16_t color, uint16_t bgcolor,
    bool successful, num t)
{
    prof_enter(azrp_perf_cmdgen);

    t = num_clamp(t, 0, TIME_PTRACE_ANIMATION);

    int frag_first, frag_offset, frag_count;
    int height = 2 * CMAX_R + 1;
    azrp_config_get_lines(y, height, &frag_first, &frag_offset, &frag_count);

    struct command *cmd = (struct command *)
        azrp_new_command(sizeof *cmd, frag_first, frag_count);
    if(cmd) {
        cmd->shader_id = BOSONX_SHADER_PTRACE;
        cmd->y = frag_offset;
        cmd->x = x;
        cmd->pty = 0;
        cmd->height = height;
        cmd->bgcolor = bgcolor;
        cmd->t = t;
    }

    prof_leave(azrp_perf_cmdgen);

    /* Start with 4/8 ticks and go up to 12/20 if un/successful */
    int N1 = successful ? 8 : 4, N2 = successful ? 20 : 12;
    int N = N1 +(N2 * t / TIME_PTRACE_ANIMATION).ifloor();

    std::minstd_rand r(seed);
    for(int i = 0; i < N; i++) {
        num speed = num_rand_between_std(r, -1, 1);
        num orig  = num_rand_between_std(r, 0, num(6.28));
        render_tick(x + CMAX_R, y + CMAX_R, CBEAMS1_R, orig + t * speed);
        render_tick(x + CMAX_R, y + CMAX_R, CBEAMS2_R, -orig - t * speed);
    }

    generate_traces(color, successful, t);
    if(img_trace) {
        azrp_image(x + CMAX_R - img_trace->width / 2,
                   y + CMAX_R - img_trace->height / 2,
                   img_trace);
    }
}