#ifndef COMMON_RPIBM_H
#define COMMON_RPIBM_H

// Header definitions for memory mapped IO and common utilities.

// Setup the right base memory address for different types of RPI.
#if RPI_VERSION == 3
#define MEM_BASE 0x3F000000
#elif RPI_VERSION == 4
#define MEM_BASE 0xFE000000
#else
#define MEM_BASE 0
#error RPI_VERSION NOT DEFINED
#endif

#include "shorthand.h"

//
// Utils.
//

// Wait for N processor cycles before returning. Implemented in start.s.
void delay(u64 ticks);

// Clear N bytes at the given address 8 bytes at a time. Implemented in start.s.
void memzero64(void *src, u64 n);

static void
memzero8(void *src, u64 n) {
    u8 * ptr = src;
    for (size_t i = 0; i < n; i++) {
        ptr[i] = 0;
    }
}

//
// GPIO Registers.
//

typedef struct GpioPinData {
    vu32 reserved;
    vu32 data[2];
} GpioPinData;

typedef struct GpioRegs {
    vu32 func_select[6];
    GpioPinData output_set;
    GpioPinData output_clear;
    GpioPinData level;
    GpioPinData ev_detect_status;
    GpioPinData re_detect_enable;
    GpioPinData fe_detect_enable;
    GpioPinData hi_detect_enable;
    GpioPinData lo_detect_enable;
    GpioPinData async_re_detect;
    GpioPinData async_fe_detect;
    vu32 reserved;
    vu32 pupd_enable;
    vu32 pupd_enable_clocks[2];
} GpioRegs;

typedef enum GpioFunc {
    GPIO_INPUT  = 0,
    GPIO_OUTPUT = 1,
    GPIO_ALT0   = 4,
    GPIO_ALT1   = 5,
    GPIO_ALT2   = 6,
    GPIO_ALT3   = 7,
    GPIO_ALT4   = 3,
    GPIO_ALT5   = 2,
} GpioFunc;

#define GPIO ((GpioRegs *)(MEM_BASE + 0x00200000))

static inline void
gpio_pin_set_func(u8 pin, GpioFunc func) {
    u8 start = (pin * 3) % 30;
    u8 reg = pin / 10;
    u32 selector = GPIO->func_select[reg];
    selector &= ~(7 << start);
    selector |= (func << start);
    GPIO->func_select[reg] = selector;
}

static inline void
gpio_pin_enable(u8 pin) {
    GPIO->pupd_enable = 0;
    delay(150);
    GPIO->pupd_enable_clocks[pin / 32] = 1 << (pin % 32);
    delay(150);
    GPIO->pupd_enable = 0;
    GPIO->pupd_enable_clocks[pin / 32] = 0;
}

//
// AUX Registers.
//

typedef struct AuxRegs {
    vu32 irq_status;
    vu32 enables;
    vu32 reserved[14];
    vu32 mu_io;
    vu32 mu_ier;
    vu32 mu_iir;
    vu32 mu_lcr;
    vu32 mu_mcr;
    vu32 mu_lsr;
    vu32 mu_msr;
    vu32 mu_scratch;
    vu32 mu_control;
    vu32 mu_status;
    vu32 mu_baud_rate;
} AuxRegs;

#define AUX ((AuxRegs *)(MEM_BASE + 0x00215000))

static inline void
uart_init() {
    gpio_pin_set_func(14, GPIO_ALT5);
    gpio_pin_set_func(15, GPIO_ALT5);
    gpio_pin_enable(14);
    gpio_pin_enable(15);

    AUX->enables = 1;
    AUX->mu_control = 0;
    AUX->mu_ier = 0;
    AUX->mu_lcr = 3;
    AUX->mu_mcr = 0;

#if RPI_VERSION == 3
    AUX->mu_baud_rate = 270; // RPI3: 115200 @ 250 MHz
#endif

#if RPI_VERSION == 4
    AUX->mu_baud_rate = 541; // RPI4: 115200 @ 500 MHz
#endif

    AUX->mu_control = 3;
    return;
}

static inline void
uart_putc(char c) {
    while (!(AUX->mu_lsr & (1 << 5)));
    AUX->mu_io = c;
    return;
}

static inline char
uart_getc() {
    while (!(AUX->mu_lsr & 1));
    u8 c = AUX->mu_io & 0xFF;
    return c == '\r' ? '\n' : c;
}

static inline void
uart_puts(char *s) {
    while (*s) {
        if (*s == '\n') {
            uart_putc('\r');
        }
        uart_putc(*s++);
    }
}

static inline void
uart_hex(unsigned int d) {
    unsigned int n;
    uart_puts("0x");
    for (int c = 28; c >= 0; c -= 4) {
        n = (d>>c) & 0xF;
        n += n> 9 ? 0x37 : 0x30;
        uart_putc(n);
    }
}

//
// VideoCore Mailboxes.
//

typedef struct MailboxRegs {
    vu32 read;
    vu32 reserved[5];
    vu32 status;
    vu32 config;
    vu32 write;
} MailboxRegs;

#define MBOX ((MailboxRegs *)(MEM_BASE + 0x0000B880))

typedef enum MboxChannels {
    MBOX_CH_POWER = 0,
    MBOX_CH_FB    = 1,
    MBOX_CH_VUART = 2,
    MBOX_CH_VCHIQ = 3,
    MBOX_CH_LEDS  = 4,
    MBOX_CH_BTNS  = 5,
    MBOX_CH_TOUCH = 6,
    MBOX_CH_COUNT = 7,
    MBOX_CH_PROP  = 8,
} MboxChannels;

// Property channel response type.
#define MBOX_REQUEST      0
#define MBOX_RESPONSE_OK  0x80000000
#define MBOX_RESPONSE_ERR 0x80000001

// Mailbox status codes.
#define MBOX_FULL  0x80000000
#define MBOX_EMPTY 0x40000000

// Mailbox tags.
#define  MBOX_TAG_SET_SCREEN  0x48003
#define  MBOX_TAG_SET_VSCREEN 0x48004
#define  MBOX_TAG_SET_DEPTH   0x48005
#define  MBOX_TAG_SET_RGB     0x48006
#define  MBOX_TAG_GET_FB      0x40001
#define  MBOX_TAG_END         0

// Mailbox request buffer. Needs to be 16 bit aligned to avoid issues.
volatile u32 __attribute__((aligned(16))) mbox[64];

static inline void
mb_write(u8 channel) {
    while (MBOX->status & MBOX_FULL);
    MBOX->write = ((uintptr_t)mbox & ~0xF) | (channel & 0xF);
}

static inline u32
mb_read(u8 channel) {
    while (true) {
        while (MBOX->status & MBOX_EMPTY);
        u32 data = MBOX->read;
        if ((u8)(data & 0xF) == channel) {
            return data & 0xFFFFFFF0;
        }
    }
}

static inline int
mb_call(unsigned char ch) {
    mb_write(ch);
    return mb_read(ch);
}

//
// Timers.
//

typedef struct TimerRegs {
    vu32 ctrl_status;
    vu32 counter_low;
    vu32 counter_high;
    vu32 compare[4];
} TimerRegs;

#define TIMER ((TimerRegs *)(MEM_BASE + 0x00003000))

static u64
timer_get_ticks(void) {
    u32 high = TIMER->counter_high;
    u32 low = TIMER->counter_low;

    // Reading is not atomic, since we need to read two 32 bit values. We check
    // if the high bits changed after reading the low bits for safety.
    if (high != TIMER->counter_high) {
        high = TIMER->counter_high;
        low = TIMER->counter_low;
    }
    return ((u64)high << 32) | low;
}

static void
wait(u64 usec) {
    u64 current_ticks = timer_get_ticks();
    // The system timer clock runs at 1 Mhz.
    while (timer_get_ticks() < (current_ticks + usec));
}

//
// Utilities.
//

static void *
memcpy(void *dest, const void *src, u32 n) {
    u8 *from = (u8*)src;
    u8 *to = (u8*)dest;
    for (size_t i = 0; i < n; i++) {
        to[i] = from[i];
    }
    return dest;
}

#endif // COMMON_RPIBM_H