use alloc::sync::Arc;
use legacy::PciLegacyConfigurationSpace;

use super::{PciAddress, PciBaseAddress, PciCapability, PciCapabilityId, PciEcam};
use crate::{device::PciInterruptPin, PciCommandRegister, PciStatusRegister};

pub(super) mod ecam;
pub(super) mod legacy;

macro_rules! pci_config_field_getter {
    ($self:ident, u32, $offset:expr) => {
        $self.read_u32($offset)
    };

    ($self:ident, u16, $offset:expr) => {
        $self.read_u16($offset)
    };

    ($self:ident, u8, $offset:expr) => {
        $self.read_u8($offset)
    };
}

macro_rules! pci_config_field_setter {
    ($self:ident, u32, $offset:expr, $value:expr) => {
        $self.write_u32($offset, $value)
    };

    ($self:ident, u16, $offset:expr, $value:expr) => {{
        $self.write_u16($offset, $value)
    }};

    ($self:ident, u8, $offset:expr, $value:expr) => {
        $self.write_u8($offset, $value)
    };
}

macro_rules! pci_config_field {
    (
        $offset:expr => $ty:ident,
        $(#[$getter_meta:meta])* $getter:ident
        $(, $(#[$setter_meta:meta])* $setter:ident)?
    ) => {
        $(#[$getter_meta])*
        fn $getter(&self) -> $ty {
            pci_config_field_getter!(self, $ty, $offset)
        }

        $(
            $(#[$setter_meta])*
            fn $setter(&self, value: $ty) {
                pci_config_field_setter!(self, $ty, $offset, value)
            }
        )?
    };
}

/// Describes a configuration space access method for a PCI device
#[derive(Debug, Clone)]
pub enum PciConfigSpace {
    /// Legacy configuration space.
    ///
    /// See [PciLegacyConfigurationSpace].
    Legacy(PciLegacyConfigurationSpace),

    /// Enhanced Configuration Access Mechanism (PCIe).
    ///
    /// See [PciEcam].
    Ecam(PciEcam),
}

pub struct CapabilityIterator<'s, S: PciConfigurationSpace + ?Sized> {
    space: &'s S,
    current: Option<usize>,
}

impl<S: PciConfigurationSpace + ?Sized> Iterator for CapabilityIterator<'_, S> {
    type Item = (Option<PciCapabilityId>, usize, usize);

    fn next(&mut self) -> Option<Self::Item> {
        let offset = self.current? & !0x3;

        let id = PciCapabilityId::try_from(self.space.read_u8(offset)).ok();
        let len = self.space.read_u8(offset + 2);
        let next_pointer = self.space.read_u8(offset + 1);

        self.current = if next_pointer != 0 {
            Some(next_pointer as usize)
        } else {
            None
        };

        Some((id, offset, len as usize))
    }
}

/// Interface for accessing the configuration space of a device
pub trait PciConfigurationSpace {
    /// Reads a 32-bit value from the device configuration space.
    ///
    /// # Note
    ///
    /// The `offset` must be u32-aligned.
    fn read_u32(&self, offset: usize) -> u32;

    /// Writes a 32-bit value to the device configuration space.
    ///
    /// # Note
    ///
    /// The `offset` must be u32-aligned.
    fn write_u32(&self, offset: usize, value: u32);

    /// Reads a 16-bit value from the device configuration space.
    ///
    /// # Note
    ///
    /// The `offset` must be u16-aligned.
    fn read_u16(&self, offset: usize) -> u16 {
        assert_eq!(offset & 1, 0);
        let value = self.read_u32(offset & !3);
        (value >> ((offset & 3) * 8)) as u16
    }

    /// Reads a byte from the device configuration space
    fn read_u8(&self, offset: usize) -> u8 {
        let value = self.read_u32(offset & !3);
        (value >> ((offset & 3) * 8)) as u8
    }

    /// Writes a 16-bit value to the device configuration space.
    ///
    /// # Note
    ///
    /// The `offset` must be u16-aligned.
    fn write_u16(&self, offset: usize, value: u16) {
        let shift = ((offset >> 1) & 1) << 4;
        assert_eq!(offset & 1, 0);
        let mut tmp = self.read_u32(offset & !3);
        tmp &= !(0xFFFF << shift);
        tmp |= (value as u32) << shift;
        self.write_u32(offset & !3, tmp);
    }

    /// Writes a byte to the device configuration space
    fn write_u8(&self, _offset: usize, _value: u16) {
        todo!()
    }

    /// Returns `true` if the device is present on the bus (i.e. configuration space is not filled
    /// with only 1's)
    fn is_valid(&self) -> bool {
        self.vendor_id() != 0xFFFF && self.device_id() != 0xFFFF
    }

    pci_config_field!(
        0x00 => u16,
        #[doc = "Returns the Vendor ID"] vendor_id
    );
    pci_config_field!(0x02 => u16,
        #[doc = "Returns the Device ID"] device_id
    );
    pci_config_field!(
        0x04 => u16,
        #[doc = "Returns the value of the command register"] command,
        #[doc = "Writes to the command word register"] set_command
    );
    pci_config_field!(
        0x06 => u16,
        #[doc = "Returns the value of the status register"] status
    );

    pci_config_field!(
        0x08 => u8,
        #[doc = "Returns the device Revision ID"]
        rev_id
    );
    pci_config_field!(
        0x09 => u8,
        #[doc = "Returns the device Prog IF field"]
        prog_if
    );
    pci_config_field!(
        0x0A => u8,
        #[doc = "Returns the device Subclass field"]
        subclass
    );
    pci_config_field!(
        0x0B => u8,
        #[doc = "Returns the device Class Code field"]
        class_code
    );

    // ...
    pci_config_field!(
        0x0E => u8,
        #[doc = "Returns the header type of the device"]
        header_type
    );
    pci_config_field!(
        0x19 => u8,
        #[doc = r#"
            Returns the secondary bus number associated with this device

            # Note

            The function is only valid for devices with `header_type() == 1`
        "#]
        secondary_bus
    );
    pci_config_field!(
        0x34 => u8,
        #[doc =
            r"Returns the offset within the configuration space where the Capabilities List
            is located. Only valid if the corresponding Status Register bit is set"
        ]
        capability_pointer
    );

    fn interrupt_pin(&self) -> Option<PciInterruptPin> {
        PciInterruptPin::try_from(self.read_u8(0x3D) as u32).ok()
    }

    fn interrupt_line(&self) -> Option<u8> {
        let value = self.read_u8(0x3C);
        if value < 16 {
            Some(value)
        } else {
            None
        }
    }

    /// # Safety
    ///
    /// This function is only meant to be called before the device has seen any use by the OS,
    /// it has not been tested outside of this use case.
    unsafe fn bar_size(&self, index: usize) -> usize {
        let cmd = self.command();

        // Disable I/O and memory
        self.set_command(
            cmd & !(PciCommandRegister::ENABLE_IO | PciCommandRegister::ENABLE_MEMORY).bits(),
        );

        let orig_value = self.bar(index).unwrap();
        // TODO preserve prefetch bit
        let mask_value = match orig_value {
            PciBaseAddress::Io(_) => PciBaseAddress::Io(0xFFFC),
            PciBaseAddress::Memory32(_) => PciBaseAddress::Memory32(0xFFFFFFF0),
            PciBaseAddress::Memory64(_) => PciBaseAddress::Memory64(0xFFFFFFFFFFFFFFF0),
        };
        self.set_bar(index, mask_value);
        let new_value = self.bar(index).unwrap();

        let size = match new_value {
            PciBaseAddress::Io(address) if address != 0 => ((!address) + 1) as usize,
            PciBaseAddress::Memory32(address) if address != 0 => ((!address) + 1) as usize,
            PciBaseAddress::Memory64(address) if address != 0 => ((!address) + 1) as usize,
            _ => 0,
        };

        self.set_bar(index, orig_value);
        self.set_command(cmd);

        size
    }

    /// Updates the value of the Base Address Register with given index.
    ///
    /// # Note
    ///
    /// The function is only valid for devices with `header_type() == 0`
    ///
    /// The `index` corresponds to the actual configuration space BAR index.
    ///
    /// # Safety
    ///
    /// Precondition: the device must have memory access disabled through its command register
    /// prior to setting a BAR.
    unsafe fn set_bar(&self, index: usize, value: PciBaseAddress) {
        assert!(index < 6);

        match value {
            PciBaseAddress::Io(value) => {
                self.write_u32(0x10 + index * 4, ((value as u32) & !0x3) | 1)
            }
            PciBaseAddress::Memory32(address) => self.write_u32(0x10 + index * 4, address & !0xF),
            PciBaseAddress::Memory64(address) => {
                self.write_u32(0x10 + index * 4, ((address as u32) & !0xF) | (2 << 1));
                self.write_u32(0x10 + (index + 1) * 4, (address >> 32) as u32);
            }
        }
    }

    /// Returns the value of the Base Address Register with given index.
    ///
    /// # Note
    ///
    /// The function is only valid for devices with `header_type() == 0`
    ///
    /// The `index` corresponds to the actual configuration space BAR index, i.e. if a 64-bit
    /// address occupies [BAR0, BAR1] and BAR 1 is requested, the function will return [None].
    fn bar(&self, index: usize) -> Option<PciBaseAddress> {
        assert!(index < 6);

        if index % 2 == 0 {
            let w0 = self.read_u32(0x10 + index * 4);

            match w0 & 1 {
                0 => match (w0 >> 1) & 3 {
                    0 => {
                        // 32-bit memory BAR
                        Some(PciBaseAddress::Memory32(w0 & !0xF))
                    }
                    2 => {
                        // 64-bit memory BAR
                        let w1 = self.read_u32(0x10 + (index + 1) * 4);
                        Some(PciBaseAddress::Memory64(
                            ((w1 as u64) << 32) | ((w0 as u64) & !0xF),
                        ))
                    }
                    _ => unimplemented!(),
                },
                1 => Some(PciBaseAddress::Io((w0 as u16) & !0x3)),
                _ => unreachable!(),
            }
        } else {
            let prev_w0 = self.read_u32(0x10 + (index - 1) * 4);
            if prev_w0 & 0x7 == 0x4 {
                // Previous BAR is 64-bit memory and this one is its continuation
                return None;
            }

            let w0 = self.read_u32(0x10 + index * 4);

            match w0 & 1 {
                0 => match (w0 >> 1) & 3 {
                    0 => {
                        // 32-bit memory BAR
                        Some(PciBaseAddress::Memory32(w0 & !0xF))
                    }
                    // TODO can 64-bit BARs not be on a 64-bit boundary?
                    2 => None,
                    _ => unimplemented!(),
                },
                1 => todo!(),
                _ => unreachable!(),
            }
        }
    }

    /// Returns an iterator over the PCI capabilities
    fn capability_iter(&self) -> CapabilityIterator<'_, Self> {
        let status = PciStatusRegister::from_bits_retain(self.status());

        let current = if status.contains(PciStatusRegister::CAPABILITIES_LIST) {
            let ptr = self.capability_pointer() as usize;

            if ptr != 0 {
                Some(self.capability_pointer() as usize)
            } else {
                None
            }
        } else {
            // Return an empty iterator
            None
        };

        CapabilityIterator {
            space: self,
            current,
        }
    }

    /// Locates a capability within this configuration space
    fn capability<C: PciCapability>(&self) -> Option<C::CapabilityData<'_, Self>> {
        self.capability_iter().find_map(|(id, offset, len)| {
            if id == Some(C::ID) && C::check(self, offset, len) {
                Some(C::data(self, offset, len))
            } else {
                None
            }
        })
    }
}

impl<T: PciConfigurationSpace> PciConfigurationSpace for Arc<T> {
    fn read_u32(&self, offset: usize) -> u32 {
        T::read_u32(self.as_ref(), offset)
    }

    fn write_u32(&self, offset: usize, value: u32) {
        T::write_u32(self.as_ref(), offset, value);
    }
}
