#include <stdarg.h>
#include <inttypes.h>

#include <dtb/dtb.h>

#include <opensbi/opensbi.h>
#include <opensbi/extensions/legacy.h>
#include <opensbi/extensions/base.h>
#include <opensbi/extensions/timer.h>
#include <opensbi/extensions/ipi.h>
#include <opensbi/extensions/rfence.h>
#include <opensbi/extensions/hsm.h>
#include <opensbi/extensions/srst.h>
#include <opensbi/extensions/pmu.h>

#define read_csr(csr) \
    ({ \
        unsigned long __v; \
        asm volatile ("csrr %0, " csr : "=r"(__v) : : "memory");\
        __v;\
    })

#define write_csr(csr, v) \
    ({ \
        unsigned long __v = (v); \
        asm volatile ("csrw " csr ", %0" : : "r"(__v) : "memory");\
    })

#define set_csr_bits(csr, bits) \
    ({ \
        unsigned long __v = (bits); \
        asm volatile ("csrs " csr ", %0" : : "r"(__v) : "memory");\
    })

#define clear_csr_bits(csr, bits) \
    ({ \
        unsigned long __v = (bits); \
        asm volatile ("csrc " csr ", %0" : : "r"(__v) : "memory");\
    })

typedef unsigned long sbi_word;

void fmt(const char *f, ...) {
    namespace legacy = drivers::opensbi::legacy;

    va_list va;
    va_start(va, f);
    while(*f) {
        if(*f != '{') {
            legacy::console_putchar(*(f++));
            continue;
        }
        f++;

        while(*f != '}') {
            if(!(*f))
                return;
            f++;
        }
        f++;

        unsigned long v = va_arg(va, unsigned long);
        for(int i = 0; i < 16; ++i) {
            const char *digits = "0123456789abcdef";
            int d = (v >> (60 - i * 4)) & 0xF;
            legacy::console_putchar(digits[d]);
        }
    }
    va_end(va);
}

enum {
    pte_valid = 1 << 0,
    pte_r = 1 << 1,
    pte_w = 1 << 2,
    pte_x = 1 << 3,
    pte_user = 1 << 4,
    pte_access = 1 << 6,
    pte_dirty = 1 << 7,
};

enum {
    pte_ppn_shift = 10
};

struct pt_t {
    unsigned long ptes[512];
} __attribute__((aligned(4096)));

enum {
    satp_sv48 = 9UL << 60
};

const char *exception_strings[] = {
    "instruction misaligned",
    "instruction access fault",
    "illegal instruction",
    "breakpoint",
    "load misaligned",
    "load access fault",
    "store misaligned",
    "store access fault",
    "u-mode ecall",
    "s-mode ecall",
    "reserved (10)",
    "reserved (11)",
    "instruction page fault",
    "load page fault",
    "reserved (14)",
    "store page fault",
};

enum {
    sie_s_software = (1 << 1),
    sie_s_timer = (1 << 5)
};

enum {
    sstatus_sie = (1 << 1)
};

void cli() {
    clear_csr_bits("sstatus", sstatus_sie);
}

void sti() {
    set_csr_bits("sstatus", sstatus_sie);
}

typedef struct {
    void *sp;
} continuation_t;

struct task_t;

struct isr_frame_t {
    task_t *task;
    isr_frame_t *next;
    unsigned long ra;   // Offset 0x10.
    unsigned long a[8]; // Offset 0x18.
    unsigned long t[7]; // Offset 0x58.
    unsigned long s1;   // Offset 0x90.
    unsigned long sstatus; // Offset 0x98.
    unsigned long sepc; // Offset 0xA0.
};

enum {
    kernel_stack_size = 0x10000
};

struct task_t {
    continuation_t cont;
    isr_frame_t isr_frames[2];
    isr_frame_t *next_isr;
    char kernel_stack[kernel_stack_size];
};

task_t task1;
task_t task2;
task_t *current_task;

continuation_t prepare_stack(void *s_top, void (*mainfn)(continuation_t)) {
    void *sp = (void *)((uintptr_t)s_top - 0x78);
    *((uint64_t *)((uintptr_t)sp + 0x70)) = (uint64_t)mainfn;
    return (continuation_t){.sp = sp};
}

extern "C" void save_stack(void (*f)(void *, continuation_t), void *ctx);
extern "C" void restore_stack(continuation_t c);

void create_task(task_t *task, void (*mainfn)(continuation_t)) {
    task->isr_frames[0].task = task;
    task->isr_frames[1].task = task;
    task->isr_frames[0].next = &task->isr_frames[1];
    task->isr_frames[1].next = 0;
    task->next_isr = &task->isr_frames[0];

    task->cont = prepare_stack((void *)(task->kernel_stack + kernel_stack_size), mainfn);
}

void switch_task(void *ctx, continuation_t c) {
    task_t *task = (task_t *)ctx;
    if(task) {
        task->next_isr = (isr_frame_t *)read_csr("sscratch");
        task->cont = c;
    }

    task_t *next;
    if(task == &task1) {
        next = &task2;
    }else{
        next = &task1;
    }
    current_task = next;

    write_csr("sscratch", (unsigned long)next->next_isr);

    restore_stack(next->cont);
}

extern "C" void isr();

extern "C" void handle_isr(isr_frame_t *frame) {
    namespace legacy = drivers::opensbi::legacy;
    namespace timer = drivers::opensbi::timer;

    unsigned long cause = read_csr("scause");
    unsigned long code = cause & ~(1UL << 63);
    if(cause & (1UL << 63)) {
        if(code == 1) {
            fmt("it's an IPI\n");
            clear_csr_bits("sip", sie_s_software);
        }else if(code == 5) { // Timer interrupt.
            //clear_csr_bits("sie", sie_s_timer);
            unsigned long time = read_csr("time");
            //drivers::opensbi::ecall(0x54494D45, 0, time + 100000);
            timer::set_timer(time + 10000000);
            //set_timer(time + 100000);
            //save_stack(switch_task, current_task);
        }else{
            fmt("it's an unhandled interrupt, code: {}\n", code);
        }
    }else{
        if(code == 8) { // Syscall.
            legacy::console_putchar(frame->a[0]);
        }else{
            unsigned long sepc = read_csr("sepc");
            fmt("it's an exception at {}\n", sepc);
            const char *s = exception_strings[cause];
            while(*s)
                legacy::console_putchar(*(s++));
            while(1)
                ;
        }
    }
}

extern "C" void isr_frame_overflow() {
    fmt("isr frame overflow\n");
    while(1)
        ;
}

pt_t pml4;

extern "C" void enter_umode(void *entry);

void syscall(int n, unsigned long arg0) {
    register sbi_word rArg0 asm("a0") = arg0;
    register sbi_word rN asm("a7") = n;
    asm volatile("ecall" : "+r"(rArg0) : "r"(rN) : "memory");
    if(rArg0)
        __builtin_trap();
}

void task1_umode() {
    while(1)
        syscall(0, '!');
}

void task1_main(continuation_t c) {
    fmt("hello from task1\n");
    sti();
    //syscall(0, '!');
    //enter_umode((void *)task1_umode);
}

void task2_main(continuation_t c) {
    fmt("hello from task2\n");
    sti();
    //while(1)
    //fmt(".");
}

isr_frame_t global_isr_frames[2];

extern "C" void kmain(unsigned long hart, void *dtb) {
    (void)hart;
    (void)dtb;
    // D1 doesnt give me dtb ptr?
    // dtb = (void *)0x5fb38580;

    dtb::DeviceTree dt;
    int res = dtb::validate_and_parse_header(dt, dtb);
    if(!res) {
        fmt("validate_and_parse_header: {}\n", res);
        fmt("magic: {}\n", dt.meta.base);
        dtb::print_reserved_regions(dt);
        dtb::print_structure_block(dt);
    }

    namespace legacy = drivers::opensbi::legacy;
    namespace base = drivers::opensbi::base;
    namespace timer = drivers::opensbi::timer;
    namespace ipi = drivers::opensbi::ipi;

    fmt("test\n");

    auto ret = base::get_spec_version();
    if(ret.error) {
        fmt("get_spec_version failed\n");
    } else {
        fmt("spec version: {}\n", ret.value);
    }

    ret = base::get_impl_id();
    if(ret.error) {
        fmt("get_impl_id failed\n");
    } else {
        fmt("impl id: {}\n", ret.value);
    }

    ret = base::get_impl_version();
    if(ret.error) {
        fmt("get_impl_version failed\n");
    } else {
        fmt("impl version: {}\n", ret.value);
    }

    ret = base::get_mvendor_id();
    if(ret.error) {
        fmt("get_mvendor_id failed\n");
    } else {
        fmt("vendor id: {}\n", ret.value);
    }

    ret = base::get_march_id();
    if(ret.error) {
        fmt("get_march_id failed\n");
    } else {
        fmt("arch id: {}\n", ret.value);
    }

    ret = base::get_mimp_id();
    if(ret.error) {
        fmt("get_mimp_id failed\n");
    } else {
        fmt("imp id: {}\n", ret.value);
    }


    fmt("uintmax: {} bits\n", sizeof(uintmax_t) * 8);
    fmt("HART: {}, DTB: {}\n", hart, dtb);

    volatile uint32_t *cfg = (uint32_t *)0x2000060;
    volatile uint32_t *dat = (uint32_t *)0x2000070;

    *cfg = 0x10;
    *dat = 0xFF;

    //while(1);

    cli();

    global_isr_frames[0].next = &global_isr_frames[1];
    write_csr("sscratch", (unsigned long)&global_isr_frames[0]);
    write_csr("stvec", ((unsigned long)&isr));

    const char *s;

    set_csr_bits("sie", sie_s_software | sie_s_timer);

    // Self-IPI.
    ipi::send_ipi({ 1UL << hart, 0 });

    sti();
    unsigned long timea = read_csr("time");
    ret = timer::set_timer(timea + 100000);

    drivers::opensbi::rfence::remote_fence_i({0, 0});
    //ret = drivers::opensbi::hsm::hart_stop();
    using drivers::opensbi::srst::ResetType;
    using drivers::opensbi::srst::ResetReason;
    // ret = drivers::opensbi::srst::system_reset(ResetType::WARM_REBOOT, ResetReason::SYSTEM_FAILURE);

    namespace pmu = drivers::opensbi::pmu;
    ret = pmu::num_counters();

    fmt("Number of counters: {}\n", ret.value);
    for(int i = 0; i < ret.value; i++) {
        ret = pmu::counter_get_info(i);
        fmt("{} err: {}\n", i, ret.error);
        fmt("{} val: {}\n", i, ret.value);
    }
    while(1);
    cli();

    unsigned long time;
    //volatile uint32_t *meme = (uint32_t *)0x6011010;
    //while(true) {
    time = read_csr("time");
    fmt("hello world {}\n", time);
    //*meme = *meme | (0xA57 << 1) | (1 << 0);
    //}

    // Identity map the first 512GiB.
    pml4.ptes[0] = ((0UL >> 12) << pte_ppn_shift) | pte_valid
            | pte_r | pte_w | pte_x | pte_access | pte_dirty | pte_user;
    // Map 512GiB -> 0.
    pml4.ptes[1] = ((0UL >> 12) << pte_ppn_shift) | pte_valid
            | pte_r | pte_w | pte_x | pte_access | pte_dirty | pte_user;

    fmt("Writing paging memes\n");
    unsigned long pml4p = (unsigned long)&pml4;
    fmt("pml4p @ {}\n", pml4p);
    fmt("Setting SStatus\n");
    set_csr_bits("sstatus", 1 << 18);
    fmt("Setting Satp\n");
    write_csr("satp", (pml4p >> 12) | satp_sv48);

    fmt("Paging might work?...\n");

    s = "paging works!\n";
    s += 512UL * 1024 * 1024 * 1024;
    while(*s)
        legacy::console_putchar(*(s++));

    create_task(&task1, task1_main);
    create_task(&task2, task2_main);

    time = read_csr("time");
    //set_timer(time + 100000);
    //drivers::opensbi::ecall(0x54494D45, 0, time + 100000);
    //timer::set_timer(time + 100000);

    //save_stack(switch_task, 0);

    namespace hsm = drivers::opensbi::hsm;
    hsm::hart_stop();
    // drivers::opensbi::ecall(8, 0, 0);

    while(1);
}