#![no_std]
#![no_main]

mod ansi;
mod font;

use ansi::{Action, Parser};
use lancer_user::show;
use lancer_user::syscall;

const FB_CAP_SLOT: u64 = 2;
const CONSOLE_RING_BASE_SLOT: u64 = 64;
const CONSOLE_RING_MAX_FRAMES: u64 = 16;
const FB_MAP_VADDR: u64 = 0x4000_0000_0000;
const CONSOLE_RING_VADDR: u64 = 0x5000_0000_0000;

const DEFAULT_FG: u32 = 0x00CC_CCCC;
const DEFAULT_BG: u32 = 0x0000_0000;

struct Console {
    fb_base: *mut u8,
    fb_pitch: u32,
    fb_byte_size: usize,
    col: usize,
    row: usize,
    max_cols: usize,
    max_rows: usize,
    fg: u32,
    bg: u32,
    bold: bool,
    saved_col: usize,
    saved_row: usize,
    scroll_top: usize,
    scroll_bottom: usize,
    parser: Parser,
}

impl Console {
    fn new(base: *mut u8, pitch: u32, byte_size: usize, max_cols: usize, max_rows: usize) -> Self {
        Self {
            fb_base: base,
            fb_pitch: pitch,
            fb_byte_size: byte_size,
            col: 0,
            row: 0,
            max_cols,
            max_rows,
            fg: DEFAULT_FG,
            bg: DEFAULT_BG,
            bold: false,
            saved_col: 0,
            saved_row: 0,
            scroll_top: 0,
            scroll_bottom: max_rows.saturating_sub(1),
            parser: Parser::new(),
        }
    }

    fn put_glyph(&self, col: usize, row: usize, ch: u8, fg: u32, bg: u32) {
        let glyph_offset = (ch as usize) * 16;
        let pixel_x = col * 8;
        let pixel_y = row * 16;

        let last_byte =
            ((pixel_y + 15) as u64) * (self.fb_pitch as u64) + ((pixel_x + 7) as u64) * 4 + 4;
        if last_byte as usize > self.fb_byte_size {
            return;
        }

        (0u32..16).fold((), |(), glyph_row| {
            let bits = font::FONT_8X16[glyph_offset + glyph_row as usize];
            let row_offset = ((pixel_y as u32 + glyph_row) as u64) * (self.fb_pitch as u64);
            let col_base = (pixel_x as u64) * 4;

            (0u32..8).fold((), |(), bit| {
                let color = match bits & (0x80 >> bit) != 0 {
                    true => fg,
                    false => bg,
                };
                let offset = row_offset + col_base + (bit as u64) * 4;
                unsafe {
                    core::ptr::write_volatile(self.fb_base.add(offset as usize) as *mut u32, color);
                }
            });
        });
    }

    fn effective_fg(&self) -> u32 {
        match self.bold {
            true => brighten(self.fg),
            false => self.fg,
        }
    }

    fn scroll_region_up(&mut self, top: usize, bottom: usize, count: usize) {
        let row_bytes = (self.fb_pitch as usize) * 16;
        let region_rows = bottom.saturating_sub(top) + 1;
        if count >= region_rows {
            self.clear_rows(top, bottom);
            return;
        }

        let src_start = (top + count) * row_bytes;
        let dst_start = top * row_bytes;
        let copy_rows = region_rows - count;
        let copy_size = copy_rows * row_bytes;

        if src_start + copy_size > self.fb_byte_size {
            return;
        }

        unsafe {
            core::ptr::copy(
                self.fb_base.add(src_start),
                self.fb_base.add(dst_start),
                copy_size,
            );
        }

        let clear_start = (bottom + 1 - count) * row_bytes;
        let clear_size = count * row_bytes;
        if clear_start + clear_size <= self.fb_byte_size {
            unsafe {
                core::ptr::write_bytes(self.fb_base.add(clear_start), 0u8, clear_size);
            }
        }
    }

    fn scroll_region_down(&mut self, top: usize, bottom: usize, count: usize) {
        let row_bytes = (self.fb_pitch as usize) * 16;
        let region_rows = bottom.saturating_sub(top) + 1;
        if count >= region_rows {
            self.clear_rows(top, bottom);
            return;
        }

        let src_start = top * row_bytes;
        let dst_start = (top + count) * row_bytes;
        let copy_rows = region_rows - count;
        let copy_size = copy_rows * row_bytes;

        if dst_start + copy_size > self.fb_byte_size {
            return;
        }

        unsafe {
            core::ptr::copy(
                self.fb_base.add(src_start),
                self.fb_base.add(dst_start),
                copy_size,
            );
        }

        let clear_size = count * row_bytes;
        if src_start + clear_size <= self.fb_byte_size {
            unsafe {
                core::ptr::write_bytes(self.fb_base.add(src_start), 0u8, clear_size);
            }
        }
    }

    fn clear_rows(&mut self, top: usize, bottom: usize) {
        let row_bytes = (self.fb_pitch as usize) * 16;
        let start = top * row_bytes;
        let end = (bottom + 1) * row_bytes;
        if end <= self.fb_byte_size {
            unsafe {
                core::ptr::write_bytes(self.fb_base.add(start), 0u8, end - start);
            }
        }
    }

    fn clear_row_cols(&self, row: usize, from_col: usize, to_col: usize) {
        (from_col..to_col).fold((), |(), col| {
            self.put_glyph(col, row, b' ', DEFAULT_FG, DEFAULT_BG);
        });
    }

    fn newline(&mut self) {
        self.col = 0;
        self.row += 1;
        if self.row > self.scroll_bottom {
            self.scroll_region_up(self.scroll_top, self.scroll_bottom, 1);
            self.row = self.scroll_bottom;
        }
    }

    fn advance_cursor(&mut self) {
        self.col += 1;
        if self.col >= self.max_cols {
            self.newline();
        }
    }

    fn handle_action(&mut self, action: Action) {
        match action {
            Action::Print(ch) => {
                let fg = self.effective_fg();
                let bg = self.bg;
                self.put_glyph(self.col, self.row, ch, fg, bg);
                self.advance_cursor();
            }
            Action::Newline => self.newline(),
            Action::CarriageReturn => {
                self.col = 0;
            }
            Action::Backspace => {
                if self.col > 0 {
                    self.col -= 1;
                    self.put_glyph(self.col, self.row, b' ', DEFAULT_FG, DEFAULT_BG);
                }
            }
            Action::Tab => {
                let target = (self.col + 8) & !7;
                let target = target.min(self.max_cols.saturating_sub(1));
                self.col = target;
            }
            Action::SetFg(color) => {
                self.fg = color;
            }
            Action::SetBg(color) => {
                self.bg = color;
            }
            Action::SetBold => {
                self.bold = true;
            }
            Action::SetDim => {
                self.bold = false;
            }
            Action::ResetStyle => {
                self.fg = DEFAULT_FG;
                self.bg = DEFAULT_BG;
                self.bold = false;
            }
            Action::DefaultFg => {
                self.fg = DEFAULT_FG;
            }
            Action::DefaultBg => {
                self.bg = DEFAULT_BG;
            }
            Action::CursorTo { row, col } => {
                self.row = (row as usize)
                    .saturating_sub(1)
                    .min(self.max_rows.saturating_sub(1));
                self.col = (col as usize)
                    .saturating_sub(1)
                    .min(self.max_cols.saturating_sub(1));
            }
            Action::CursorUp(n) => {
                self.row = self.row.saturating_sub(n as usize);
            }
            Action::CursorDown(n) => {
                self.row = (self.row + n as usize).min(self.max_rows.saturating_sub(1));
            }
            Action::CursorForward(n) => {
                self.col = (self.col + n as usize).min(self.max_cols.saturating_sub(1));
            }
            Action::CursorBack(n) => {
                self.col = self.col.saturating_sub(n as usize);
            }
            Action::CursorColumn(n) => {
                self.col = (n as usize)
                    .saturating_sub(1)
                    .min(self.max_cols.saturating_sub(1));
            }
            Action::SaveCursor => {
                self.saved_col = self.col;
                self.saved_row = self.row;
            }
            Action::RestoreCursor => {
                self.col = self.saved_col;
                self.row = self.saved_row;
            }
            Action::HideCursor | Action::ShowCursor => {}
            Action::ClearToEnd => {
                self.clear_row_cols(self.row, self.col, self.max_cols);
                if self.row + 1 < self.max_rows {
                    self.clear_rows(self.row + 1, self.max_rows - 1);
                }
            }
            Action::ClearFromStart => {
                if self.row > 0 {
                    self.clear_rows(0, self.row - 1);
                }
                self.clear_row_cols(self.row, 0, self.col + 1);
            }
            Action::ClearScreen => {
                self.clear_rows(0, self.max_rows.saturating_sub(1));
                self.col = 0;
                self.row = 0;
            }
            Action::ClearLineToEnd => {
                self.clear_row_cols(self.row, self.col, self.max_cols);
            }
            Action::ClearLineFromStart => {
                self.clear_row_cols(self.row, 0, self.col + 1);
            }
            Action::ClearLine => {
                self.clear_row_cols(self.row, 0, self.max_cols);
            }
            Action::SetScrollRegion { top, bottom } => {
                let t = (top as usize)
                    .saturating_sub(1)
                    .min(self.max_rows.saturating_sub(1));
                let b = (bottom as usize)
                    .saturating_sub(1)
                    .min(self.max_rows.saturating_sub(1));
                if t < b {
                    self.scroll_top = t;
                    self.scroll_bottom = b;
                    self.col = 0;
                    self.row = 0;
                }
            }
            Action::ResetScrollRegion => {
                self.scroll_top = 0;
                self.scroll_bottom = self.max_rows.saturating_sub(1);
            }
            Action::ScrollUp(n) => {
                self.scroll_region_up(self.scroll_top, self.scroll_bottom, n as usize);
            }
            Action::ScrollDown(n) => {
                self.scroll_region_down(self.scroll_top, self.scroll_bottom, n as usize);
            }
        }
    }

    fn process_byte(&mut self, byte: u8) {
        let mut action_buf = [None; 4];
        let mut action_count = 0;

        self.parser.feed(byte, |action| {
            if action_count < action_buf.len() {
                action_buf[action_count] = Some(action);
                action_count += 1;
            }
        });

        (0..action_count).fold((), |(), i| {
            if let Some(action) = action_buf[i] {
                self.handle_action(action);
            }
        });
    }
}

fn brighten(color: u32) -> u32 {
    let r = ((color >> 16) & 0xFF).min(204) + 51;
    let g = ((color >> 8) & 0xFF).min(204) + 51;
    let b = (color & 0xFF).min(204) + 51;
    (r << 16) | (g << 8) | b
}

struct ConsoleRing {
    base: *const u8,
    capacity: u32,
    last_tail: u32,
}

impl ConsoleRing {
    fn attach(vaddr: u64) -> Self {
        let base = vaddr as *const u8;
        let capacity = unsafe { (base.add(8) as *const u32).read_volatile() };
        let write_head = unsafe { (base as *const u32).read_volatile() };
        Self {
            base,
            capacity,
            last_tail: write_head,
        }
    }

    fn drain(&mut self, mut handler: impl FnMut(u8)) {
        let write_head = unsafe { (self.base as *const u32).read_volatile() };
        let available = write_head.wrapping_sub(self.last_tail);

        let effective_tail = match available > self.capacity {
            true => write_head.wrapping_sub(self.capacity),
            false => self.last_tail,
        };

        let data_base = unsafe { self.base.add(16) };
        let count = write_head.wrapping_sub(effective_tail);

        (0..count).fold(effective_tail, |tail, _| {
            let idx = tail % self.capacity;
            let byte = unsafe { data_base.add(idx as usize).read_volatile() };
            handler(byte);
            tail.wrapping_add(1)
        });

        self.last_tail = write_head;
        unsafe {
            (self.base as *mut u32).add(1).write_volatile(write_head);
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lancer_main() -> ! {
    let info = match syscall::fb_info(FB_CAP_SLOT) {
        Some(i) => i,
        None => syscall::exit(),
    };

    let mapped_pages = syscall::fb_map(FB_CAP_SLOT, FB_MAP_VADDR);
    if mapped_pages < 0 {
        syscall::exit();
    }
    let mapped_bytes = (mapped_pages as usize) * 4096;

    let byte_size = (info.byte_size as usize).min(mapped_bytes);
    let pitch = info.pitch as u32;
    let row_pixel_bytes = (pitch as usize) * 16;

    unsafe { core::ptr::write_bytes(FB_MAP_VADDR as *mut u8, 0u8, byte_size) };

    show!(
        fbcon,
        "{}x{} pitch {} byte_size {}",
        info.width,
        info.height,
        info.pitch,
        info.byte_size
    );

    let max_cols = (info.width as usize) / 8;
    let max_rows = match row_pixel_bytes {
        0 => 0,
        rpb => byte_size / rpb,
    };

    let mut console = Console::new(
        FB_MAP_VADDR as *mut u8,
        pitch,
        byte_size,
        max_cols,
        max_rows,
    );

    let ring_frames = (0..CONSOLE_RING_MAX_FRAMES)
        .take_while(|&i| {
            syscall::frame_map(
                CONSOLE_RING_BASE_SLOT + i,
                CONSOLE_RING_VADDR + i * 4096,
                1,
            ) >= 0
        })
        .count();
    if ring_frames == 0 {
        show!(fbcon, error, "console ring map failed, exiting");
        syscall::exit();
    }

    let mut ring = ConsoleRing::attach(CONSOLE_RING_VADDR);

    ring.drain(|b| console.process_byte(b));

    loop {
        ring.drain(|b| console.process_byte(b));
        syscall::sched_yield();
    }
}