use alloc::boxed::Box;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embedded_hal::digital::ErrorType;
use embedded_hal::digital::OutputPin as EhOutputPin;
use embedded_hal_bus::util::AtomicCell;
use esp_hal::gpio::interconnect::{PeripheralInput, PeripheralOutput};
use esp_hal::gpio::{Level, Output, OutputConfig, OutputPin};
use esp_hal::spi::Mode;
use esp_hal::spi::master::{Config as SpiConfig, Spi};
use esp_hal::time::Rate;

use crate::drivers::primary_lcd::PrimaryDisplay;
use crate::drivers::primary_lcd::PrimaryDisplayPinConfiguration;
use crate::drivers::secondary_oled::SecondaryDisplay;
use crate::drivers::secondary_oled::SecondaryDisplayPinConfiguration;

use core::cell::RefCell;
use embassy_sync::blocking_mutex::Mutex as BlockingMutex;

pub struct SharedPin<'a, P> {
    mutex: &'a BlockingMutex<CriticalSectionRawMutex, RefCell<P>>,
}

impl<'a, P> Clone for SharedPin<'a, P> {
    fn clone(&self) -> Self {
        Self { mutex: self.mutex }
    }
}

impl<'a, P> SharedPin<'a, P> {
    pub fn new(pin: P) -> Self {
        let mutexed = BlockingMutex::new(RefCell::new(pin));
        let mutex = Box::leak(Box::new(mutexed));
        Self { mutex }
    }
}

impl<'a, P: EhOutputPin> ErrorType for SharedPin<'a, P> {
    type Error = P::Error;
}

impl<'a, P: EhOutputPin> EhOutputPin for SharedPin<'a, P> {
    fn set_low(&mut self) -> Result<(), Self::Error> {
        self.mutex.lock(|p| p.borrow_mut().set_low())
    }

    fn set_high(&mut self) -> Result<(), Self::Error> {
        self.mutex.lock(|p| p.borrow_mut().set_high())
    }
}

pub struct SharedResetPin<'a, P> {
    mutex: &'a BlockingMutex<CriticalSectionRawMutex, RefCell<(P, bool)>>,
}

impl<'a, P> SharedResetPin<'a, P> {
    pub fn new(pin: P, disabled: bool) -> Self {
        let mutexed = BlockingMutex::new(RefCell::new((pin, disabled)));
        let mutex = Box::leak(Box::new(mutexed));
        Self { mutex }
    }

    pub fn disable(&self, disabled: bool) {
        self.mutex.lock(|p| {
            p.borrow_mut().1 = disabled;
        });
        if disabled {
            log::debug!("Reset pin disabled (masked). No more resets now.");
        } else {
            log::debug!("Reset pin enabled (unmasked). Reset now possible.");
        }
    }
}

impl<'a, P> Clone for SharedResetPin<'a, P> {
    fn clone(&self) -> Self {
        Self { mutex: self.mutex }
    }
}

impl<'a, P: EhOutputPin> ErrorType for SharedResetPin<'a, P> {
    type Error = P::Error;
}

impl<'a, P: EhOutputPin> EhOutputPin for SharedResetPin<'a, P> {
    fn set_low(&mut self) -> Result<(), Self::Error> {
        self.mutex.lock(|p| {
            let mut guard = p.borrow_mut();
            let disabled = guard.1;
            let pin = &mut guard.0;
            if !disabled { pin.set_low() } else { Ok(()) }
        })
    }

    fn set_high(&mut self) -> Result<(), Self::Error> {
        self.mutex.lock(|p| {
            let mut guard = p.borrow_mut();
            let disabled = guard.1;
            let pin = &mut guard.0;
            if !disabled { pin.set_high() } else { Ok(()) }
        })
    }
}

pub type SharedOutput<'a> = SharedPin<'a, Output<'a>>;
pub type SharedReset<'a> = SharedResetPin<'a, Output<'a>>;
pub type DisplaySpiBus = Spi<'static, esp_hal::Blocking>;

pub struct DisplayPinConfiguration<SPI, SCK, MOSI, MISO, CS0, CS1, CS2, CS3, RES1, RES2, DC> {
    pub spi_peripheral: SPI,
    pub sck: SCK,
    pub mosi: MOSI,
    pub miso: MISO,
    pub cs_primary: CS0,
    pub cs_secondary_1: CS1,
    pub cs_secondary_2: CS2,
    pub cs_secondary_3: CS3,
    pub reset_primary: RES1,
    pub reset_secondary: RES2,
    pub dc_pin: DC,
}

impl<SPI, SCK, MOSI, MISO, CS0, CS1, CS2, CS3, RES1, RES2, DC>
    DisplayPinConfiguration<SPI, SCK, MOSI, MISO, CS0, CS1, CS2, CS3, RES1, RES2, DC>
where
    SPI: esp_hal::spi::master::Instance + 'static,
    SCK: PeripheralOutput<'static>,
    MOSI: PeripheralOutput<'static>,
    MISO: PeripheralInput<'static>,
    CS0: OutputPin + 'static,
    CS1: OutputPin + 'static,
    CS2: OutputPin + 'static,
    CS3: OutputPin + 'static,
    RES1: OutputPin + 'static,
    RES2: OutputPin + 'static,
    DC: OutputPin + 'static,
{
    pub fn build(self) -> Result<(PrimaryDisplay<'static>, [SecondaryDisplay<'static>; 3]), Self> {
        let spi = Spi::new(
            self.spi_peripheral,
            SpiConfig::default()
                .with_frequency(Rate::from_khz(20_000))
                .with_mode(Mode::_0),
        )
        .unwrap()
        .with_sck(self.sck)
        .with_mosi(self.mosi)
        .with_miso(self.miso);

        let bus_static: &'static AtomicCell<_> = Box::leak(Box::new(AtomicCell::new(spi)));
        let reset_secondary = SharedResetPin::new(
            Output::new(self.reset_secondary, Level::High, OutputConfig::default()),
            false,
        );
        let dc = SharedPin::new(Output::new(
            self.dc_pin,
            Level::High,
            OutputConfig::default(),
        ));

        // Secondary displays must come first or they don't work in wokwi for some reason.
        let secondaries = SecondaryDisplayPinConfiguration {
            spi_bus: bus_static,
            cs_pins: [
                low_active(self.cs_secondary_1),
                low_active(self.cs_secondary_2),
                low_active(self.cs_secondary_3),
            ],
            res: reset_secondary,
            dc: dc.clone(),
        }
        .build();
        
        let primary = PrimaryDisplayPinConfiguration {
            spi_bus: bus_static,
            cs: low_active(self.cs_primary),
            reset: low_active(self.reset_primary),
            dc,
        }
        .build();

        Ok((primary, secondaries))
    }
}

fn low_active<'a, CS: OutputPin + 'a>(pin: CS) -> Output<'a> {
    Output::new(pin, Level::High, OutputConfig::default())
}