use core::{
    ops::Deref,
    sync::atomic::{AtomicU32, Ordering},
};

use alloc::{collections::BTreeMap, format, sync::Arc, vec::Vec};
use device_api::{
    i2c::{I2CController, I2CDevice},
    pwm::{PwmChannelDevice, PwmController, PwmControllerDevice},
    spi::SpiController,
};
use libk_util::{OneTimeInit, sync::spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::{error::Error, io::FileMode};

use crate::{
    config,
    debug::{self, DebugSink},
    device::char::spi::SpiControllerDevice,
    fs::devfs,
    vfs::NodeRef,
};

use super::{
    block::{
        BlockDevice,
        partition::{self, Partition},
    },
    char::CharDevice,
    display::{DisplayDevice, DisplayWrapper},
};

struct GenericRegistry<T> {
    members: IrqSafeRwLock<BTreeMap<u32, T>>,
    last_index: AtomicU32,
}

pub type DisplayRegisterCallback = fn(&Arc<dyn DisplayDevice>, Option<(&str, NodeRef)>);

// Manages devices: fb<N>
pub struct DisplayDeviceRegistry {
    registry: GenericRegistry<Arc<DisplayWrapper>>,
    callback: OneTimeInit<DisplayRegisterCallback>,
}

// Manages devices: tty<N>
pub struct TerminalRegistry {
    registry: GenericRegistry<Arc<dyn CharDevice>>,
}

// Manages devices: ttyS<N>
pub struct SerialTerminalRegistry {
    registry: GenericRegistry<Arc<dyn CharDevice>>,
}

// Manages devices: i2c<N>, i2c<N>-<M>
pub struct I2CDeviceRegistry {
    registry: GenericRegistry<Arc<I2CDevice>>,
}

// Manages devices: spi<X>.<Y>
pub struct SpiDeviceRegistry {
    registry: GenericRegistry<Arc<SpiControllerDevice>>,
}

// Manages devices: pwm<N>, pwm<N>-<M>
#[allow(unused)]
struct PwmDeviceRegistryEntry {
    controller: Arc<PwmControllerDevice>,
    controller_node: Option<NodeRef>,
    channels: BTreeMap<u32, (Arc<PwmChannelDevice>, Option<NodeRef>)>,
}

pub struct PwmDeviceRegistry {
    registry: GenericRegistry<PwmDeviceRegistryEntry>,
}

pub struct DeviceRegistry {
    pub display: DisplayDeviceRegistry,
    pub terminal: TerminalRegistry,
    pub serial_terminal: SerialTerminalRegistry,
    pub i2c: I2CDeviceRegistry,
    pub spi: SpiDeviceRegistry,
    pub pwm: PwmDeviceRegistry,
}

pub static DEVICE_REGISTRY: DeviceRegistry = DeviceRegistry {
    display: DisplayDeviceRegistry::new(),
    terminal: TerminalRegistry::new(),
    serial_terminal: SerialTerminalRegistry::new(),
    i2c: I2CDeviceRegistry::new(),
    spi: SpiDeviceRegistry::new(),
    pwm: PwmDeviceRegistry::new(),
};

impl TerminalRegistry {
    pub const fn new() -> Self {
        Self {
            registry: GenericRegistry::new(),
        }
    }

    pub fn register(&self, terminal: Arc<dyn CharDevice>) -> Result<u32, Error> {
        let id = self.registry.register(terminal.clone())?;
        devfs::add_named_char_device(terminal, format!("tty{id}"), FileMode::new(0o660)).ok();
        Ok(id)
    }
}

impl SerialTerminalRegistry {
    pub const fn new() -> Self {
        Self {
            registry: GenericRegistry::new(),
        }
    }

    pub fn register(
        &self,
        terminal: Arc<dyn CharDevice>,
        sink: Option<Arc<dyn DebugSink>>,
    ) -> Result<u32, Error> {
        let id = self.registry.register(terminal.clone())?;
        devfs::add_named_char_device(terminal, format!("ttyS{id}"), FileMode::new(0o660)).ok();
        if let Some(sink) = sink {
            debug::add_serial_sink(sink, config::get().debug.serial_level);
        }
        Ok(id)
    }
}

impl DisplayDeviceRegistry {
    pub const fn new() -> Self {
        Self {
            registry: GenericRegistry::new(),
            callback: OneTimeInit::new(),
        }
    }

    pub fn set_callback(&self, callback: DisplayRegisterCallback) {
        self.callback.init(callback);
    }

    pub fn register(&self, display: Arc<dyn DisplayDevice>, boot: bool) -> Result<u32, Error> {
        let wrapper = Arc::new(DisplayWrapper::new(display, boot)?);
        let id = self.registry.register(wrapper.clone())?;
        let name = format!("fb{id}");
        let node = devfs::add_named_block_device(wrapper.clone(), &name, FileMode::new(0o600));
        if let Some(callback) = self.callback.try_get() {
            callback(&wrapper.device, node.ok().map(|node| (name.as_str(), node)));
        }
        Ok(id)
    }
}

impl I2CDeviceRegistry {
    pub const fn new() -> Self {
        Self {
            registry: GenericRegistry::new(),
        }
    }

    pub fn register_bus(&self, device: Arc<dyn I2CController>) -> Result<u32, Error> {
        let wrapper = Arc::new(I2CDevice::from(device));
        let id = self.registry.register(wrapper.clone())?;
        let name = format!("i2c{id}");
        devfs::add_named_char_device(wrapper.clone(), &name, FileMode::new(0o600))?;
        Ok(id)
    }
}

impl SpiDeviceRegistry {
    pub const fn new() -> Self {
        Self {
            registry: GenericRegistry::new(),
        }
    }

    pub fn register_bus(&self, device: Arc<dyn SpiController>) -> Result<u32, Error> {
        let wrapper = Arc::new(SpiControllerDevice::from(device));
        let id = self.registry.register(wrapper.clone())?;
        let name = format!("spi{id}");
        devfs::add_named_char_device(wrapper.clone(), &name, FileMode::new(0o600))?;
        Ok(id)
    }
}

impl PwmDeviceRegistry {
    pub const fn new() -> Self {
        Self {
            registry: GenericRegistry::new(),
        }
    }

    pub fn register_controller(&self, controller: Arc<dyn PwmController>) -> Result<u32, Error> {
        let controller_info = controller.controller_info()?;

        let id = self.registry.register_with(|id| {
            let controller_device = Arc::new(PwmControllerDevice::from(controller.clone()));
            let controller_name = format!("pwm{id}");
            log::info!("Register PWM controller {controller_name:?}");
            let controller_node = devfs::add_named_char_device(
                controller_device.clone(),
                controller_name,
                FileMode::new(0o600),
            )
            .ok();

            let start_channel = if controller_info.per_channel_period {
                0
            } else {
                1
            };
            let mut channels = BTreeMap::new();
            for channel in start_channel..controller_info.channel_count {
                let channel_name = format!("pwm{id}-{channel}");
                log::info!("Register PWM channel {channel_name:?}");
                let channel_device = Arc::new(PwmChannelDevice {
                    device: controller.clone(),
                    channel: channel as u32,
                });
                let channel_node = devfs::add_named_char_device(
                    channel_device.clone(),
                    channel_name,
                    FileMode::new(0o600),
                )
                .ok();
                channels.insert(channel as u32, (channel_device, channel_node));
            }

            let entry = PwmDeviceRegistryEntry {
                controller: controller_device,
                controller_node,
                channels,
            };

            entry
        })?;

        Ok(id)
    }
}

impl<T> GenericRegistry<T> {
    pub const fn new() -> Self {
        Self {
            members: IrqSafeRwLock::new(BTreeMap::new()),
            last_index: AtomicU32::new(0),
        }
    }

    pub fn register(&self, member: T) -> Result<u32, Error> {
        let mut members = self.members.write();
        let id = self.last_index.fetch_add(1, Ordering::Release);
        members.insert(id, member);
        Ok(id)
    }

    pub fn register_with<F: FnOnce(u32) -> T>(&self, mapper: F) -> Result<u32, Error> {
        let mut members = self.members.write();
        let id = self.last_index.fetch_add(1, Ordering::Release);
        members.insert(id, (mapper)(id));
        Ok(id)
    }
}

impl Deref for DisplayWrapper {
    type Target = dyn DisplayDevice;

    fn deref(&self) -> &Self::Target {
        self.device.as_ref()
    }
}

impl DeviceRegistry {
    // pub fn add_pending_initialization(&self, device: Arc<dyn Device>, irq_only: bool) {
    //     self.pending_initialization.write().push(PendingDevice {
    //         device,
    //         irq_only,
    //         failed: false,
    //     });
    // }

    // pub fn run_initialization(&self) {
    //     let mut devices = self.pending_initialization.write();

    //     for pending in devices.iter_mut() {
    //         if pending.irq_only {
    //             continue;
    //         }

    //         log::debug!("Init device: {:?}", pending.device.display_name());
    //         if let Err(error) = unsafe { pending.device.clone().init() } {
    //             log::error!("{:?} init error: {error:?}", pending.device.display_name());
    //             pending.failed = true;
    //             continue;
    //         }
    //     }

    //     for pending in devices.drain(..) {
    //         if pending.failed {
    //             continue;
    //         }

    //         log::debug!("Init IRQ: {:?}", pending.device.display_name());
    //         if let Err(error) = unsafe { pending.device.clone().init_irq() } {
    //             log::error!(
    //                 "{:?} IRQ init error: {error:?}",
    //                 pending.device.display_name()
    //             );
    //         }
    //     }
    // }
}

async fn probe_partition_table(
    dev: &Arc<dyn BlockDevice>,
) -> Result<Option<Vec<Partition>>, Error> {
    if let Some(list) = partition::probe_gpt(dev).await? {
        return Ok(Some(list));
    }
    if let Some(list) = partition::probe_mbr(dev).await? {
        return Ok(Some(list));
    }

    Ok(None)
}

pub async fn probe_partitions<F: Fn(usize, Partition)>(
    dev: Arc<dyn BlockDevice>,
    callback: F,
) -> Result<(), Error> {
    match probe_partition_table(&dev).await {
        Ok(Some(partitions)) => {
            for (i, partition) in partitions.into_iter().enumerate() {
                callback(i, partition);
            }
            Ok(())
        }
        Ok(None) => {
            log::warn!("Unrecognized or missing partition table");
            Err(Error::DoesNotExist)
        }
        Err(error) => {
            log::warn!("Could not probe partition table: {error:?}");
            Err(error)
        }
    }
}
