//! Console device interfaces

use core::time::Duration;

use abi::{error::Error, io::TerminalSize, primitive_enum};
use alloc::{vec, vec::Vec};
use bitflags::bitflags;
use libk::{task::runtime, vfs::TerminalOutput};
use libk_util::{sync::IrqSafeSpinlock, StaticVector};

use crate::debug::DebugSink;

const CONSOLE_ROW_LEN: usize = 80;
const MAX_CSI_ARGS: usize = 8;

const DEFAULT_FG_COLOR: ColorAttribute = ColorAttribute::White;
const DEFAULT_BG_COLOR: ColorAttribute = ColorAttribute::Blue;

primitive_enum! {
    #[allow(missing_docs)]
    #[doc = "Color attribute of a console character"]
    pub enum ColorAttribute: u8 {
        Black = 0,
        Red = 1,
        Green = 2,
        Yellow = 3,
        Blue = 4,
        Magenta = 5,
        Cyan = 6,
        White = 7,
    }
}

bitflags! {
    #[doc = "Extra attributes of a console character"]
    #[derive(Clone, Copy)]
    pub struct Attributes: u8 {
        #[allow(missing_docs)]
        const BOLD = 1 << 0;
    }
}

impl ColorAttribute {
    fn from_vt100(val: u8) -> Self {
        match val {
            0..=7 => Self::try_from(val).unwrap(),
            _ => ColorAttribute::Red,
        }
    }

    /// Converts the attribute to RGBA representation
    pub fn as_rgba(&self, bold: bool) -> u32 {
        let color = match self {
            Self::Black => 0x000000,
            Self::Red => 0x7F0000,
            Self::Green => 0x007F00,
            Self::Yellow => 0x7F7F00,
            Self::Blue => 0x00007F,
            Self::Magenta => 0x7F007F,
            Self::Cyan => 0x007F7F,
            Self::White => 0x7F7F7F,
        };
        if bold {
            color * 2
        } else {
            color
        }
    }
}

/// Represents a single character with its attributes
#[derive(Clone, Copy)]
#[repr(C)]
pub struct ConsoleChar {
    attributes: u16,
    char: u8,
    _pad: u8,
}

/// Represents a single line in the console buffer
#[derive(Clone, Copy)]
pub struct ConsoleRow {
    dirty: u8,
    chars: [ConsoleChar; CONSOLE_ROW_LEN],
}

/// Buffer that contains text rows of the console with their attributes + tracks dirty rows which
/// need to be flushed to the display
pub struct ConsoleBuffer {
    rows: Vec<ConsoleRow>,
    height: u32,
}

/// Console wrapper for adding it into devfs as a char device
pub struct ConsoleWrapper<'a>(pub &'a dyn DisplayConsole);

enum EscapeState {
    Normal,
    Escape,
    Csi,
}

/// Common state for console output devices
pub struct ConsoleState {
    /// Current cursor row
    pub cursor_row: u32,
    /// Current cursor column
    pub cursor_col: u32,
    /// Current foreground color
    pub fg_color: ColorAttribute,
    /// Current background color
    pub bg_color: ColorAttribute,
    /// Current set of attributes
    pub attributes: Attributes,

    esc_args: StaticVector<u32, MAX_CSI_ARGS>,
    esc_state: EscapeState,

    /// Row buffer
    pub buffer: ConsoleBuffer,
}

/// Helper type to iterate over dirty rows in the buffer
pub struct RowIter<'a> {
    buffer: &'a mut ConsoleBuffer,
    index: u32,
}

/// Interface to implement buffered console semantics on an abstract console output device
pub trait DisplayConsole: Sync {
    /// Returns the state lock
    fn state(&self) -> &IrqSafeSpinlock<ConsoleState>;

    /// Flushes the data from console buffer to the display
    fn flush(&self, state: &mut ConsoleState);

    /// Writes characters to the backing buffer + handles special control sequences
    fn write_char(&self, c: u8) {
        let mut state = self.state().lock();
        if state.putc(c) {
            self.flush(&mut state);
        }
    }

    /// Returns the dimensions of the console in chars: (rows, columns)
    fn text_dimensions(&self) -> (usize, usize) {
        let state = self.state().lock();
        (state.buffer.height as _, CONSOLE_ROW_LEN as _)
    }
}

impl ConsoleChar {
    /// Empty character
    pub const BLANK: Self = Self {
        attributes: 0,
        char: 0,
        _pad: 0,
    };

    /// Constructs a console character from a char and its attributes
    #[inline(always)]
    pub fn from_parts(
        char: u8,
        fg: ColorAttribute,
        bg: ColorAttribute,
        attributes: Attributes,
    ) -> Self {
        let attributes = ((attributes.bits() as u16) << 8)
            | ((u8::from(bg) as u16) << 4)
            | (u8::from(fg) as u16);
        Self {
            attributes,
            char,
            _pad: 0,
        }
    }

    /// Returns the attributes of the character
    #[inline(always)]
    pub fn attributes(self) -> (ColorAttribute, ColorAttribute, Attributes) {
        let fg =
            ColorAttribute::try_from((self.attributes & 0xF) as u8).unwrap_or(DEFAULT_FG_COLOR);
        let bg = ColorAttribute::try_from(((self.attributes >> 4) & 0xF) as u8)
            .unwrap_or(DEFAULT_BG_COLOR);
        let attributes =
            Attributes::from_bits((self.attributes >> 8) as u8).unwrap_or(Attributes::empty());

        (fg, bg, attributes)
    }

    /// Returns the character data of this [ConsoleChar]
    #[inline(always)]
    pub const fn character(self) -> u8 {
        self.char
    }
}

impl<'a> RowIter<'a> {
    /// Returns the next dirty row
    pub fn next_dirty(&mut self) -> Option<(u32, &[ConsoleChar])> {
        loop {
            if self.index == self.buffer.height {
                return None;
            }

            if !self.buffer.rows[self.index as usize].clear_dirty() {
                self.index += 1;
                continue;
            }

            let row_index = self.index;
            let row = &self.buffer.rows[self.index as usize];

            self.index += 1;

            return Some((row_index, &row.chars));
        }
    }
}

impl ConsoleRow {
    /// Constructs a row filled with blank characters
    pub const fn zeroed() -> Self {
        Self {
            dirty: 1,
            chars: [ConsoleChar {
                attributes: ((DEFAULT_BG_COLOR as u8) as u16) << 4,
                char: b' ',
                _pad: 0,
            }; CONSOLE_ROW_LEN],
        }
    }

    /// Returns `true` if the row's dirty flag is set
    #[inline]
    pub const fn is_dirty(&self) -> bool {
        self.dirty != 0
    }

    /// Clears "dirty" flag for the row
    #[inline]
    pub fn clear_dirty(&mut self) -> bool {
        let old = self.dirty;
        self.dirty = 0;
        old == 1
    }

    /// Clears the console row with blank characters
    pub fn clear(&mut self, bg: ColorAttribute) {
        self.dirty = 1;
        self.chars
            .fill(ConsoleChar::from_parts(b' ', bg, bg, Attributes::empty()));
    }
}

impl ConsoleBuffer {
    /// Constructs a fixed-size console buffer
    pub fn new(height: u32) -> Result<Self, Error> {
        // let size = size_of::<ConsoleRow>() * (height as usize);
        let mut rows = vec![ConsoleRow::zeroed(); height as usize];

        for row in rows.iter_mut() {
            row.clear(DEFAULT_BG_COLOR);
        }

        Ok(Self { rows, height })
        // let size = size_of::<ConsoleRow>() * (height as usize);
        // let page_count = (size + 0xFFF) / 0x1000;
        // let pages = phys::alloc_pages_contiguous(page_count, PageUsage::Used)?;

        // let rows = unsafe {
        //     core::slice::from_raw_parts_mut(pages.virtualize() as *mut ConsoleRow, height as usize)
        // };

        // for row in rows.iter_mut() {
        //     row.clear(DEFAULT_BG_COLOR);
        // }

        // Ok(Self { rows, height })
    }

    #[inline(never)]
    fn set_char(&mut self, row: u32, col: u32, c: ConsoleChar) {
        self.rows[row as usize].dirty = 1;
        self.rows[row as usize].chars[col as usize] = c;
    }

    #[inline(never)]
    fn set_dirty(&mut self, row: u32) {
        self.rows[row as usize].dirty = 1;
    }

    /// Returns an iterator over dirty rows, while clearing dirty flag for them
    pub fn flush_rows(&mut self) -> RowIter {
        RowIter {
            buffer: self,
            index: 0,
        }
    }

    fn clear(&mut self, bg: ColorAttribute) {
        for row in self.rows.iter_mut() {
            row.clear(bg);
        }
    }

    fn clear_row(&mut self, row: u32, bg: ColorAttribute) {
        self.rows[row as usize].dirty = 1;
        self.rows[row as usize].clear(bg);
    }

    fn erase_in_row(&mut self, row: u32, start: usize, bg: ColorAttribute) {
        self.rows[row as usize].dirty = 1;
        self.rows[row as usize].chars[start..].fill(ConsoleChar::from_parts(
            b' ',
            DEFAULT_FG_COLOR,
            bg,
            Attributes::empty(),
        ));
    }

    fn scroll_once(&mut self, bg: ColorAttribute) {
        self.rows.copy_within(1.., 0);
        self.rows[(self.height - 1) as usize].clear(bg);

        // Mark everything dirty
        self.rows.iter_mut().for_each(|row| {
            row.dirty = 1;
        });
    }
}

impl ConsoleState {
    /// Constructs a new console state with given buffer
    pub fn new(buffer: ConsoleBuffer) -> Self {
        Self {
            cursor_row: 0,
            cursor_col: 0,
            fg_color: DEFAULT_FG_COLOR,
            bg_color: DEFAULT_BG_COLOR,
            attributes: Attributes::empty(),

            esc_args: StaticVector::new(),
            esc_state: EscapeState::Normal,

            buffer,
        }
    }

    fn putc_normal(&mut self, c: u8) -> bool {
        let mut flush = false;

        match c {
            c if c >= 127 => {
                self.buffer.set_char(
                    self.cursor_row,
                    self.cursor_col,
                    ConsoleChar::from_parts(
                        b'?',
                        self.fg_color,
                        ColorAttribute::Red,
                        self.attributes,
                    ),
                );

                self.cursor_col += 1;
            }
            b'\x1b' => {
                self.esc_state = EscapeState::Escape;
                return false;
            }
            b'\r' => {
                self.cursor_col = 0;
            }
            b'\n' => {
                self.cursor_row += 1;
                self.cursor_col = 0;
                flush = true;
            }
            _ => {
                self.buffer.set_char(
                    self.cursor_row,
                    self.cursor_col,
                    ConsoleChar::from_parts(c, self.fg_color, self.bg_color, self.attributes),
                );

                self.cursor_col += 1;
            }
        }

        if self.cursor_col == CONSOLE_ROW_LEN as u32 {
            self.cursor_col = 0;
            self.cursor_row += 1;
        }

        if self.cursor_row == self.buffer.height {
            self.buffer.scroll_once(self.bg_color);
            self.cursor_row = self.buffer.height - 1;
            flush = true;
        }

        flush
    }

    fn handle_csi(&mut self, c: u8) -> bool {
        match c {
            // Move back one character
            b'D' => {
                if self.cursor_col > 0 {
                    self.cursor_col -= 1;
                }
            }
            // Manipulate display attributes
            b'm' => {
                if let Some(arg) = self.esc_args.first() {
                    match arg {
                        // Reset
                        0 => {
                            self.fg_color = DEFAULT_FG_COLOR;
                            self.bg_color = DEFAULT_BG_COLOR;
                            self.attributes = Attributes::empty();
                        }
                        // Bold
                        1 => {
                            self.attributes |= Attributes::BOLD;
                        }
                        // Foreground colors
                        30..=39 => {
                            let vt_color = self.esc_args[0] % 10;
                            if vt_color == 9 {
                                self.fg_color = DEFAULT_FG_COLOR;
                            } else {
                                self.fg_color = ColorAttribute::from_vt100(vt_color as u8);
                            }
                        }
                        // Background colors
                        40..=49 => {
                            let vt_color = self.esc_args[0] % 10;
                            if vt_color == 9 {
                                self.bg_color = DEFAULT_BG_COLOR;
                            } else {
                                self.bg_color = ColorAttribute::from_vt100(vt_color as u8);
                            }
                        }
                        _ => (),
                    }
                }
            }
            // Move cursor to position
            b'f' => {
                let row = self.esc_args[0].clamp(1, self.buffer.height) - 1;
                let col = self.esc_args[1].clamp(1, CONSOLE_ROW_LEN as u32) - 1;

                self.buffer.set_dirty(row);

                self.cursor_row = row;
                self.cursor_col = col;
            }
            // Clear rows/columns/screen
            b'J' => match self.esc_args[0] {
                // Erase lines down
                0 => (),
                // Erase lines up
                1 => (),
                // Erase all
                2 => {
                    self.buffer.clear(self.bg_color);
                }
                _ => (),
            },
            // Erase in Line
            b'K' => match self.esc_args[0] {
                // Erase to Right
                0 => {
                    self.buffer
                        .erase_in_row(self.cursor_row, self.cursor_col as _, self.bg_color);
                }
                // Erase All
                2 => {
                    self.buffer.clear_row(self.cursor_row, self.bg_color);
                }
                _ => (),
            },
            _ => (),
        }

        self.esc_state = EscapeState::Normal;
        false
    }

    fn handle_csi_byte(&mut self, c: u8) -> bool {
        match c {
            b'0'..=b'9' => {
                let arg = self.esc_args.last_mut().unwrap();
                *arg *= 10;
                *arg += (c - b'0') as u32;
                false
            }
            b';' => {
                self.esc_args.push(0);
                false
            }
            _ => self.handle_csi(c),
        }
    }

    fn putc(&mut self, c: u8) -> bool {
        match self.esc_state {
            EscapeState::Normal => self.putc_normal(c),
            EscapeState::Escape => match c {
                b'[' => {
                    self.esc_state = EscapeState::Csi;
                    self.esc_args.clear();
                    self.esc_args.push(0);
                    false
                }
                _ => {
                    self.esc_state = EscapeState::Normal;
                    self.esc_args.clear();
                    false
                }
            },
            EscapeState::Csi => self.handle_csi_byte(c),
        }
    }
}

impl DebugSink for dyn DisplayConsole {
    fn putc(&self, c: u8) -> Result<(), Error> {
        self.write_char(c);
        Ok(())
    }

    fn supports_control_sequences(&self) -> bool {
        true
    }
}

impl TerminalOutput for ConsoleWrapper<'_> {
    fn size(&self) -> TerminalSize {
        let (rows, columns) = self.0.text_dimensions();
        TerminalSize { rows, columns }
    }

    fn write(&self, byte: u8) -> Result<(), Error> {
        self.0.write_char(byte);
        Ok(())
    }
}

static CONSOLES: IrqSafeSpinlock<Vec<&'static dyn DisplayConsole>> =
    IrqSafeSpinlock::new(Vec::new());

/// Adds a console device to a auto-flush list
pub fn add_console_autoflush(console: &'static dyn DisplayConsole) {
    CONSOLES.lock().push(console);
}

/// Flushes console buffers to their displays
pub fn flush_consoles() {
    for console in CONSOLES.lock().iter() {
        let mut state = console.state().lock();
        console.flush(&mut state);
    }
}

/// Periodically flushes data from console buffers onto their displays
pub async fn update_consoles_task() {
    loop {
        flush_consoles();

        runtime::sleep(Duration::from_millis(20)).await;
    }
}