#include <stdint.h>

// --- BCM2835 PERIPHERAL ADDRESSES ---
// Base address for Pi 1 is 0x20000000
#define PERIPHERAL_BASE 0x20000000
#define GPIO_BASE       (PERIPHERAL_BASE + 0x200000)
#define PWM_BASE        (PERIPHERAL_BASE + 0x20C000)
#define CLK_BASE        (PERIPHERAL_BASE + 0x101000)
#define UART0_BASE      (PERIPHERAL_BASE + 0x201000) // PL011 UART

// --- MAILBOX REGISTERS (VIDEO) ---
#define MBOX_READ       (PERIPHERAL_BASE + 0xB880)
#define MBOX_STATUS     (PERIPHERAL_BASE + 0xB898)
#define MBOX_WRITE      (PERIPHERAL_BASE + 0xB8A0)

// --- REGISTER ACCESS MACROS ---
#define MMIO32(addr) (*(volatile uint32_t *)(addr))

// --- GPIO REGISTERS ---
#define GPFSEL0   (GPIO_BASE + 0x00)
#define GPFSEL1   (GPIO_BASE + 0x04)
#define GPFSEL4   (GPIO_BASE + 0x10) // For Pin 40, 45

// --- PWM REGISTERS ---
#define PWM_CTL   (PWM_BASE + 0x00)
#define PWM_STA   (PWM_BASE + 0x04)
#define PWM_RNG1  (PWM_BASE + 0x10)
#define PWM_DAT1  (PWM_BASE + 0x14)
#define PWM_RNG2  (PWM_BASE + 0x20)
#define PWM_DAT2  (PWM_BASE + 0x24)
#define PWM_FIFO  (PWM_BASE + 0x18)

#define PWM_CTL_PWEN1   (1 << 0)
#define PWM_CTL_USEF1   (1 << 5) // Use FIFO for channel 1
#define PWM_CTL_PWEN2   (1 << 8)
#define PWM_CTL_USEF2   (1 << 13) // Use FIFO for channel 2
#define PWM_CTL_CLRF    (1 << 6)
#define PWM_STA_FULL1   (1 << 0)

// --- CLOCK MANAGER ---
#define CM_PWMCTL (CLK_BASE + 0xA0)
#define CM_PWMDIV (CLK_BASE + 0xA4)
#define CM_PASSWD (0x5A << 24)

// --- UART REGISTERS ---
#define UART0_DR    (UART0_BASE + 0x00)
#define UART0_FR    (UART0_BASE + 0x18)
#define UART0_IBRD  (UART0_BASE + 0x24)
#define UART0_FBRD  (UART0_BASE + 0x28)
#define UART0_LCRH  (UART0_BASE + 0x2C)
#define UART0_CR    (UART0_BASE + 0x30)

// --- UTILS ---
void delay(int32_t count) {
    while (count--) asm volatile("nop");
}

void gpio_set_func(uint32_t pin, uint32_t func) {
    uint32_t offset = (pin / 10) * 4;
    uint32_t shift = (pin % 10) * 3;
    uint32_t addr = GPIO_BASE + offset;
    
    uint32_t val = MMIO32(addr);
    val &= ~(7 << shift);
    val |= (func << shift);
    MMIO32(addr) = val;
}

void uart_init() {
    // Disable UART
    MMIO32(UART0_CR) = 0;
    
    // Setup GPIO 14/15 as Alt0 (UART0 TX/RX)
    gpio_set_func(14, 4); // Alt0
    gpio_set_func(15, 4); // Alt0

    // Baud rate setup (assuming 3MHz default UART clock)
    // For 115200: 3000000 / (16 * 115200) = 1.627
    // IBRD = 1
    // FBRD = 0.627 * 64 = 40
    MMIO32(UART0_IBRD) = 1;
    MMIO32(UART0_FBRD) = 40;

    // 8 bit, no parity, FIFO enabled
    MMIO32(UART0_LCRH) = (1 << 4) | (1 << 5) | (1 << 6);

    // Enable UART, TX, RX
    MMIO32(UART0_CR) = (1 << 0) | (1 << 8) | (1 << 9);
}

char uart_read_byte() {
    // Check RX FIFO empty flag
    if (MMIO32(UART0_FR) & (1 << 4)) return 0;
    return (char)(MMIO32(UART0_DR) & 0xFF);
}

void pwm_audio_init() {
    // 1. Setup GPIO 40 and 45 to Alt0 (PWM)
    // Pi 1 Model B uses 40/45. Model B+ uses 40/41.
    // We enable all of them to cover both revisions.
    gpio_set_func(40, 4); // PWM0
    gpio_set_func(41, 4); // PWM1
    gpio_set_func(45, 4); // PWM1

    // 2. Setup PWM Clock
    // Kill the clock completely first
    MMIO32(CM_PWMCTL) = CM_PASSWD | 0;
    // Wait for busy flag to clear
    while (MMIO32(CM_PWMCTL) & 0x80) delay(1);

    // Set Divider. 19.2MHz / 2 = 9.6MHz.
    // Using divisor 2 is safer for the clock generator.
    MMIO32(CM_PWMDIV) = CM_PASSWD | (2 << 12);
    
    // Start Clock (Source 1 = Oscillator 19.2MHz)
    MMIO32(CM_PWMCTL) = CM_PASSWD | 0x11; // Enable + Source 1

    // 3. Configure PWM
    // Clock is 9.6MHz. Target ~44.1kHz.
    // 9600000 / 44100 = ~218
    MMIO32(PWM_RNG1) = 218;
    MMIO32(PWM_RNG2) = 218;

    // Enable PWM0 and PWM1, Use FIFO for both
    MMIO32(PWM_CTL) = PWM_CTL_CLRF; // Clear FIFO
    delay(100);
    MMIO32(PWM_CTL) = PWM_CTL_PWEN1 | PWM_CTL_USEF1 | PWM_CTL_PWEN2 | PWM_CTL_USEF2;
}

// --- VIDEO / FRAMEBUFFER DRIVER ---
// The GPU handles composite/HDMI. We just ask for a memory buffer to draw in.

volatile uint32_t __attribute__((aligned(16))) mailbox[36];

uint32_t fb_width, fb_height, fb_pitch;
uint16_t *fb_ptr = 0; // Pointer to the screen pixels (RGB565)

void mbox_write(uint8_t channel, uint32_t value) {
    while (MMIO32(MBOX_STATUS) & 0x80000000); // Wait while full
    MMIO32(MBOX_WRITE) = (value & ~0xF) | (channel & 0xF);
}

uint32_t mbox_read(uint8_t channel) {
    while (1) {
        while (MMIO32(MBOX_STATUS) & 0x40000000); // Wait while empty
        uint32_t data = MMIO32(MBOX_READ);
        if ((data & 0xF) == channel) return data & ~0xF;
    }
}

void video_init() {
    // Property Tag Interface to GPU
    mailbox[0] = 35 * 4; // Total size
    mailbox[1] = 0;      // Request code

    mailbox[2] = 0x00048003; // Set Physical Size
    mailbox[3] = 8; mailbox[4] = 8; mailbox[5] = 640; mailbox[6] = 480;

    mailbox[7] = 0x00048004; // Set Virtual Size
    mailbox[8] = 8; mailbox[9] = 8; mailbox[10] = 640; mailbox[11] = 480;

    mailbox[12] = 0x00048005; // Set Depth
    mailbox[13] = 4; mailbox[14] = 4; mailbox[15] = 16; // 16-bit (RGB565)

    mailbox[16] = 0x00040001; // Allocate Buffer
    mailbox[17] = 8; mailbox[18] = 8; mailbox[19] = 4096; mailbox[20] = 0;

    mailbox[21] = 0x00040008; // Get Pitch
    mailbox[22] = 4; mailbox[23] = 4; mailbox[24] = 0;

    mailbox[25] = 0; // End Tag

    // Send to Mailbox Channel 8 (Tags)
    // Add 0x40000000 to address to tell GPU to bypass L2 cache (ensure coherency)
    mbox_write(8, (uint32_t)&mailbox | 0x40000000);
    mbox_read(8);

    if (mailbox[1] == 0x80000000) { // Success
        // Convert GPU bus address (0x4XXXXXXX) to ARM physical address (0x0XXXXXXX)
        fb_ptr = (uint16_t*)(mailbox[19] & 0x3FFFFFFF);
        fb_width = mailbox[5];
        fb_height = mailbox[6];
        fb_pitch = mailbox[24];
    }
}

void draw_rect(int x, int y, int w, int h, uint16_t color) {
    if (!fb_ptr) return;
    for (int j = y; j < y + h; j++) {
        if (j < 0 || j >= fb_height) continue;
        // Calculate row start: fb_ptr + (j * pitch / 2) because pitch is bytes
        uint16_t *row = (uint16_t*)((uint8_t*)fb_ptr + j * fb_pitch);
        for (int i = x; i < x + w; i++) {
            if (i >= 0 && i < fb_width) {
                row[i] = color;
            }
        }
    }
}

// --- MAIN KERNEL ---
void kernel_main(void) {
    uart_init();
    pwm_audio_init();
    video_init();

    // Draw a background
    draw_rect(0, 0, 640, 480, 0x0000); // Black
    draw_rect(10, 10, 620, 460, 0x001F); // Blue border box

    // USB Note
    // (USB requires a massive host stack, sticking to UART for input)

    // Synth State
    uint32_t phase = 0;
    uint32_t increment = 0;
    int amplitude = 0; // Use signed for decay
    
    // Frequencies (scaled for 44.1kHz sample rate and 32-bit phase accumulator)
    // Inc = (Freq * 2^32) / SampleRate
    // C4 (261.63) -> ~25466000
    const uint32_t note_C4 = 25466000;
    const uint32_t note_E4 = 32090000;
    const uint32_t note_G4 = 38160000;
    const uint32_t note_A4 = 42837000;
    const uint32_t note_C5 = 50932000;
    const uint32_t notes[] = { note_C4, note_E4, note_G4, note_A4, note_C5 };
    
    int bar_height = 0;
    uint32_t sample_count = 0;
    uint32_t next_event = 0;
    uint32_t rng = 0x12345678;

    while (1) {
        // 1. Sequencer and Envelope
        if (sample_count >= next_event) {
            rng = rng * 1664525 + 1013904223; // Simple LCG
            int note_idx = (rng >> 24) % 5;
            increment = notes[note_idx];
            amplitude = 200; // Start amplitude (max height for visualizer)
            
            // Duration: ~0.1s to 0.5s (4410 to 22050 samples)
            next_event = sample_count + 4410 + ((rng >> 10) % 17640);
        } else {
            // Simple amplitude decay
            // Decay every 200 samples (~4.5ms) to prevent clicking
            if (sample_count % 200 == 0) {
                if (amplitude > 0) {
                    amplitude--;
                } else {
                    increment = 0; // Stop note when amplitude is zero
                }
            }
        }
        sample_count++;

        // 2. Update Graphics (Simple visualizer)
        // Only update every 1000 samples (~22ms / 44fps) to keep audio smooth
        if (sample_count % 1000 == 0 && amplitude != bar_height) {
            if (amplitude > bar_height) {
                // Growing: Draw red on top
                draw_rect(300, 240 - amplitude, 40, amplitude - bar_height, 0xF800);
            } else {
                // Shrinking: Draw black on top
                draw_rect(300, 240 - bar_height, 40, bar_height - amplitude, 0x0000);
            }
            bar_height = amplitude;
        }

        // 3. Wait for FIFO space
        // The hardware consumes data at 44.1kHz. 
        // By waiting for space, we automatically sync to that speed.
        // Added timeout to prevent hang if clock fails
        int timeout = 2000;
        while ((MMIO32(PWM_STA) & PWM_STA_FULL1) && timeout--) {
            asm volatile("nop");
        }

        // 3. Generate Sample
        phase += increment;

        // Signed 8-bit sawtooth
        int8_t s_saw = (phase >> 24) - 128;

        // Apply volume and convert to unsigned for PWM
        int32_t mixed = (s_saw * amplitude) / 255;
        uint32_t sample = 128 + mixed;

        // 4. Write to FIFO
        // Write same sample to Left (PWM1) and Right (PWM2)
        MMIO32(PWM_FIFO) = sample; // Channel 1
        MMIO32(PWM_FIFO) = sample; // Channel 2
    }
}