use core::sync::atomic::{AtomicU64, Ordering};

use abi::error::Error;
use acpi::HpetInfo;
use alloc::{collections::btree_map::BTreeMap, format, string::String, sync::Arc};
use device_api::{
    device::Device,
    interrupt::{InterruptHandler, Irq, IrqOptions, IrqTrigger},
};
use libk::{device::external_interrupt_controller, task::runtime, time};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite},
};

register_bitfields! {
    u64,
    GENERAL_CAPABILITIES [
        COUNTER_CLK_PERIOD OFFSET(32) NUMBITS(32) [],
        VENDOR_ID OFFSET(16) NUMBITS(16) [],
        COUNT_SIZE_CAP OFFSET(13) NUMBITS(1) [],
        NUM_TIM_CAP OFFSET(8) NUMBITS(4) [],
        REV_ID OFFSET(0) NUMBITS(8) [],
    ],
    GENERAL_CONFIG [
        LEG_RT_CNF OFFSET(1) NUMBITS(1) [],
        ENABLE_CNF OFFSET(0) NUMBITS(1) [],
    ],
    Tn_CONFIG [
        Tn_INT_ROUTE_CAP OFFSET(32) NUMBITS(32) [],
        Tn_FSB_INT_DEL_CAP OFFSET(15) NUMBITS(1) [],
        Tn_FSB_EN_CNF OFFSET(14) NUMBITS(1) [],
        Tn_INT_ROUTE_CNF OFFSET(9) NUMBITS(5) [],
        Tn_32MODE_CNF OFFSET(8) NUMBITS(1) [],
        Tn_VAL_SET_CNF OFFSET(6) NUMBITS(1) [],
        Tn_SIZE_CAP OFFSET(5) NUMBITS(1) [],
        Tn_PER_INT_CAP OFFSET(4) NUMBITS(1) [],
        Tn_TYPE_CNF OFFSET(3) NUMBITS(1) [
            NonPeriodic = 0,
            Periodic = 1,
        ],
        Tn_INT_ENB_CNF OFFSET(2) NUMBITS(1) [],
        Tn_INT_TYPE_CNF OFFSET(1) NUMBITS(1) [
            Edge = 0,
            Level = 1,
        ],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    GeneralRegs {
        (0x000 => GENERAL_CAPABILITIES: ReadOnly<u64, GENERAL_CAPABILITIES::Register>),
        (0x008 => _0),
        (0x010 => GENERAL_CONFIG: ReadWrite<u64, GENERAL_CONFIG::Register>),
        (0x018 => _1),
        (0x020 => GENERAL_ISR: ReadWrite<u64>),
        (0x028 => _2),
        (0x0F0 => MAIN_COUNTER: ReadWrite<u64>),
        (0x0F8 => _3),
        (0x100 => @END),
    }
}

register_structs! {
    #[allow(non_snake_case)]
    TimerRegs {
        (0x00 => Tn_CONFIG: ReadWrite<u64, Tn_CONFIG::Register>),
        (0x08 => Tn_COMPARATOR: ReadWrite<u64>),
        (0x10 => Tn_FSB_ROUTE: ReadWrite<u64>),
        (0x18 => _0),
        (0x20 => @END),
    }
}

pub struct HpetTimer {
    name: String,
    parent: Arc<Hpet>,
    regs: IrqSafeRwLock<DeviceMemoryIo<'static, TimerRegs>>,
    system_clock_source: bool,
    period: u64,

    last_ticks: AtomicU64,

    periodic: bool,
    bits64: bool,
}

pub struct Hpet {
    base: PhysicalAddress,
    regs: IrqSafeRwLock<DeviceMemoryIo<'static, GeneralRegs>>,
    timers: IrqSafeRwLock<BTreeMap<usize, Arc<HpetTimer>>>,

    base_frequency: u64,
    base_period: u64,
}

const FS_IN_S: u64 = 10u64.pow(15);
const FS_IN_NS: u64 = 10u64.pow(6);

impl InterruptHandler for HpetTimer {
    fn handle_irq(self: Arc<Self>, _vector: Option<usize>) -> bool {
        let now = self.parent.read_counter();
        let last = self.last_ticks.swap(now, Ordering::Relaxed);
        let delta = now.wrapping_sub(last);

        if !self.periodic {
            self.regs.write().rearm(now.wrapping_add(self.period));
        }

        if self.system_clock_source {
            let dt = delta * (self.parent.base_period / FS_IN_NS);
            time::add_nanoseconds(dt);
            runtime::tick();
        }

        true
    }
}

impl Device for HpetTimer {
    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
        if self.period > u32::MAX as u64 && !self.bits64 {
            log::error!(
                "HPET period is >32bit and the HPET itself does not support 64bit counters"
            );
            return Err(Error::InvalidArgument);
        }

        let intc = external_interrupt_controller()?;

        let regs = self.regs.write();

        if regs.is_fsb_capable() {
            log::warn!("TODO: HPET supports FSB interrupt delivery, but the kernel doesn't yet");
        }

        let irq = regs.select_irq().ok_or(Error::InvalidArgument)?;

        log::info!(
            "Set up {:?}, period={}t, periodic={}, irq={:?}",
            self.name,
            self.period,
            self.periodic,
            irq
        );

        intc.register_irq(
            irq,
            IrqOptions {
                trigger: IrqTrigger::Edge,
                ..Default::default()
            },
            self.clone(),
        )?;
        intc.enable_irq(irq)?;
        let irq = IrqDestination::Irq(irq);

        let now = self.parent.read_counter();

        // Pause the timer
        regs.stop();
        if self.periodic {
            regs.configure_periodic(now.wrapping_add(self.period), self.period, irq);
        } else {
            regs.configure_oneshot(now.wrapping_add(self.period), irq);
        }

        Ok(())
    }

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        todo!()
    }

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

impl Device for Hpet {
    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
        let regs = self.regs.write();

        // Reset the device until at least one timer is created
        regs.stop();
        regs.GENERAL_CONFIG
            .modify(GENERAL_CONFIG::LEG_RT_CNF::CLEAR);
        regs.MAIN_COUNTER.set(0);

        Ok(())
    }

    fn display_name(&self) -> &str {
        "HPET"
    }
}

impl HpetTimer {
    fn create(
        hpet: Arc<Hpet>,
        index: usize,
        frequency: u64,
        system_clock_source: bool,
    ) -> Result<Arc<Self>, Error> {
        let base = hpet.base.add(0x100 + 0x20 * index);
        let regs = unsafe { DeviceMemoryIo::<TimerRegs>::map(base, Default::default()) }?;

        let period = hpet.base_frequency / frequency;

        let periodic = regs.is_periodic();
        let bits64 = regs.is_64_bit();

        let timer = Arc::new(HpetTimer {
            name: format!("hpet{index}"),
            regs: IrqSafeRwLock::new(regs),
            parent: hpet,
            last_ticks: AtomicU64::new(0),
            system_clock_source,
            periodic,
            period,
            bits64,
        });

        Ok(timer)
    }
}

impl Hpet {
    fn new(base: PhysicalAddress) -> Result<Self, Error> {
        let regs = unsafe { DeviceMemoryIo::<GeneralRegs>::map(base, Default::default()) }?;
        let base_period = regs.timer_period();
        let base_frequency = FS_IN_S / base_period;
        Ok(Self {
            base,
            base_period,
            base_frequency,
            regs: IrqSafeRwLock::new(regs),
            timers: IrqSafeRwLock::new(BTreeMap::new()),
        })
    }

    pub fn setup_from_acpi(acpi: &HpetInfo) -> Result<Arc<Self>, Error> {
        let base = PhysicalAddress::from_usize(acpi.base_address);
        let hpet = Arc::new(Self::new(base)?);
        unsafe { hpet.clone().init() }?;
        Ok(hpet)
    }

    pub fn setup_timer(
        self: Arc<Self>,
        index: usize,
        frequency: u64,
        system_clock_source: bool,
    ) -> Result<Arc<HpetTimer>, Error> {
        let regs = self.regs.write();
        let mut timers = self.timers.write();
        let timer_count = regs.timer_count();

        if index >= timer_count {
            log::warn!("HPET #{index} does not exist");
            return Err(Error::DoesNotExist);
        }

        if timers.contains_key(&index) {
            log::warn!("HPET #{index} already configured");
            return Err(Error::AlreadyExists);
        }

        // Pause the timer
        regs.stop();

        drop(regs);

        let timer = HpetTimer::create(self.clone(), index, frequency, system_clock_source);
        let timer = if let Ok(timer) = timer {
            timers.insert(index, timer.clone());

            match unsafe { timer.clone().init() } {
                Ok(()) => Ok(timer),
                Err(error) => Err(error),
            }
        } else {
            timer
        };

        let regs = self.regs.write();

        regs.start();

        timer
    }

    fn read_counter(&self) -> u64 {
        self.regs.read().counter()
    }
}

impl TimerRegs {
    fn is_periodic(&self) -> bool {
        self.Tn_CONFIG.matches_all(Tn_CONFIG::Tn_PER_INT_CAP::SET)
    }

    fn is_64_bit(&self) -> bool {
        self.Tn_CONFIG.matches_all(Tn_CONFIG::Tn_SIZE_CAP::SET)
    }

    fn is_fsb_capable(&self) -> bool {
        self.Tn_CONFIG
            .matches_all(Tn_CONFIG::Tn_FSB_INT_DEL_CAP::SET)
    }

    fn select_irq(&self) -> Option<Irq> {
        let mask = self.Tn_CONFIG.read(Tn_CONFIG::Tn_INT_ROUTE_CAP);
        for i in 9..32 {
            if mask & (1 << i) != 0 {
                return Some(Irq::External(i));
            }
        }
        None
    }

    fn stop(&self) {
        self.Tn_CONFIG.modify(Tn_CONFIG::Tn_INT_ENB_CNF::CLEAR);
    }

    fn enable(&self, irq: IrqDestination) {
        let config = match irq {
            IrqDestination::Irq(Irq::External(irq)) => {
                Tn_CONFIG::Tn_INT_ROUTE_CNF.val(irq as u64) + Tn_CONFIG::Tn_INT_TYPE_CNF::Edge
            }
            _ => unreachable!(),
        };

        self.Tn_CONFIG.modify(config);
        self.Tn_CONFIG.modify(Tn_CONFIG::Tn_INT_ENB_CNF::SET);
    }

    fn configure_oneshot(&self, compare_value: u64, irq: IrqDestination) {
        self.Tn_CONFIG
            .modify(Tn_CONFIG::Tn_32MODE_CNF::CLEAR + Tn_CONFIG::Tn_TYPE_CNF::NonPeriodic);

        self.Tn_COMPARATOR.set(compare_value);

        self.enable(irq);
    }

    fn configure_periodic(&self, compare_value: u64, period: u64, irq: IrqDestination) {
        self.Tn_CONFIG.modify(
            Tn_CONFIG::Tn_VAL_SET_CNF::SET
                + Tn_CONFIG::Tn_32MODE_CNF::CLEAR
                + Tn_CONFIG::Tn_TYPE_CNF::Periodic,
        );
        self.Tn_COMPARATOR.set(compare_value);
        self.Tn_COMPARATOR.set(period);

        self.enable(irq);
    }

    fn rearm(&self, compare_value: u64) {
        self.Tn_COMPARATOR.set(compare_value);
    }
}

impl GeneralRegs {
    fn stop(&self) {
        self.GENERAL_CONFIG
            .modify(GENERAL_CONFIG::ENABLE_CNF::CLEAR);
    }
    fn start(&self) {
        self.GENERAL_CONFIG.modify(GENERAL_CONFIG::ENABLE_CNF::SET);
    }

    fn timer_count(&self) -> usize {
        self.GENERAL_CAPABILITIES
            .read(GENERAL_CAPABILITIES::NUM_TIM_CAP) as usize
            + 1
    }

    fn timer_period(&self) -> u64 {
        self.GENERAL_CAPABILITIES
            .read(GENERAL_CAPABILITIES::COUNTER_CLK_PERIOD)
    }

    fn counter(&self) -> u64 {
        self.MAIN_COUNTER.get()
    }
}

enum IrqDestination {
    Irq(Irq),
}