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

use alloc::{collections::BTreeMap, format, sync::Arc, vec::Vec};
use libk_util::{sync::spin_rwlock::IrqSafeRwLock, OneTimeInit};
use yggdrasil_abi::{error::Error, io::FileMode};

use crate::{
    config,
    debug::{self, DebugSink},
    fs::devfs,
    vfs::NodeRef,
};

use super::{
    block::{
        partition::{self, Partition},
        BlockDevice,
    },
    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>>,
}

pub struct DeviceRegistry {
    pub display: DisplayDeviceRegistry,
    pub terminal: TerminalRegistry,
    pub serial_terminal: SerialTerminalRegistry,
}

pub static DEVICE_REGISTRY: DeviceRegistry = DeviceRegistry {
    display: DisplayDeviceRegistry::new(),
    terminal: TerminalRegistry::new(),
    serial_terminal: SerialTerminalRegistry::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<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)
    }
}

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)
        }
    }
}
