use core::{
    future::Future,
    ops::Deref,
    ops::DerefMut,
    pin::Pin,
    task::{Context, Poll, Waker},
};

use alloc::boxed::Box;
use alloc::sync::Arc;
use mammoth::{cap::Capability, sync::Mutex, syscall, zion::ZError};

use super::ahci_command::{HostToDeviceRegisterFis, SataCommand};

pub enum CommandStatus {
    Empty,
    NotSent,
    Pending(Waker),
    Complete,
}

pub struct CommandFuture {
    status: Arc<Mutex<CommandStatus>>,
    trigger: Box<dyn Fn() + Sync + Send>,
}

impl CommandFuture {
    pub fn new(status: Arc<Mutex<CommandStatus>>, trigger: Box<dyn Fn() + Sync + Send>) -> Self {
        Self { status, trigger }
    }
}

impl Drop for CommandFuture {
    fn drop(&mut self) {
        *self.status.lock() = CommandStatus::Empty;
    }
}

impl Future for CommandFuture {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let s = self.deref_mut();
        let mut status = s.status.lock();
        match status.deref() {
            CommandStatus::NotSent => {
                *status = CommandStatus::Pending(cx.waker().clone());
                (s.trigger)();
                Poll::Pending
            }
            CommandStatus::Pending(_) => Poll::Pending,
            CommandStatus::Complete => Poll::Ready(()),
            CommandStatus::Empty => panic!("Polling empty command slot"),
        }
    }
}

pub struct Command {
    pub command: SataCommand,
    pub lba: u64,
    pub sector_cnt: u16,
    pub paddr: u64,

    memory_region: Option<Capability>,
}

impl Command {
    pub fn identify() -> Result<Self, ZError> {
        let (memory_region, paddr) = syscall::memory_object_contiguous_physical(512)?;

        Ok(Self {
            command: SataCommand::IdentifyDevice,
            lba: 0,
            sector_cnt: 1,
            paddr,
            memory_region: Some(memory_region),
        })
    }

    pub fn read(lba: u64, lba_count: u16) -> Result<Self, ZError> {
        let (memory_region, paddr) =
            syscall::memory_object_contiguous_physical(512 * (lba_count as u64))?;

        Ok(Self {
            command: SataCommand::DmaReadExt,
            lba,
            sector_cnt: lba_count,
            paddr,
            memory_region: Some(memory_region),
        })
    }

    pub fn read_manual(lba: u64, lba_count: u16, paddr: u64) -> Self {
        Self {
            command: SataCommand::DmaReadExt,
            lba,
            sector_cnt: lba_count,
            paddr: paddr,
            memory_region: None,
        }
    }

    pub fn release_mem_cap(&mut self) -> u64 {
        self.memory_region.take().unwrap().release()
    }
}

impl From<&Command> for HostToDeviceRegisterFis {
    fn from(val: &Command) -> Self {
        HostToDeviceRegisterFis::new_command(val.command, val.lba, val.sector_cnt)
    }
}