#![allow(unused)]

use core::{
    ops::Range,
    pin::Pin,
    task::{Context, Poll},
};

use alloc::boxed::Box;
use futures_util::{task::AtomicWaker, Future};
use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageBox, PageProvider};
use libk_util::waker::QueueWaker;
use yggdrasil_abi::{error::Error, io::DeviceRequest};

use crate::{
    request::{IoOperation, IoRequest, IoSubmissionId},
    BlockDevice,
};

pub trait CompletionNotify {
    fn wait_for_completion<'a, D: NgBlockDevice + 'a>(
        &'a self,
        device: &'a D,
        id: IoSubmissionId,
    ) -> impl Future<Output = Result<(), Error>> + Send + '_;
}

pub trait NgBlockDevice: Sync {
    type CompletionNotify: CompletionNotify;

    fn bus_id(&self) -> u32; // HBA, controller ID, etc.
    fn unit_id(&self) -> u32; // Drive, slot, connector ID, etc.

    fn block_size(&self) -> u64;
    fn block_count(&self) -> u64;
    fn max_blocks_per_request(&self) -> u64;

    fn submit_request(
        &self,
        request: IoRequest,
    ) -> impl Future<Output = Result<IoSubmissionId, Error>> + Send;

    fn poll_completion(&self, id: IoSubmissionId) -> Poll<Result<(), Error>>;
    fn completion_notify(&self, id: IoSubmissionId) -> &Self::CompletionNotify;

    fn wait_for_completion(
        &self,
        id: IoSubmissionId,
    ) -> impl Future<Output = Result<(), Error>> + Send + '_
    where
        Self: Sized,
    {
        self.completion_notify(id).wait_for_completion(self, id)
    }
}

pub struct NgBlockDeviceWrapper<'a, D: NgBlockDevice + 'a> {
    device: &'a D,

    pub(crate) block_size: u64,
    pub(crate) block_count: u64,
    #[allow(unused)]
    max_blocks_per_request: u64,
}

#[derive(Debug, PartialEq)]
struct BlockChunk {
    lba_start: u64,
    lba_count: usize,
    buffer_offset: usize,
    lba_offset: usize,
    byte_count: usize,
}

struct BlockChunkIter {
    remaining: usize,
    buffer_offset: usize,
    position: u64,

    block_size: u64,
    max_blocks_per_request: u64,
}

impl CompletionNotify for QueueWaker {
    fn wait_for_completion<'a, D: NgBlockDevice + 'a>(
        &'a self,
        device: &'a D,
        id: IoSubmissionId,
    ) -> impl Future<Output = Result<(), Error>> + Send + '_ {
        struct F<'f, D: NgBlockDevice + 'f> {
            device: &'f D,
            notify: &'f QueueWaker,
            id: IoSubmissionId,
        }

        impl<'f, D: NgBlockDevice + 'f> Future for F<'f, D> {
            type Output = Result<(), Error>;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                self.notify.register(cx.waker());
                match self.device.poll_completion(self.id) {
                    Poll::Ready(result) => {
                        self.notify.remove(cx.waker());
                        Poll::Ready(result)
                    }
                    Poll::Pending => Poll::Pending,
                }
            }
        }

        F {
            notify: self,
            device,
            id,
        }
    }
}

impl CompletionNotify for AtomicWaker {
    fn wait_for_completion<'a, D: NgBlockDevice + 'a>(
        &'a self,
        device: &'a D,
        id: IoSubmissionId,
    ) -> impl Future<Output = Result<(), Error>> + Send + '_ {
        struct F<'f, D: NgBlockDevice + 'f> {
            device: &'f D,
            notify: &'f AtomicWaker,
            id: IoSubmissionId,
        }

        impl<'f, D: NgBlockDevice + 'f> Future for F<'f, D> {
            type Output = Result<(), Error>;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                if let Poll::Ready(result) = self.device.poll_completion(self.id) {
                    return Poll::Ready(result);
                }
                self.notify.register(cx.waker());
                self.device.poll_completion(self.id)
            }
        }

        F {
            notify: self,
            device,
            id,
        }
    }
}

impl BlockChunk {
    pub fn block_range(&self) -> Range<usize> {
        self.lba_offset..self.lba_offset + self.byte_count
    }

    pub fn buffer_range(&self) -> Range<usize> {
        self.buffer_offset..self.buffer_offset + self.byte_count
    }
}

impl BlockChunkIter {
    pub fn new(pos: u64, count: usize, lba_size: u64, max_lba_per_request: u64) -> Self {
        Self {
            remaining: count,
            buffer_offset: 0,
            position: pos,

            block_size: lba_size,
            max_blocks_per_request: max_lba_per_request,
        }
    }
}

impl Iterator for BlockChunkIter {
    type Item = BlockChunk;

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            return None;
        }

        let lba_start = self.position / self.block_size;
        let lba_end =
            (self.position + self.remaining as u64 + self.block_size - 1) / self.block_size;

        let lba_count = core::cmp::min(lba_end - lba_start, self.max_blocks_per_request);

        let lba_offset = (self.position % self.block_size) as usize;
        let byte_count = core::cmp::min(
            (lba_count * self.block_size) as usize - lba_offset,
            self.remaining,
        );

        let buffer_offset = self.buffer_offset;

        self.position += byte_count as u64;
        self.buffer_offset += byte_count;
        self.remaining -= byte_count;

        Some(BlockChunk {
            lba_start,
            lba_count: lba_count as usize,
            buffer_offset,
            lba_offset,
            byte_count,
        })
    }
}

impl<'a, D: NgBlockDevice + 'a> NgBlockDeviceWrapper<'a, D> {
    pub fn new(device: &'a D) -> &'a Self {
        let block_size = device.block_size();
        let block_count = device.block_count();
        let max_blocks_per_request = device.max_blocks_per_request();

        Box::leak(Box::new(Self {
            device,
            block_size,
            block_count,
            max_blocks_per_request,
        }))
    }

    async fn read_range_inner(&self, lba: u64, count: usize) -> Result<PageBox<[u8]>, Error> {
        let mut data = PageBox::new_uninit_slice(self.block_size as usize * count)?;

        let id = self
            .device
            .submit_request(IoRequest {
                operation: IoOperation::Read { lba, count },
                data: &mut data,
            })
            .await?;

        self.device.wait_for_completion(id).await?;

        Ok(unsafe { data.assume_init_slice() })
    }
}

impl<'a, D: NgBlockDevice + 'a> PageProvider for NgBlockDeviceWrapper<'a, D> {
    fn get_page(&self, _offset: u64) -> Result<PhysicalAddress, Error> {
        todo!()
    }

    fn release_page(&self, _offset: u64, _phys: PhysicalAddress) -> Result<(), Error> {
        todo!()
    }

    fn clone_page(
        &self,
        _offset: u64,
        _src_phys: PhysicalAddress,
        _src_attrs: MapAttributes,
    ) -> Result<PhysicalAddress, Error> {
        todo!()
    }
}

impl<'a, D: NgBlockDevice + 'a> BlockDevice for NgBlockDeviceWrapper<'a, D> {
    fn poll_read(
        &self,
        cx: &mut Context<'_>,
        pos: u64,
        buf: &mut [u8],
    ) -> Poll<Result<usize, Error>> {
        todo!()
    }

    fn poll_write(&self, cx: &mut Context<'_>, pos: u64, buf: &[u8]) -> Poll<Result<usize, Error>> {
        todo!()
    }

    // fn read(&'static self, pos: u64, buf: &mut [u8]) -> Result<usize, Error> {
    //     // TODO block cache
    //     block! {
    //         let mut bytes_read = 0;

    //         for chunk in
    //             BlockChunkIter::new(pos, buf.len(), self.block_size, self.max_blocks_per_request)
    //         {
    //             log::debug!(
    //                 "Read chunk: lba_start={}, lba_count={}",
    //                 chunk.lba_start,
    //                 chunk.lba_count
    //             );

    //             let block = self.read_range_inner(chunk.lba_start, chunk.lba_count).await?;

    //             buf[chunk.buffer_range()].copy_from_slice(&block[chunk.block_range()]);

    //             bytes_read += chunk.byte_count;
    //         }

    //         Ok(bytes_read)
    //     }?
    // }

    // fn write(&'static self, _pos: u64, _buf: &[u8]) -> Result<usize, Error> {
    //     todo!()
    // }

    fn size(&self) -> Result<u64, Error> {
        Ok(self.block_size * self.block_count)
    }

    fn device_request(&self, _req: &mut DeviceRequest) -> Result<(), Error> {
        todo!()
    }
}

#[cfg(test)]
mod tests {
    use crate::device::BlockChunk;

    use super::BlockChunkIter;

    #[test]
    fn block_chunk_iter() {
        let mut it = BlockChunkIter {
            remaining: 512 * 9 + 1,
            position: 123,
            block_size: 512,
            buffer_offset: 0,
            max_blocks_per_request: 2,
        };

        assert_eq!(
            it.next().unwrap(),
            BlockChunk {
                lba_start: 0,
                lba_count: 2,
                buffer_offset: 0,
                lba_offset: 123,
                byte_count: 901
            }
        );
        assert_eq!(
            it.next().unwrap(),
            BlockChunk {
                lba_start: 2,
                lba_count: 2,
                buffer_offset: 1024 - 123,
                lba_offset: 0,
                byte_count: 1024
            }
        );
        assert_eq!(
            it.next().unwrap(),
            BlockChunk {
                lba_start: 4,
                lba_count: 2,
                buffer_offset: 2 * 1024 - 123,
                lba_offset: 0,
                byte_count: 1024
            }
        );
        assert_eq!(
            it.next().unwrap(),
            BlockChunk {
                lba_start: 6,
                lba_count: 2,
                buffer_offset: 3 * 1024 - 123,
                lba_offset: 0,
                byte_count: 1024
            }
        );
        assert_eq!(
            it.next().unwrap(),
            BlockChunk {
                lba_start: 8,
                lba_count: 2,
                buffer_offset: 4 * 1024 - 123,
                lba_offset: 0,
                byte_count: 512 + 123 + 1
            }
        );
    }
}