use alloc::boxed::Box;
use alloc::sync::Arc;
use mammoth::cap::Capability;
use mammoth::sync::Mutex;
use mammoth::task::Spawner;
use mammoth::task::Task;

use super::registers::{self};
use crate::xhci::data_structures::CommandCompletionCode;
use crate::xhci::data_structures::TransferRequestBlock;
use crate::xhci::data_structures::TrbAddressDeviceCommand;
use crate::xhci::data_structures::TrbCommandCompletion;
use crate::xhci::data_structures::TrbEnableSlotCommand;
use crate::xhci::data_structures::TrbNoOp;
use crate::xhci::data_structures::TrbPortStatusChangeEvent;
use crate::xhci::data_structures::TrbType;
use crate::xhci::data_structures::TypedTrb;
use crate::xhci::device_context_base_array::DeviceSlotManager;
use crate::xhci::interrupter::Interrupter;
use crate::xhci::registers::DoorbellPointer;
use crate::xhci::registers::HostControllerOperationalWrapper;
use crate::xhci::registers::InterrupterRegisterSet;
use crate::xhci::registers::PortStatusAndControl;
use crate::xhci::trb_ring::CommandRing;

pub struct XHCIDriver {
    #[allow(dead_code)]
    pci_device: pci::PciDevice,
    capabilities: registers::HostControllerCapabilities,
    operational: HostControllerOperationalWrapper,
    command_ring: Mutex<CommandRing>,
    // TODO: Add multiple interrupters.
    interrupter: Mutex<Interrupter>,
    device_slot_manager: Mutex<DeviceSlotManager>,
}

impl XHCIDriver {
    pub fn from_pci_device(mut pci_device: pci::PciDevice) -> Self {
        let address =
            ((pci_device.header().bars[1] as usize) << 32) | (pci_device.header().bars[0] as usize);
        let irq_port_cap = pci_device.register_msi().unwrap();

        let (operational, capabilities) = HostControllerOperationalWrapper::new(address as usize);

        let max_slots = capabilities.params_1.max_device_slots();
        let doorbell_physical = address + capabilities.doorbell_offset as usize;
        let (command_doorbell, slot_doorbells) =
            DoorbellPointer::create_command_and_slots(doorbell_physical, max_slots);

        // Offset to skip the mfindex register.
        let interrupter_registers = mammoth::mem::map_direct_physical_and_leak(
            address + capabilities.runtime_register_space_offset as usize,
            size_of::<InterrupterRegisterSet>() * 2,
        );
        let interrupter_registers = unsafe { interrupter_registers.add(1) };

        let mut driver = Self {
            pci_device,
            capabilities,
            operational,
            command_ring: Mutex::new(CommandRing::new(command_doorbell)),
            interrupter: Mutex::new(Interrupter::new(interrupter_registers, irq_port_cap)),
            device_slot_manager: Mutex::new(DeviceSlotManager::new(max_slots, slot_doorbells)),
        };
        driver.initialize();
        driver
    }

    fn initialize(&mut self) {
        #[cfg(feature = "debug")]
        mammoth::debug!("Stopping XHCI Controller.");

        // Stop the host controller.
        self.operational
            .update_command(|cmd| cmd.with_run_stop(false));

        #[cfg(feature = "debug")]
        mammoth::debug!("Waiting for controller to halt.");

        // Sleep until the controller is halted.
        let mut status = self.operational.read_status();
        while !status.host_controller_halted() {
            // TODO: Sleep for how long?
            mammoth::syscall::thread_sleep(50).unwrap();
            status = self.operational.read_status();
        }

        #[cfg(feature = "debug")]
        mammoth::debug!("Resetting Controller.");

        self.operational
            .update_command(|cmd| cmd.with_host_controller_reset(true));

        let mut command: registers::UsbCommand = self.operational.read_command();
        while command.host_controller_reset() {
            // TODO: Sleep for how long?
            mammoth::syscall::thread_sleep(50).unwrap();
            command = self.operational.read_command();
        }

        #[cfg(feature = "debug")]
        mammoth::debug!("XHCI Controller Reset, waiting ready.");

        let mut status = self.operational.read_status();
        while status.controller_not_ready() {
            // TODO: Sleep for how long?
            mammoth::syscall::thread_sleep(50).unwrap();
            status = self.operational.read_status();
        }

        #[cfg(feature = "debug")]
        mammoth::debug!("XHCI Controller Ready.");

        #[cfg(feature = "debug")]
        mammoth::debug!("Setting Command Ring");

        self.operational.set_command_ring_dequeue_pointer(
            self.command_ring.lock().trb_ring.physical_base_address(),
            true,
        );

        #[cfg(feature = "debug")]
        mammoth::debug!("Setting DCBA.");

        self.operational
            .set_device_context_base_address_array_pointer(
                self.device_slot_manager
                    .lock()
                    .device_context_base_array_physical_address(),
            );
        // We tell the controller that we can support as many slots as it does because
        // we allocate a full 4K page to the DCBA, which is 256 entries and the max
        // slots are 255.
        self.operational.update_configure(|cfg| {
            cfg.with_max_device_slots_enabled(self.capabilities.params_1.max_device_slots())
        });

        assert!(
            self.capabilities.params_2.max_scratchpad_buffers() == 0,
            "Unsupported scratchpad buffers."
        );

        #[cfg(feature = "debug")]
        mammoth::debug!("Resetting event ring.");
        // SAFETY: The HC is stopped.
        unsafe { self.interrupter.lock().reset() };

        self.operational
            .update_command(|cmd| cmd.with_run_stop(true).with_interrupter_enable(true));

        #[cfg(feature = "debug")]
        mammoth::debug!("Enabled interrupts and controller.");
    }

    pub fn interrupt_loop(self: Arc<Self>, spawner: Spawner) {
        let completion_handler = |trb: TransferRequestBlock| {
            self.clone().handle_completion(spawner.clone(), trb);
        };

        self.interrupter.lock().interrupt_loop(completion_handler);
    }

    fn handle_completion(self: Arc<XHCIDriver>, spawner: Spawner, trb: TransferRequestBlock) {
        match trb.trb_type() {
            TrbType::TransferEvent => {
                todo!("Handle Transfer")
            }
            TrbType::CommandCompletionEvent => {
                self.command_ring
                    .lock()
                    .trb_ring
                    .handle_completion(TrbCommandCompletion::from_trb(trb));
            }
            TrbType::PortStatusChangeEvent => {
                let trb = TrbPortStatusChangeEvent::from_trb(trb);
                let self_clone = self.clone();
                spawner.spawn(Task::new(async move {
                    self_clone.port_status_change(trb).await
                }));
            }
            _ => {
                panic!("Unhandled event type: {:?}", trb.trb_type());
            }
        }
    }

    async fn send_command(&self, trb: impl TypedTrb) -> TrbCommandCompletion {
        // Split the future and the await so the lock is dropped before we await.
        let future = { self.command_ring.lock().enqueue_command(trb) };
        future.await
    }

    pub async fn startup(&self) {
        #[cfg(feature = "debug")]
        mammoth::debug!("Sending no op command.");

        let result = self.send_command(TrbNoOp::new()).await;

        assert!(result.completion_code() == CommandCompletionCode::Success.into_bits());

        #[cfg(feature = "debug")]
        mammoth::debug!("Successfully tested no op command.");

        #[cfg(feature = "debug")]
        mammoth::debug!("Resetting all connected ports.");
        for port_index in 0..self.operational.num_ports() {
            self.operational
                .update_port_status(port_index, |p| p.clear_change_bits());
        }

        for port_index in 0..self.operational.num_ports() {
            let status = self.operational.get_port(port_index).status_and_control;
            if status.port_power() && status.current_connect_status() {
                mammoth::debug!("Resetting port {}", port_index);
                self.operational.update_port_status(port_index, |_| {
                    PortStatusAndControl::new()
                        .with_port_reset(true)
                        .with_port_power(true)
                });
            }
        }
    }

    async fn port_status_change(self: Arc<Self>, status_change: TrbPortStatusChangeEvent) {
        // Ports are indexed from 1.
        let port_id = status_change.port_id();
        let port_index = (port_id - 1) as usize;

        let port_status = self
            .operational
            .get_port(port_index as usize)
            .status_and_control;

        #[cfg(feature = "debug")]
        mammoth::debug!("Port status change for port {}", port_id);

        if !port_status.port_reset_change() {
            mammoth::debug!(
                "Unknown port status event, not handling. status= {:?}",
                port_status
            );
            return;
        }

        self.operational
            .update_port_status(port_index, |s| s.clear_change_bits());

        #[cfg(feature = "debug")]
        mammoth::debug!("Enabling slot.");

        let resp = self.send_command(TrbEnableSlotCommand::new()).await;
        assert!(resp.completion_code() == CommandCompletionCode::Success.into_bits());

        let slot = resp.slot_id();

        #[cfg(feature = "debug")]
        mammoth::debug!("Creating slot data structures in slot {}.", slot);

        let input_context = self
            .device_slot_manager
            .lock()
            .prep_slot_for_address_device(slot, port_id);

        #[cfg(feature = "debug")]
        mammoth::debug!("Sending address device.");

        let resp = self
            .send_command(
                TrbAddressDeviceCommand::new()
                    .with_slot_id(slot)
                    .with_input_context_pointer(input_context.physical_address() as u64),
            )
            .await;
        assert!(resp.completion_code() == CommandCompletionCode::Success.into_bits());
    }
}