#![allow(missing_docs)]

use abi::{
    error::Error,
    io::{FileMode, device::i2c::I2CMasterInfo},
};
use alloc::{boxed::Box, string::String, sync::Arc};
use async_trait::async_trait;
use device_api::{
    clock::Hertz,
    device::{Device, DeviceInitContext},
    i2c::{I2CAddress, I2CController, I2CDevice, I2CTransfer},
};
use device_tree::driver::{Node, ProbeContext, device_tree_driver, lookup_phandle};
use libk::fs::devfs;

pub struct I2CMuxChild {
    name: String,
    channel: u32,
    parent: Arc<dyn I2CController>,
}

pub struct I2CMux {
    name: &'static str,
}

impl I2CMux {
    pub fn from_fdt(node: &Arc<Node>, _context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
        let i2c_parent = node.prop_usize("i2c-parent")? as u32;
        let i2c_parent = lookup_phandle(i2c_parent, true)?;
        let i2c_parent = i2c_parent.as_i2c_controller()?;
        let name = node.name().unwrap_or("i2c-mux-pinctrl");

        let bus_index = i2c_parent.bus_number();

        for (index, child) in node.children().enumerate() {
            node.setup_pins(index).unwrap();
            let channel_index = child.prop_usize("reg")? as u32;
            let name = alloc::format!("{name}@{channel_index}");
            let channel: Arc<dyn I2CController> = Arc::new(I2CMuxChild {
                name,
                channel: channel_index,
                parent: i2c_parent.clone(),
            });

            let device = Arc::new(I2CDevice::from(channel));
            let fs_name = alloc::format!("i2c{bus_index}-{channel_index}");
            devfs::add_named_char_device(device, fs_name, FileMode::new(0o600)).ok();
        }

        Some(Arc::new(I2CMux { name }))
    }
}

impl Device for I2CMux {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        Ok(())
    }

    fn display_name(&self) -> &str {
        self.name
    }
}

impl Device for I2CMuxChild {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        unreachable!()
    }

    fn display_name(&self) -> &str {
        &self.name
    }
}

#[async_trait]
impl I2CController for I2CMuxChild {
    fn bus_number(&self) -> u32 {
        self.parent.bus_number()
    }

    fn child_number(&self) -> Option<u32> {
        Some(self.channel)
    }

    fn set_speed(&self, speed: Hertz) -> Result<Hertz, Error> {
        self.parent.set_speed(speed)
    }

    fn capabilities(&self) -> I2CMasterInfo {
        self.parent.capabilities()
    }

    async fn i2c_transfer(
        &self,
        address: I2CAddress,
        transfer: &mut I2CTransfer,
    ) -> Result<usize, Error> {
        self.parent.i2c_transfer(address, transfer).await
    }
}

device_tree_driver! {
    compatible: ["i2c-mux-pinctrl"],
    config: {
        auto_pinctrl: false,
    },
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            I2CMux::from_fdt(node, context)
        }
    }
}
