use core::ops::{Deref, DerefMut};

use device_api::dma::DmaAllocator;
use libk::dma::{BusAddress, DmaBuffer};
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use xhci_lib::context::{self, DeviceHandler, InputHandler};
use ygg_driver_usb::error::UsbError;

use crate::regs::{ContextSize, PortNumber};

pub enum XhciDeviceContext {
    Context32(IrqSafeRwLock<DmaBuffer<context::Device32Byte>>),
    Context64(IrqSafeRwLock<DmaBuffer<context::Device64Byte>>),
}

pub enum XhciInputContext {
    Context32(DmaBuffer<context::Input32Byte>),
    Context64(DmaBuffer<context::Input64Byte>),
}

impl XhciDeviceContext {
    pub fn new(dma: &dyn DmaAllocator, size: ContextSize) -> Result<Self, UsbError> {
        match size {
            ContextSize::Context32 => DmaBuffer::new(dma, context::Device::new_32byte())
                .map(IrqSafeRwLock::new)
                .map(Self::Context32),
            ContextSize::Context64 => DmaBuffer::new(dma, context::Device::new_64byte())
                .map(IrqSafeRwLock::new)
                .map(Self::Context64),
        }
        .map_err(UsbError::MemoryError)
    }

    pub fn map<T, F: FnOnce(&dyn DeviceHandler) -> T>(&self, mapper: F) -> T {
        match self {
            Self::Context32(cx) => mapper(&**cx.read()),
            Self::Context64(cx) => mapper(&**cx.read()),
        }
    }

    pub fn bus_address(&self) -> BusAddress {
        match self {
            Self::Context32(cx) => cx.read().bus_address(),
            Self::Context64(cx) => cx.read().bus_address(),
        }
    }
    // pub fn physical_address(&self) -> PhysicalAddress {
    // }
}

impl XhciInputContext {
    pub fn new(dma: &dyn DmaAllocator, size: ContextSize) -> Result<Self, UsbError> {
        match size {
            ContextSize::Context32 => {
                DmaBuffer::new(dma, context::Input::new_32byte()).map(Self::Context32)
            }
            ContextSize::Context64 => {
                DmaBuffer::new(dma, context::Input::new_64byte()).map(Self::Context64)
            }
        }
        .map_err(UsbError::MemoryError)
    }

    pub fn new_address_device(
        dma: &dyn DmaAllocator,
        size: ContextSize,
        root_hub_port_number: PortNumber,
        max_packet_size: usize,
        speed: u8,
        dequeue_pointer: BusAddress,
    ) -> Result<Self, UsbError> {
        let mut cx = Self::new(dma, size)?;

        {
            let control = cx.control_mut();

            control.set_add_context_flag(0); // Enable slot context
            control.set_add_context_flag(1); // Enable endpoint 0 context
        }

        {
            let slot = cx.device_mut().slot_mut();

            slot.set_context_entries(1);
            slot.set_interrupter_target(0);
            slot.set_root_hub_port_number(root_hub_port_number.into());
            slot.set_speed(speed);
        }

        {
            let ep0 = cx.device_mut().endpoint_mut(1);

            ep0.set_endpoint_type(context::EndpointType::Control);
            ep0.set_error_count(3);
            // Use the provided max_packet_size, or the default one for the given speed
            ep0.set_max_packet_size(max_packet_size as u16);
            ep0.set_tr_dequeue_pointer(dequeue_pointer.into_u64());
            ep0.set_dequeue_cycle_state();
        }

        Ok(cx)
    }

    // pub fn physical_address(&self) -> PhysicalAddress {
    //     match self {
    //         Self::Context32(cx) => unsafe { cx.as_physical_address() },
    //         Self::Context64(cx) => unsafe { cx.as_physical_address() },
    //     }
    // }
    pub fn bus_address(&self) -> BusAddress {
        match self {
            Self::Context32(cx) => cx.bus_address(),
            Self::Context64(cx) => cx.bus_address(),
        }
    }
}

impl Deref for XhciInputContext {
    type Target = dyn InputHandler;

    fn deref(&self) -> &Self::Target {
        match self {
            Self::Context32(cx) => &**cx,
            Self::Context64(cx) => &**cx,
        }
    }
}

impl DerefMut for XhciInputContext {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            Self::Context32(cx) => &mut **cx,
            Self::Context64(cx) => &mut **cx,
        }
    }
}
