use core::{any::Any, fmt, future::poll_fn, task::Poll};

use alloc::sync::Arc;
use libk_util::{sync::spin_rwlock::IrqSafeRwLock, waker::QueueWaker};
use yggdrasil_abi::{
    error::Error,
    io::{FileMode, OpenOptions},
};

use crate::{
    task::sync::AsyncMutex,
    vfs::{CommonImpl, Filename, InstanceData, Metadata, Node, NodeFlags, NodeRef, RegularImpl},
};

const RECORD_SIZE: usize = 128;
const FIXED_LOG_SIZE: usize = 256;

struct LogFile;
pub(super) struct RingLoggerSink;

struct WriteBuffer {
    text: [u8; 255],
    len: usize,
}

struct Record {
    len: usize,
    level: log::Level,
}

impl fmt::Write for WriteBuffer {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        let bytes = s.as_bytes();
        let available = self.text.len() - self.len;
        let amount = available.min(bytes.len());
        self.text[self.len..self.len + amount].copy_from_slice(&bytes[..amount]);
        self.len += amount;
        Ok(())
    }
}

static mut FIXED_LOG_BUFFER: [u8; FIXED_LOG_SIZE * RECORD_SIZE] = [0; FIXED_LOG_SIZE * RECORD_SIZE];
static mut FIXED_LOG_RECORDS: [Record; FIXED_LOG_SIZE] = [const {
    Record {
        len: 0,
        level: log::Level::Trace,
    }
}; FIXED_LOG_SIZE];
#[allow(static_mut_refs)]
static LOG_BUFFER: IrqSafeRwLock<LogBuffer> = IrqSafeRwLock::new(unsafe {
    LogBuffer::new(LogBufferStorage::Fixed {
        meta: &mut FIXED_LOG_RECORDS,
        data: &mut FIXED_LOG_BUFFER,
    })
});
static LOG_NOTIFY: QueueWaker = QueueWaker::new();

enum LogBufferStorage {
    Fixed {
        meta: &'static mut [Record; FIXED_LOG_SIZE],
        data: &'static mut [u8; FIXED_LOG_SIZE * RECORD_SIZE],
    },
}

struct LogBuffer {
    storage: LogBufferStorage,
    head: usize,
    tail: usize,
    filled: bool,
}

struct OpenLogFile {
    seq: AsyncMutex<usize>,
}

impl LogBufferStorage {
    fn record_mut(&mut self, index: usize) -> (&mut Record, &mut [u8]) {
        match self {
            Self::Fixed { meta, data } => {
                let index = index % meta.len();
                (&mut meta[index], &mut data[index * RECORD_SIZE..])
            }
        }
    }

    fn record(&self, index: usize) -> (&Record, &[u8]) {
        match self {
            Self::Fixed { meta, data } => {
                let index = index % meta.len();
                (&meta[index], &data[index * RECORD_SIZE..])
            }
        }
    }

    fn len(&self) -> usize {
        match self {
            Self::Fixed { meta, .. } => meta.len(),
        }
    }
}

impl LogBuffer {
    const fn new(storage: LogBufferStorage) -> Self {
        Self {
            storage,
            head: 0,
            tail: 0,
            filled: false,
        }
    }

    fn write(&mut self, level: log::Level, text: &[u8]) {
        let len = text.len().min(RECORD_SIZE);
        let (meta, data) = self.storage.record_mut(self.head);
        meta.len = len;
        meta.level = level;
        data[..len].copy_from_slice(&text[..len]);
        self.head = self.head.wrapping_add(1);
        if self.filled {
            self.tail = self.head.wrapping_sub(self.storage.len() - 1);
        } else if self.head == self.storage.len() {
            self.filled = true;
        }
    }
}

impl CommonImpl for LogFile {
    fn metadata(&self, _node: &NodeRef) -> Result<Metadata, Error> {
        unreachable!()
    }

    fn size(&self, _node: &NodeRef) -> Result<u64, Error> {
        Ok(0)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

impl RegularImpl for LogFile {
    fn open(
        &self,
        _node: &NodeRef,
        opts: OpenOptions,
    ) -> Result<(u64, Option<InstanceData>), Error> {
        if opts.contains_any(OpenOptions::WRITE | OpenOptions::APPEND | OpenOptions::TRUNCATE) {
            return Err(Error::ReadOnly);
        }
        let instance = Arc::new(OpenLogFile {
            seq: AsyncMutex::new(LOG_BUFFER.read().tail),
        });
        Ok((0, Some(instance)))
    }

    fn read(
        &self,
        _node: &NodeRef,
        instance: Option<&InstanceData>,
        _pos: u64,
        buf: &mut [u8],
    ) -> Result<usize, Error> {
        if buf.len() < 4 {
            return Err(Error::InvalidArgument);
        }

        let instance: &OpenLogFile = instance
            .and_then(|t| t.downcast_ref())
            .ok_or(Error::InvalidFile)?;

        block! {
            let mut seq = instance.seq.lock().await;
            let (len, new_seq) = read_log_line(*seq, buf).await;
            *seq = new_seq;
            Ok(len)
        }?
    }

    fn close(&self, _node: &NodeRef, _instance: Option<&InstanceData>) -> Result<(), Error> {
        Ok(())
    }
}

impl log::Log for RingLoggerSink {
    fn enabled(&self, metadata: &log::Metadata) -> bool {
        metadata.level() != log::Level::Trace
    }

    fn log(&self, record: &log::Record) {
        use core::fmt::Write;

        if !self.enabled(record.metadata()) {
            return;
        }

        let mut write = WriteBuffer {
            text: [0; 255],
            len: 0,
        };
        let result = if record.target() == ":raw" || record.target() == ":program" {
            write!(write, "{}", record.args())
        } else {
            writeln!(write, "{}", record.args())
        };

        if let Ok(()) = result {
            LOG_BUFFER
                .write()
                .write(record.level(), &write.text[..write.len]);
            LOG_NOTIFY.wake_all();
        }
    }

    fn flush(&self) {}
}

/// Adds a kernel log file to the parent node
pub fn add_kernel_log_file(parent: &NodeRef) {
    let node = Node::regular(
        LogFile,
        NodeFlags::IN_MEMORY_SIZE | NodeFlags::IN_MEMORY_PROPS,
        Some(Metadata::now_root(FileMode::new(0o400), 0)),
        None,
    );

    let filename = unsafe { Filename::from_str_unchecked("log") };
    parent.add_child(filename, node).ok();
}

async fn read_log_line(seq: usize, buffer: &mut [u8]) -> (usize, usize) {
    poll_fn(|cx| {
        if let Some(out) = try_read_log_line(seq, buffer) {
            LOG_NOTIFY.remove(cx.waker());
            Poll::Ready(out)
        } else {
            LOG_NOTIFY.register(cx.waker());
            Poll::Pending
        }
    })
    .await
}

fn try_read_log_line(seq: usize, buffer: &mut [u8]) -> Option<(usize, usize)> {
    let lock = LOG_BUFFER.read();
    if seq == lock.head {
        return None;
    }
    let available = lock.head.wrapping_sub(seq);
    let ridx = if available < lock.storage.len() {
        seq
    } else {
        lock.tail
    };

    let (meta, data) = lock.storage.record(ridx);
    let new_seq = ridx.wrapping_add(1);
    let len = core::cmp::min(meta.len, buffer.len() - 2);

    buffer[0] = match meta.level {
        log::Level::Trace => b't',
        log::Level::Debug => b'd',
        log::Level::Info => b'i',
        log::Level::Warn => b'w',
        log::Level::Error => b'e',
    };
    buffer[1] = b':';
    buffer[2..len + 2].copy_from_slice(&data[..len]);

    Some((len + 2, new_seq))
}
