use alloc::sync::Arc;
use mammoth::{
    cap::Capability,
    mem::{self, MemoryRegion},
    sync::Mutex,
    thread,
    zion::ZError,
};

use crate::ahci::{
    port::{AhciDeviceDetection, AhciInterfacePowerManagement},
    Command,
};

use super::{hba::AhciHba, port::AhciPortHba, port_controller::PortController};

#[derive(Debug)]
#[repr(C, packed)]
pub struct PciDeviceHeader {
    pub vendor_id: u16,
    pub device_id: u16,
    pub command_reg: u16,
    pub status_reg: u16,
    pub revision: u8,
    pub prog_interface: u8,
    pub subclass: u8,
    pub class_code: u8,
    pub cache_line_size: u8,
    pub latency_timer: u8,
    pub header_type: u8,
    pub bist: u8,
    pub bars: [u32; 5],
    pub abar: u32,
    __: u32,
    pub subsystem_id: u32,
    pub expansion_rom: u32,
    pub cap_ptr: u8,
    ___: [u8; 7],
    pub interrupt_line: u8,
    pub interrupt_pin: u8,
    pub min_grant: u8,
    pub max_latency: u8,
}

pub struct AhciController {
    pci_header: &'static mut PciDeviceHeader,
    hba: Mutex<&'static mut AhciHba>,
    ports: [Option<PortController>; 32],
    hba_vaddr: u64,
}

impl AhciController {
    pub fn new(pci_memory: Capability) -> Self {
        let pci_vaddr = mem::map_cap_and_leak(pci_memory);
        let pci_header = unsafe { (pci_vaddr as *mut PciDeviceHeader).as_mut().unwrap() };

        let hba_vaddr = mem::map_direct_physical_and_leak(pci_header.abar as u64, 0x1100);
        let hba = unsafe { (hba_vaddr as *mut AhciHba).as_mut().unwrap() };
        let mut controller = Self {
            pci_header,
            hba: Mutex::new(hba),
            ports: [const { None }; 32],
            hba_vaddr,
        };
        controller.init();
        controller
    }

    fn init(&mut self) {
        self.hba.lock().init();
        let hba = self.hba.lock();
        for i in 0..(hba.capabilities.read().num_ports() as usize) {
            let port_index = 1 << i;
            if (hba.port_implemented.read() & port_index) != port_index {
                mammoth::debug!("Skipping port {}, not implemented", i);
                continue;
            }

            let port_offset: usize = 0x100 + (0x80 * i);
            let port = unsafe {
                ((self.hba_vaddr as usize + port_offset) as *mut AhciPortHba)
                    .as_mut()
                    .unwrap()
            };

            let sata_status = port.sata_status.read();
            if (sata_status.device_detection() != AhciDeviceDetection::CommunicationEstablished)
                || (sata_status.interface_power_management()
                    != AhciInterfacePowerManagement::Active)
            {
                mammoth::debug!(
                    "Skipping port {}, no communcation. Status: {:?}",
                    i,
                    sata_status
                );
                continue;
            }

            self.ports[i] = Some(PortController::new(port));
        }
    }

    fn irq_num(&self) -> u64 {
        match self.pci_header.interrupt_pin {
            1 => mammoth::zion::kZIrqPci1,
            2 => mammoth::zion::kZIrqPci2,
            3 => mammoth::zion::kZIrqPci3,
            4 => mammoth::zion::kZIrqPci4,
            _ => panic!(
                "Unrecognized pci interrupt pin {}",
                self.pci_header.interrupt_pin
            ),
        }
    }

    fn handle_irq(&self) {
        let mut hba = self.hba.lock();
        for i in 0..hba.capabilities.read().num_ports() {
            let int_offset = 1 << i;
            if (hba.interrupt_status.read() & int_offset) == int_offset {
                if let Some(port) = &self.ports[i as usize] {
                    port.handle_interrupt();
                    hba.interrupt_status.update(|is| {
                        *is &= !int_offset;
                    });
                }
            }
        }
    }

    pub async fn identify_ports(&self) -> Result<(), ZError> {
        for port in self.ports.iter().flatten() {
            let sig = port.get_signature();
            if sig == 0x101 {
                let mut command = Command::identify()?;
                mammoth::debug!("IDENT!");
                port.issue_command(&command)?.await;
                let memory_region =
                    MemoryRegion::from_cap(Capability::take(command.release_mem_cap()))?;
                let ident = memory_region.slice::<u16>();
                let new_sector_size = if ident[106] & (1 << 12) != 0 {
                    ident[117] as u32 | ((ident[118] as u32) << 16)
                } else {
                    512
                };

                let lba_count = if ident[83] & (1 << 10) != 0 {
                    ident[100] as u64
                        | (ident[101] as u64) << 16
                        | (ident[102] as u64) << 32
                        | (ident[103] as u64) << 48
                } else {
                    ident[60] as u64 | (ident[61] as u64) << 16
                };
                mammoth::debug!("Sector size: {:#0x}", new_sector_size);
                mammoth::debug!("LBA Count: {:#0x}", lba_count);

                // We hardcode assumptions about sector size in a few places, better to panic if we
                // are wrong.
                assert_eq!(
                    new_sector_size, 512,
                    "Sector size {} differs from 512.",
                    new_sector_size
                );
            } else {
                mammoth::debug!("Skipping non-sata sig: {:#0x}", sig);
            }
        }
        Ok(())
    }

    pub async fn issue_command(&self, port_num: usize, command: &Command) -> Result<(), ZError> {
        assert!(port_num < 32);

        self.ports[port_num]
            .as_ref()
            .ok_or(ZError::INVALID_ARGUMENT)?
            .issue_command(command)?
            .await;

        Ok(())
    }
}

pub fn spawn_irq_thread(controller: Arc<AhciController>) -> thread::JoinHandle {
    let irq_thread = move || {
        let irq_num = controller.irq_num();
        let irq_port =
            mammoth::port::PortServer::from_cap(mammoth::syscall::register_irq(irq_num).unwrap());
        controller.hba.lock().global_host_control.update(|ghc| {
            ghc.set_interrupt_enable(true);
        });
        loop {
            irq_port.recv_null().unwrap();
            controller.handle_irq();
        }
    };
    thread::spawn(irq_thread)
}

pub async fn identify_ports(controller: Arc<AhciController>) {
    controller.identify_ports().await.unwrap();
}