use core::{
    future::poll_fn,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll},
};

use alloc::{sync::Arc, vec, vec::Vec};
use libk_util::{
    sync::{IrqSafeSpinlock, IrqSafeSpinlockGuard},
    waker::QueueWaker,
};
use yggdrasil_abi::error::Error;

use crate::vfs::FileReadiness;

struct PipeInner {
    data: Vec<u8>,
    capacity: usize,
    rd: usize,
    wr: usize,
}

pub struct Pipe {
    inner: IrqSafeSpinlock<PipeInner>,
    shutdown: AtomicBool,
    read_notify: QueueWaker,
    write_notify: QueueWaker,
}

pub enum PipeEnd {
    Read(Arc<Pipe>),
    Write(Arc<Pipe>),
}

impl PipeInner {
    pub fn new(capacity: usize) -> Self {
        Self {
            data: vec![0; capacity],
            capacity,
            rd: 0,
            wr: 0,
        }
    }

    pub fn can_write(&self) -> bool {
        (self.wr + 1) % self.capacity != self.rd
    }

    pub fn can_read(&self) -> bool {
        self.rd != self.wr
    }

    pub unsafe fn write(&mut self, val: u8) {
        self.data[self.wr] = val;
        self.wr = (self.wr + 1) % self.capacity;
    }

    pub unsafe fn read(&mut self) -> u8 {
        let val = self.data[self.rd];
        self.rd = (self.rd + 1) % self.capacity;
        val
    }

    fn try_write(&mut self, val: u8) -> bool {
        if self.can_write() {
            unsafe {
                self.write(val);
            }
            true
        } else {
            false
        }
    }

    fn try_read(&mut self) -> Option<u8> {
        if self.can_read() {
            Some(unsafe { self.read() })
        } else {
            None
        }
    }
}

impl Pipe {
    pub fn new(capacity: usize) -> Self {
        Self {
            inner: IrqSafeSpinlock::new(PipeInner::new(capacity)),
            shutdown: AtomicBool::new(false),
            read_notify: QueueWaker::new(),
            write_notify: QueueWaker::new(),
        }
    }

    pub fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Release);
        self.read_notify.wake_all();
        self.write_notify.wake_all();
    }

    pub async fn blocking_write(&self, val: u8) -> Result<(), Error> {
        poll_fn(|cx| {
            let mut lock = self.inner.lock();
            if self.shutdown.load(Ordering::Acquire) {
                self.write_notify.remove(cx.waker());
                // TODO broken pipe
                Poll::Ready(Err(Error::ReadOnly))
            } else if lock.try_write(val) {
                self.write_notify.remove(cx.waker());
                self.read_notify.wake_one();
                Poll::Ready(Ok(()))
            } else {
                self.write_notify.register(cx.waker());
                Poll::Pending
            }
        })
        .await
    }

    fn poll_read_end(
        &self,
        cx: &mut Context<'_>,
    ) -> Poll<Option<IrqSafeSpinlockGuard<'_, PipeInner>>> {
        let lock = self.inner.lock();

        if lock.can_read() {
            self.read_notify.remove(cx.waker());
            Poll::Ready(Some(lock))
        } else if self.shutdown.load(Ordering::Acquire) {
            self.read_notify.remove(cx.waker());
            Poll::Ready(None)
        } else {
            self.read_notify.register(cx.waker());
            Poll::Pending
        }
    }

    async fn read_async(&self, buf: &mut [u8]) -> Result<usize, Error> {
        let mut pos = 0;
        let lock = poll_fn(|cx| self.poll_read_end(cx)).await;
        match lock {
            Some(mut lock) => {
                while pos < buf.len()
                    && let Some(byte) = lock.try_read()
                {
                    buf[pos] = byte;
                    pos += 1;
                }
                if pos != 0 {
                    self.write_notify.wake_all();
                }
                Ok(pos)
            }
            None => Ok(0),
        }
    }
}

impl FileReadiness for Pipe {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        if self.shutdown.load(Ordering::Acquire) || self.inner.lock().can_read() {
            self.read_notify.remove(cx.waker());
            Poll::Ready(Ok(()))
        } else {
            self.read_notify.register(cx.waker());
            Poll::Pending
        }
    }
}

impl PipeEnd {
    pub fn new_pair(capacity: usize) -> (PipeEnd, PipeEnd) {
        let pipe = Arc::new(Pipe::new(capacity));
        let read = PipeEnd::Read(pipe.clone());
        let write = PipeEnd::Write(pipe);

        (read, write)
    }

    pub fn read(&self, buf: &mut [u8], nonblocking: bool) -> Result<usize, Error> {
        let PipeEnd::Read(read) = self else {
            return Err(Error::InvalidOperation);
        };

        if nonblocking {
            let mut lock = read.inner.lock();
            let mut pos = 0;
            while pos < buf.len()
                && let Some(ch) = lock.try_read()
            {
                buf[pos] = ch;
                pos += 1;
            }
            read.write_notify.wake_all();
            if pos == 0 && !read.shutdown.load(Ordering::Acquire) {
                Err(Error::WouldBlock)
            } else {
                Ok(pos)
            }
        } else {
            block!(read.read_async(buf).await)?
        }
    }

    pub fn write(&self, buf: &[u8]) -> Result<usize, Error> {
        let PipeEnd::Write(write) = self else {
            return Err(Error::InvalidOperation);
        };

        block! {
            let mut pos = 0;
            let mut rem = buf.len();

            while rem != 0 {
                write.blocking_write(buf[pos]).await?;
                pos += 1;
                rem -= 1;
            }

            Ok(pos)
        }?
    }
}

impl FileReadiness for PipeEnd {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        match self {
            Self::Read(pipe) => pipe.poll_read(cx),
            Self::Write(_) => Poll::Ready(Err(Error::NotImplemented)),
        }
    }
}

impl Drop for PipeEnd {
    fn drop(&mut self) {
        match self {
            Self::Read(read) => read.shutdown(),
            Self::Write(write) => write.shutdown(),
        }
    }
}
