use alloc::{sync::Arc, vec::Vec};
use device_api::{
    clock::{ClockHandle, ResetHandle},
    device::{Device, DeviceInitContext},
    gpio::{
        GpioController, GpioInterruptEvent, GpioPinLevel, OutputPinBias, PinHandle,
        SinglePinDirection,
    },
};
use device_tree::{
    driver::{
        device_tree_driver,
        util::{generic_gpio_config, GenericPinctrlBiasConfig, GenericPinctrlConfig},
        DeviceTreeGpioPins, DeviceTreePinController, Node, ProbeContext,
    },
    DeviceTreePropertyRead, TProp,
};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIoMut};
use libk_util::{bit::BitField, sync::IrqSafeSpinlock, OneTimeInit};
use yggdrasil_abi::error::Error;

#[derive(Debug)]
struct PinMuxConfig {
    pin: u8,
    function: u8,
    doen: u8,
    dout: u8,
    din: u8,
}

struct SysRegs(DeviceMemoryIoMut<'static, [u32]>);
struct AonRegs(DeviceMemoryIoMut<'static, [u32]>);

unsafe impl Send for SysRegs {}
unsafe impl Sync for SysRegs {}
unsafe impl Send for AonRegs {}
unsafe impl Sync for AonRegs {}

struct Gpio<R: GpioRegs> {
    regs: IrqSafeSpinlock<R>,
    init: OneTimeInit<Result<(), Error>>,
    clocks: Vec<ClockHandle>,
    resets: Vec<ResetHandle>,
}

trait GpioRegs: Sized + Send + 'static {
    const NAME: &'static str;

    const DOUT_OFFSET: usize;
    const DOEN_OFFSET: usize;
    const GPI_OFFSET: usize;

    const DOUT_MASK: u32;
    const DOEN_MASK: u32;
    const GPI_MASK: u32;

    unsafe fn map(base: PhysicalAddress) -> Result<Self, Error>;
    fn reg_mut(&mut self, index: usize) -> *mut u32;
    fn pin_functions(&self, pin: usize) -> Option<(usize, usize, u32)>;

    fn write_output_pin(&mut self, pin: usize, value: bool) {
        let reg = pin / 4;
        let shift = (pin % 4) * 8;
        let reg_dout = self.reg_mut(Self::DOUT_OFFSET + reg);
        unsafe {
            let mut val = reg_dout.read_volatile();
            val.set_bit(shift, value);
            reg_dout.write_volatile(val);
        }
    }

    fn configure_pin(&mut self, pin: usize, doen: u32, dout: u32, din: u32) {
        let reg = pin / 4;
        let shift = (pin % 4) * 8;

        let reg_doen = self.reg_mut(Self::DOEN_OFFSET + reg);
        let reg_dout = self.reg_mut(Self::DOUT_OFFSET + reg);

        unsafe {
            let mut val = reg_dout.read_volatile();
            val &= !(Self::DOUT_MASK << shift);
            val |= dout << shift;
            reg_dout.write_volatile(val);

            let mut val = reg_doen.read_volatile();
            val &= !(Self::DOEN_MASK << shift);
            val |= doen << shift;
            reg_doen.write_volatile(val);

            if din != 0xFF {
                let din_reg = din as usize / 4;
                let din_shift = (din % 4) * 8;

                let reg_din = self.reg_mut(Self::GPI_OFFSET + din_reg);
                let mut val = reg_din.read_volatile();
                val &= !(Self::GPI_MASK << din_shift);
                val |= (pin as u32 + 2) << din_shift;
                reg_din.write_volatile(val);
            }
        }
    }

    fn configure_pin_function(&mut self, pin: usize, func: u32) {
        let Some((offset, shift, max)) = self.pin_functions(pin) else {
            return;
        };

        if func > max {
            return;
        }

        log::info!("jh7110-sys-pinctrl: set pad{pin} function={func}");
        let reg = self.reg_mut(offset / 4);
        unsafe {
            let mut val = reg.read_volatile();
            val &= !(0x3 << shift);
            val |= func << shift;
            reg.write_volatile(val);
        }
    }

    fn set_input_enabled(&mut self, _pin: usize, _enabled: bool) {}

    fn set_output_enabled(&mut self, _pin: usize, _enabled: bool) {}

    fn set_bias(&mut self, _pin: usize, _bias: GenericPinctrlBiasConfig) {}
}

// sys-pinctrl pads
const fn pad_gpio(n: usize) -> usize {
    n
}
const PAD_GMAC1_RXC: usize = 82;
const SYS_PIN_COUNT: usize = 95;

const SYS_PIN_FUNCTIONS: [Option<(usize, usize, u32)>; SYS_PIN_COUNT] = const {
    let mut t = [const { None }; SYS_PIN_COUNT];

    t[PAD_GMAC1_RXC] = Some((0x29C, 0, 1));
    t[pad_gpio(10)] = Some((0x29C, 2, 3));
    t[pad_gpio(11)] = Some((0x29C, 5, 3));
    t[pad_gpio(12)] = Some((0x29C, 8, 3));
    t[pad_gpio(13)] = Some((0x29C, 11, 3));
    t[pad_gpio(14)] = Some((0x29C, 14, 3));
    t[pad_gpio(15)] = Some((0x29C, 17, 3));
    t[pad_gpio(16)] = Some((0x29C, 20, 3));
    t[pad_gpio(17)] = Some((0x29C, 23, 3));
    t[pad_gpio(18)] = Some((0x29C, 26, 3));
    t[pad_gpio(19)] = Some((0x29C, 29, 3));

    t[pad_gpio(20)] = Some((0x2A0, 0, 3));
    t[pad_gpio(21)] = Some((0x2A0, 3, 3));
    t[pad_gpio(22)] = Some((0x2A0, 6, 3));
    t[pad_gpio(23)] = Some((0x2A0, 9, 3));
    t[pad_gpio(24)] = Some((0x2A0, 12, 3));
    t[pad_gpio(25)] = Some((0x2A0, 15, 3));
    t[pad_gpio(26)] = Some((0x2A0, 18, 3));
    t[pad_gpio(27)] = Some((0x2A0, 21, 3));
    t[pad_gpio(28)] = Some((0x2A0, 24, 3));
    t[pad_gpio(29)] = Some((0x2A0, 27, 3));

    t[pad_gpio(30)] = Some((0x2A4, 0, 3));
    t[pad_gpio(31)] = Some((0x2A4, 3, 3));
    t[pad_gpio(32)] = Some((0x2A4, 6, 3));
    t[pad_gpio(33)] = Some((0x2A4, 9, 3));
    t[pad_gpio(34)] = Some((0x2A4, 12, 3));
    t[pad_gpio(35)] = Some((0x2A4, 15, 3));
    t[pad_gpio(36)] = Some((0x2A4, 17, 3));
    t[pad_gpio(37)] = Some((0x2A4, 20, 3));
    t[pad_gpio(38)] = Some((0x2A4, 23, 3));
    t[pad_gpio(39)] = Some((0x2A4, 26, 3));
    t[pad_gpio(40)] = Some((0x2A4, 29, 3));

    t[pad_gpio(41)] = Some((0x2A8, 0, 3));
    t[pad_gpio(42)] = Some((0x2A8, 3, 3));
    t[pad_gpio(43)] = Some((0x2A8, 6, 3));
    t[pad_gpio(44)] = Some((0x2A8, 9, 3));
    t[pad_gpio(45)] = Some((0x2A8, 12, 3));
    t[pad_gpio(46)] = Some((0x2A8, 15, 3));
    t[pad_gpio(47)] = Some((0x2A8, 18, 3));
    t[pad_gpio(48)] = Some((0x2A8, 21, 3));
    t[pad_gpio(49)] = Some((0x2A8, 24, 3));
    t[pad_gpio(50)] = Some((0x2A8, 27, 3));
    t[pad_gpio(51)] = Some((0x2A8, 30, 3));

    t[pad_gpio(52)] = Some((0x2AC, 0, 3));
    t[pad_gpio(53)] = Some((0x2AC, 2, 3));
    t[pad_gpio(54)] = Some((0x2AC, 4, 3));
    t[pad_gpio(55)] = Some((0x2AC, 6, 3));
    t[pad_gpio(56)] = Some((0x2AC, 9, 3));
    t[pad_gpio(57)] = Some((0x2AC, 12, 3));
    t[pad_gpio(58)] = Some((0x2AC, 15, 3));
    t[pad_gpio(59)] = Some((0x2AC, 18, 3));
    t[pad_gpio(60)] = Some((0x2AC, 21, 3));
    t[pad_gpio(61)] = Some((0x2AC, 24, 3));
    t[pad_gpio(62)] = Some((0x2AC, 27, 3));
    t[pad_gpio(63)] = Some((0x2AC, 30, 3));

    t[pad_gpio(6)] = Some((0x2B0, 0, 3));
    t[pad_gpio(7)] = Some((0x2B0, 2, 3));
    t[pad_gpio(8)] = Some((0x2B0, 5, 3));
    t[pad_gpio(9)] = Some((0x2B0, 8, 3));

    t
};

impl PinMuxConfig {
    fn from_fdt(fdt: u32) -> Self {
        let pin = fdt as u8;
        let function = ((fdt >> 8) & 0x3) as u8;
        let doen = ((fdt >> 10) & 0x3F) as u8;
        let dout = (fdt >> 16) as u8;
        let din = (fdt >> 24) as u8;

        Self {
            pin,
            function,
            doen,
            dout,
            din,
        }
    }
}

impl<R: GpioRegs> Gpio<R> {
    fn from_device_tree(node: &Arc<Node>, context: &mut ProbeContext) -> Result<Arc<Self>, Error> {
        let base = node.map_base(context, 0).ok_or(Error::DoesNotExist)?;
        let regs = unsafe { R::map(base) }?;
        let clocks = node.clocks().map(|c| c.collect()).unwrap_or_default();
        let resets = node.resets().map(|c| c.collect()).unwrap_or_default();

        let gpio = Arc::new(Self {
            regs: IrqSafeSpinlock::new(regs),
            init: OneTimeInit::new(),
            clocks,
            resets,
        });

        Ok(gpio)
    }

    fn configure_pinmux(&self, mux: PinMuxConfig, cfg: &GenericPinctrlConfig) {
        let mut regs = self.regs.lock();
        regs.configure_pin_function(mux.pin as _, mux.function as _);
        regs.configure_pin(mux.pin as _, mux.doen as _, mux.dout as _, mux.din as _);
        if let Some(enable) = cfg.input_enable {
            regs.set_input_enabled(mux.pin as _, enable);
        }
        if let Some(enable) = cfg.output_enable {
            regs.set_output_enabled(mux.pin as _, enable);
        }
        if let Some(bias) = cfg.bias {
            regs.set_bias(mux.pin as _, bias);
        }
    }

    fn configure_pin_output(
        &self,
        pin: usize,
        _active_level: GpioPinLevel,
        initial_level: GpioPinLevel,
        _output_bias: Option<OutputPinBias>,
    ) -> Result<(), Error> {
        // TODO configure pull up/pull down?
        let mut regs = self.regs.lock();
        let dout_val = initial_level == GpioPinLevel::High;
        regs.configure_pin_function(pin, 0);
        regs.configure_pin(pin, 0, dout_val as u32, 0);
        Ok(())
    }

    fn ensure_init(&self) -> Result<(), Error> {
        let res = self.init.or_init_with(|| {
            for clock in self.clocks.iter() {
                clock.enable()?;
            }
            for reset in self.resets.iter() {
                reset.deassert()?;
            }

            // TODO disable all pin IRQs

            Ok(())
        });

        *res
    }
}

impl<R: GpioRegs> Device for Gpio<R> {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        self.ensure_init()
    }

    fn display_name(&self) -> &str {
        R::NAME
    }
}

impl<R: GpioRegs> GpioController for Gpio<R> {
    fn setup_gpio(&self, pin: &PinHandle, _event: Option<GpioInterruptEvent>) -> Result<(), Error> {
        log::warn!(
            "TODO: setup gpio #{} input={}, output={}",
            pin.index,
            pin.config.input,
            pin.config.output
        );
        // TODO bidi pins?
        match pin.config.force_single_direction() {
            Some(SinglePinDirection::Input) => todo!(),
            Some(SinglePinDirection::Output) => self.configure_pin_output(
                pin.index as usize,
                pin.config.active_level,
                pin.config.initial_level,
                pin.config.output_bias,
            ),
            None => {
                todo!("{}: support for bidi gpio", R::NAME);
            }
        }
    }

    fn read_gpio(&self, _pin: &PinHandle) -> Result<bool, Error> {
        Err(Error::NotImplemented)
    }

    fn write_gpio(&self, pin: &PinHandle, value: bool) -> Result<(), Error> {
        self.regs.lock().write_output_pin(pin.index as _, value);
        Ok(())
    }
}

impl<R: GpioRegs> DeviceTreePinController for Gpio<R> {
    fn configure_pin_group(self: Arc<Self>, pins: &Arc<Node>) -> Result<(), Error> {
        self.ensure_init()?;

        for child in pins.children() {
            let pinmux = child.property("pinmux").ok_or(Error::DoesNotExist)?;
            let cfg = GenericPinctrlConfig::from_node(child);

            for i in 0..pinmux.len() / 4 {
                let mux = pinmux.read_cell(i, 1).unwrap() as u32;
                let mux = PinMuxConfig::from_fdt(mux);
                self.configure_pinmux(mux, &cfg);
            }
        }
        Ok(())
    }

    fn map_gpio(
        self: Arc<Self>,
        property: &TProp,
        offset: usize,
        info: &DeviceTreeGpioPins,
    ) -> Option<(PinHandle, usize)> {
        let (pin, options) = property.read_cells_at(offset, (1, 1))?;
        let config = generic_gpio_config(options as u32, info);
        Some((
            PinHandle {
                index: pin as u32,
                config,
                parent: self,
            },
            2,
        ))
    }
}

impl GpioRegs for SysRegs {
    const NAME: &'static str = "jh7110-sys-pinctrl";

    const DOEN_OFFSET: usize = 0x00;
    const DOUT_OFFSET: usize = 0x40 / 4;
    const GPI_OFFSET: usize = 0x80 / 4;

    const DOEN_MASK: u32 = 0x1F;
    const DOUT_MASK: u32 = 0x3F;
    const GPI_MASK: u32 = 0x3F;

    unsafe fn map(base: PhysicalAddress) -> Result<Self, Error> {
        DeviceMemoryIoMut::map_slice(base, 192, Default::default()).map(Self)
    }

    fn pin_functions(&self, pin: usize) -> Option<(usize, usize, u32)> {
        SYS_PIN_FUNCTIONS.get(pin).copied()?
    }

    fn reg_mut(&mut self, index: usize) -> *mut u32 {
        &raw mut self.0[index]
    }
}

impl GpioRegs for AonRegs {
    const NAME: &'static str = "jh7110-aon-pinctrl";

    const DOEN_OFFSET: usize = 0x00;
    const DOUT_OFFSET: usize = 0x04 / 4;
    const GPI_OFFSET: usize = 0x08 / 4;

    const DOEN_MASK: u32 = 0x7;
    const DOUT_MASK: u32 = 0xF;
    const GPI_MASK: u32 = 0xF;

    unsafe fn map(base: PhysicalAddress) -> Result<Self, Error> {
        DeviceMemoryIoMut::map_slice(base, 64, Default::default()).map(Self)
    }

    fn pin_functions(&self, _pin: usize) -> Option<(usize, usize, u32)> {
        None
    }

    fn reg_mut(&mut self, index: usize) -> *mut u32 {
        &raw mut self.0[index]
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-sys-pinctrl"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let gpio = Gpio::<SysRegs>::from_device_tree(node, context)
                .inspect_err(|e| log::error!("jh7110-sys-pinctrl: probe error: {e:?}"))
                .ok()?;

            node.make_pin_controller(gpio.clone());

            Some(gpio)
        }
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-aon-pinctrl"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let gpio = Gpio::<AonRegs>::from_device_tree(node, context)
                .inspect_err(|e| log::error!("jh7110-aon-pinctrl: probe error: {e:?}"))
                .ok()?;

            node.make_pin_controller(gpio.clone());

            Some(gpio)
        }
    }
}
