use core::alloc::Layout;
use core::fmt::Debug;
use core::slice;

use alloc::alloc::Allocator;
use alloc::boxed::Box;
use alloc::collections::btree_map::{BTreeMap, Entry};
use alloc::string::String;
use alloc::sync::Arc;
use alloc::vec::Vec;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::channel::Channel;
use embassy_time::Instant;
use esp_alloc::MemoryCapability;
use log::{debug, error, info, warn};
use rmk::descriptor::KeyboardReport;
use rmk::hid::Report;
use rmk::{heapless, join_all};
use slint::platform::Key;
use xkbcommon::xkb::{self, FeedResult, KeyDirection, Keysym, ModMask, Status};

use crate::ram::PsramAllocator;
use crate::util::{DurationExt, get_file_name};
use crate::{KEYBOARD_REPORT_PROXY, PSRAM_ALLOCATOR};

pub static KEY_MESSAGE_CHANNEL: Channel<CriticalSectionRawMutex, KeyMessage, 16> = Channel::new();
pub static OUTPUT_STRING_CHANNEL: Channel<CriticalSectionRawMutex, String, 16> = Channel::new();

pub struct KeyMessage {
    pub keysym: Keysym,
    pub string: Option<String>,
    pub direction: KeyDirection,
}

impl Debug for KeyMessage {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("KeyMessage")
            .field("keysym", &self.keysym)
            .field("string", &self.string)
            .field_with("direction", |f| match self.direction {
                KeyDirection::Down => f.write_str("Down"),
                KeyDirection::Up => f.write_str("Up"),
            })
            .finish()
    }
}

pub fn create_hid_report_interceptor() -> impl Future<Output = ()> {
    // TODO: Use `include_dir` or similar to support multiple keymaps.
    const KEYMAP_NAME: &str = get_file_name(env!("ACID_KEYMAP_PATH"));
    const KEYMAP_STRING: &str = include_str!(env!("ACID_KEYMAP_PATH"));
    const COMPOSE_MAP_NAME: &str = get_file_name(env!("ACID_COMPOSE_PATH"));
    const COMPOSE_MAP_STRING: &str = include_str!(env!("ACID_COMPOSE_PATH"));
    // E.g. `cs_CZ.UTF-8`
    const COMPOSE_MAP_LOCALE: &str = env!("ACID_COMPOSE_LOCALE");

    info!(
        "Keymap: {KEYMAP_NAME:?}, compose: {COMPOSE_MAP_NAME:?}, compose locale: {COMPOSE_MAP_LOCALE:?}"
    );

    let keymap_string_buffer = unsafe {
        let allocation = PSRAM_ALLOCATOR.alloc_caps(
            MemoryCapability::External.into(),
            Layout::from_size_align(KEYMAP_STRING.len(), 32).unwrap(),
        );
        let slice = str::from_utf8_unchecked_mut(slice::from_raw_parts_mut(
            allocation,
            KEYMAP_STRING.len(),
        ));

        slice
            .as_bytes_mut()
            .copy_from_slice(KEYMAP_STRING.as_bytes());

        Box::from_raw_in(slice as *mut str, &PSRAM_ALLOCATOR)
    };
    let context = xkb::Context::new(xkb::CONTEXT_NO_FLAGS);
    info!("Loading XKB keymap...");
    let instant_start = Instant::now();
    let keymap = Arc::new_in(
        xkb::Keymap::new_from_string(
            &context,
            keymap_string_buffer,
            xkb::KEYMAP_FORMAT_TEXT_V1,
            xkb::KEYMAP_COMPILE_NO_FLAGS,
        )
        .unwrap(),
        &PSRAM_ALLOCATOR,
    );
    let duration = Instant::now().duration_since(instant_start);
    info!(
        "XKB keymap loaded successfully! Took {} seconds.",
        duration.display_as_secs(),
    );
    info!("Loading XKB compose map...");
    let instant_start = Instant::now();
    let compose_table = Arc::new_in(
        xkb::compose::Table::new_from_buffer(
            &context,
            COMPOSE_MAP_STRING,
            COMPOSE_MAP_LOCALE,
            xkb::compose::FORMAT_TEXT_V1,
            xkb::compose::COMPILE_NO_FLAGS,
        )
        .unwrap(),
        &PSRAM_ALLOCATOR,
    );
    let duration = Instant::now().duration_since(instant_start);
    info!(
        "XKB compose map loaded successfully! Took {} seconds.",
        duration.display_as_secs()
    );

    async move {
        join_all![
            send_keycodes_to_slint(keymap.clone(), compose_table.clone()),
            send_strings_to_host(keymap, compose_table)
        ]
        .await;
    }
}

const KEYCODES_LEN_MODIFIERS: usize = 8;
const USB_HID_LEFT_CTRL: u8 = 0xE0;

#[derive(Debug)]
struct ModifierDescriptor {
    name: &'static str,
    hid_flag: Option<u8>,
    keycode: Option<u8>,
}

impl ModifierDescriptor {
    const fn new(name: &'static str, hid_index: Option<u8>) -> Self {
        Self {
            name,
            hid_flag: if let Some(hid_index) = hid_index {
                Some(1_u8 << hid_index)
            } else {
                None
            },
            keycode: if let Some(hid_index) = hid_index {
                Some(USB_HID_LEFT_CTRL + hid_index)
            } else {
                None
            },
        }
    }
}

async fn send_strings_to_host<A: Allocator>(
    keymap: Arc<xkb::Keymap, A>,
    compose_table: Arc<xkb::compose::Table, A>,
) {
    /// Names of modifiers in the conventional order.
    const XKB_REAL_MODIFIER_DESCRIPTORS: [ModifierDescriptor; KEYCODES_LEN_MODIFIERS] = [
        ModifierDescriptor::new(xkb::MOD_NAME_SHIFT, Some(1)),
        ModifierDescriptor::new(xkb::MOD_NAME_CAPS, None),
        ModifierDescriptor::new(xkb::MOD_NAME_CTRL, Some(0)),
        ModifierDescriptor::new(xkb::MOD_NAME_ALT, Some(2)),
        ModifierDescriptor::new(xkb::MOD_NAME_NUM, None),
        ModifierDescriptor::new(xkb::MOD_NAME_MOD3, None),
        ModifierDescriptor::new(xkb::MOD_NAME_LOGO, Some(3)),
        ModifierDescriptor::new(xkb::MOD_NAME_ISO_LEVEL3_SHIFT, Some(6)),
    ];

    let modifier_xkb_index_to_descriptor = keymap
        .mods()
        .enumerate()
        .map(|(index, modifier)| {
            let descriptor = XKB_REAL_MODIFIER_DESCRIPTORS
                .iter()
                .find(|descriptor| descriptor.name == modifier);
            debug!("Modifier #{index} {descriptor:02x?}");
            descriptor
        })
        .collect::<Vec<_>>();

    info!("modifier_xkb_index_to_hid_flag: {modifier_xkb_index_to_descriptor:02x?}");

    loop {
        let string = OUTPUT_STRING_CHANNEL.receive().await;
        let string_keycodes = string_to_hid_keycodes(&keymap, &compose_table, &string);

        warn!("keycodes for {string:?}: {string_keycodes:02x?}");

        if let Ok(string_keycodes) = string_keycodes {
            for keycode in string_keycodes {
                let mut keycodes_vec = heapless::Vec::<u8, 6>::new();
                let modifier = {
                    let mut modifier = 0;
                    let mut remaining_mask = keycode.mod_mask;
                    let mut index = 0;

                    while remaining_mask != 0 {
                        if remaining_mask & 1 != 0
                            && let Some(mod_descriptor) = modifier_xkb_index_to_descriptor[index]
                        {
                            if let Some(hid_flag) = mod_descriptor.hid_flag {
                                modifier |= hid_flag;
                            }

                            if let Some(keycode) = mod_descriptor.keycode {
                                let _ = keycodes_vec.push(keycode);
                            }
                        }

                        remaining_mask >>= 1;
                        index += 1;
                    }

                    modifier
                };

                if keycodes_vec.len() + 1 > keycodes_vec.capacity() {
                    warn!("Failed to send key because the keycode buffer is full.");
                    continue;
                }

                // TODO: Check whether this is necessary, otherwise remove.
                rmk::channel::KEYBOARD_REPORT_RECEIVER
                    .send(Report::KeyboardReport(KeyboardReport {
                        modifier,
                        reserved: Default::default(),
                        leds: Default::default(),
                        keycodes: {
                            let mut keycodes = [0; _];
                            keycodes[..keycodes_vec.len()].copy_from_slice(&keycodes_vec);
                            keycodes
                        },
                    }))
                    .await;

                keycodes_vec
                    .push(keycode.keycode)
                    .expect("no space for the main keycode");
                warn!(
                    "Sending HID keycode 0x{:02x} with modifier 0x{:02x} (xkb mask 0x{:02x}) as keycode sequence {keycodes_vec:02x?}",
                    keycode.keycode, modifier, keycode.mod_mask
                );

                rmk::channel::KEYBOARD_REPORT_RECEIVER
                    .send(Report::KeyboardReport(KeyboardReport {
                        modifier,
                        reserved: Default::default(),
                        leds: Default::default(),
                        keycodes: {
                            let mut keycodes = [0; _];
                            keycodes[..keycodes_vec.len()].copy_from_slice(&keycodes_vec);
                            keycodes
                        },
                    }))
                    .await;
                rmk::channel::KEYBOARD_REPORT_RECEIVER
                    .send(Report::KeyboardReport(KeyboardReport::default()))
                    .await;
            }
        }
    }
}

async fn send_keycodes_to_slint<A: Allocator>(
    keymap: Arc<xkb::Keymap, A>,
    compose_table: Arc<xkb::compose::Table, A>,
) {
    let mut state = xkb::State::new(&keymap);
    let mut previous_state = KeyboardReport::default();
    let mut compose_state = xkb::compose::State::new(&compose_table, xkb::compose::STATE_NO_FLAGS);
    // TODO: Use a stack-allocated map instead
    // This is a map from the basic keysyms (not a composed ones) to the string that should be produced.
    let mut pressed_keys_to_strings = BTreeMap::<Keysym, Option<String>>::new();

    loop {
        let report = KEYBOARD_REPORT_PROXY.receive().await;

        info!("Intercepted HID keyboard report: {report:#02x?}");

        if let Report::KeyboardReport(report) = &report {
            const KEYCODES_LEN_REGULAR: usize = 6;
            const KEYCODES_LEN: usize = KEYCODES_LEN_MODIFIERS + KEYCODES_LEN_REGULAR;

            let mut pressed_keys = rmk::heapless::Vec::<u8, KEYCODES_LEN>::new();
            let mut released_keys = rmk::heapless::Vec::<u8, KEYCODES_LEN>::new();

            let pressed_mods_bits = !previous_state.modifier & report.modifier;
            let released_mods_bits = previous_state.modifier & !report.modifier;

            for index in 0..KEYCODES_LEN_MODIFIERS {
                let mod_bit = 1_u8 << index;
                let mod_keycode = USB_HID_LEFT_CTRL + index as u8;

                if pressed_mods_bits & mod_bit != 0 {
                    pressed_keys.push(mod_keycode).unwrap();
                }

                if released_mods_bits & mod_bit != 0 {
                    released_keys.push(mod_keycode).unwrap();
                }
            }

            // TODO: This currently depends on pressed keys not changing position in the array.
            // Should be made independent of that.
            for (&keycode_old, &keycode_new) in
                core::iter::zip(&previous_state.keycodes, &report.keycodes)
            {
                if keycode_old == keycode_new {
                    continue;
                }

                if keycode_old != 0 {
                    released_keys.push(keycode_old).unwrap();
                }

                if keycode_new != 0 {
                    pressed_keys.push(keycode_new).unwrap();
                }
            }

            previous_state = *report;

            for keycode in released_keys {
                debug!("Release: 0x{:02x} ({})", keycode, keycode);
                let keycode_xkb = hid_to_xkb_keycode(keycode);
                state.update_key(keycode_xkb, KeyDirection::Up);
                let keysym = state.key_get_one_sym(keycode_xkb);
                let string = pressed_keys_to_strings.remove(&keysym).unwrap_or_else(|| {
                            warn!("Could not determine the string of a released key: keysym={keysym:?} pressed_keys_to_strings={pressed_keys_to_strings:?}");
                            None
                        });
                KEY_MESSAGE_CHANNEL
                    .send(KeyMessage {
                        keysym,
                        string,
                        direction: KeyDirection::Up,
                    })
                    .await;
            }

            for keycode in pressed_keys {
                debug!("Pressed: 0x{:02x} ({})", keycode, keycode);
                let keycode_xkb = hid_to_xkb_keycode(keycode);
                let sym = state.key_get_one_sym(keycode_xkb);

                let result: Option<(Keysym, Keysym, Option<String>)> = match compose_state.feed(sym)
                {
                    FeedResult::Ignored => {
                        let string = state.key_get_utf8(keycode_xkb);
                        Some((sym, sym, Some(string)))
                    }
                    FeedResult::Accepted => {
                        let status = compose_state.status();
                        debug!("Compose status: {status:?}");
                        match status {
                            Status::Nothing => {
                                let string = state.key_get_utf8(keycode_xkb);
                                Some((sym, sym, Some(string)))
                            }
                            Status::Composing => None,
                            Status::Composed => {
                                let composed_sym = compose_state.keysym().unwrap_or_default();
                                let string = compose_state.utf8().unwrap_or_default();
                                Some((sym, composed_sym, Some(string)))
                            }
                            Status::Cancelled => None,
                        }
                    }
                };

                if let Some((basic_keysym, composed_keysym, string)) = result {
                    // Change `Some("")` into `None`.
                    let string = string.filter(|string| !string.is_empty());

                    info!(
                        "Basic keysym: {basic_keysym:?}, composed keysym: {composed_keysym:?}, string: {:?}",
                        string.as_ref()
                    );

                    pressed_keys_to_strings.insert(basic_keysym, string.clone());
                    KEY_MESSAGE_CHANNEL
                        .send(KeyMessage {
                            keysym: composed_keysym,
                            string,
                            direction: KeyDirection::Down,
                        })
                        .await;
                    state.update_key(keycode_xkb, KeyDirection::Down);
                }
            }
        }
    }
}

pub fn keysym_to_xkb_keycode() {}

/// Based on https://github.com/xkbcommon/libxkbcommon/blob/6c67e3d41d3215ab1edd4406de215c7bf1f20c74/tools/how-to-type.c#L389
/// Fields reordered for `Ord` based on https://github.com/xkbcommon/libxkbcommon/blob/6c67e3d41d3215ab1edd4406de215c7bf1f20c74/tools/how-to-type.c#L404
#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct KeysymEntry {
    layout: xkb::LayoutIndex,
    level: xkb::LevelIndex,
    keycode: xkb::Keycode,
    mask: xkb::ModMask,
}

type KeysymEntries = Vec<KeysymEntry, PsramAllocator>;
type KeysymMap = BTreeMap<Keysym, KeysymEntries, PsramAllocator>;

/// Based on https://github.com/xkbcommon/libxkbcommon/blob/6c67e3d41d3215ab1edd4406de215c7bf1f20c74/tools/how-to-type.c#L434
fn lookup_keysym_entries(
    entries: &mut KeysymMap,
    keysym: xkb::Keysym,
    insert: bool,
) -> Option<&mut KeysymEntries> {
    match entries.entry(keysym) {
        Entry::Occupied(occupied) => Some(occupied.into_mut()),
        Entry::Vacant(vacant) if insert => Some(vacant.insert(Vec::new_in(&PSRAM_ALLOCATOR))),
        Entry::Vacant(_) => None,
    }
}

#[derive(Debug, Clone)]
pub struct HidKeycodeWithMods {
    keycode: u8,
    mod_mask: ModMask,
}

pub fn string_to_hid_keycodes(
    keymap: &xkb::Keymap,
    _compose_table: &xkb::compose::Table,
    string: &str,
) -> Result<Vec<HidKeycodeWithMods, PsramAllocator>, char> {
    #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
    struct KeycodeChoice {
        mod_mask: ModMask,
        hid_keycode: u8,
        xkb_keycode: xkb::Keycode,
    }

    let mut hid_keycodes = Vec::new_in(&PSRAM_ALLOCATOR);

    for character in string.chars() {
        let keysym = xkbcommon::xkb::utf32_to_keysym(character as u32);
        let mut chosen_keycode = None;

        keymap.key_for_each(|keymap, xkb_keycode| {
            for layout_index in 0..keymap.num_layouts_for_key(xkb_keycode) {
                for level_index in 0..keymap.num_levels_for_key(xkb_keycode, layout_index) {
                    let [current_keysym] =
                        keymap.key_get_syms_by_level(xkb_keycode, layout_index, level_index)
                    else {
                        // Multi-keysym levels are currently not handled.
                        continue;
                    };

                    if current_keysym == &keysym {
                        let mut masks: [ModMask; 32 /* TODO: Re-evaluate the appropriate size */] = Default::default();
                        let masks_len = keymap.key_get_mods_for_level(
                            xkb_keycode,
                            layout_index,
                            level_index,
                            &mut masks,
                        );
                        let masks = &mut masks[0..masks_len];
                        let hid_keycode = xkb_to_hid_keycode(xkb_keycode);

                        info!("Candidate for {character:?} ({keysym:?}) considered: XKB = 0x{xkb_keycode:02x?}, HID = 0x{hid_keycode:02x?}, layout = {layout_index:?}, level = {level_index:?}, masks = 0x{masks:02x?}", xkb_keycode = xkb_keycode.raw());

                        let Some(hid_keycode) = hid_keycode else {
                            warn!("Could not translate XKB {xkb_keycode:?} to an HID keycode, skipping candidate.");
                            continue;
                        };

                        if (0x65..USB_HID_LEFT_CTRL).contains(&hid_keycode) {
                            warn!("Skipping candidate because is in a poorly supported range of HID keycodes.");
                            continue;
                        }

                        let Some(&simplest_mod_mask) = masks.iter().min() else {
                            error!("No mod mask found for this key.");
                            continue;
                        };

                        let keycode_choice = KeycodeChoice {
                            mod_mask: simplest_mod_mask,
                            hid_keycode,
                            xkb_keycode,
                        };

                        if let Some(found_keycode) = &mut chosen_keycode {
                            *found_keycode = core::cmp::min(*found_keycode, keycode_choice);
                        } else {
                            chosen_keycode = Some(keycode_choice);
                        }
                    }
                }
            }
        });

        let Some(KeycodeChoice {
            mod_mask,
            hid_keycode,
            xkb_keycode,
        }) = chosen_keycode
        else {
            return Err(character);
        };

        warn!(
            "Candidate for {character:?} ({keysym:?}) chosen: XKB = 0x{xkb_keycode:02x?}, HID = 0x{hid_keycode:02x?}, mod_mask = 0x{mod_mask:02x?}",
            xkb_keycode = xkb_keycode.raw()
        );

        hid_keycodes.push(HidKeycodeWithMods {
            keycode: hid_keycode,
            mod_mask,
        });
    }

    Ok(hid_keycodes)
}

const fn hid_to_xkb_keycode_inner(rmk_keycode: u8) -> u8 {
    // https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/drivers/hid/hid-input.c?id=refs/tags/v6.18#n27
    const UNK: u8 = 240;
    #[rustfmt::skip]
    const HID_KEYBOARD: [u8; 256] = [
          0,   0,   0,   0,  30,  48,  46,  32,  18,  33,  34,  35,  23,  36,  37,  38,
         50,  49,  24,  25,  16,  19,  31,  20,  22,  47,  17,  45,  21,  44,   2,   3,
          4,   5,   6,   7,   8,   9,  10,  11,  28,   1,  14,  15,  57,  12,  13,  26,
         27,  43,  43,  39,  40,  41,  51,  52,  53,  58,  59,  60,  61,  62,  63,  64,
         65,  66,  67,  68,  87,  88,  99,  70, 119, 110, 102, 104, 111, 107, 109, 106,
        105, 108, 103,  69,  98,  55,  74,  78,  96,  79,  80,  81,  75,  76,  77,  71,
         72,  73,  82,  83,  86, 127, 116, 117, 183, 184, 185, 186, 187, 188, 189, 190,
        191, 192, 193, 194, 134, 138, 130, 132, 128, 129, 131, 137, 133, 135, 136, 113,
        115, 114, UNK, UNK, UNK, 121, UNK,  89,  93, 124,  92,  94,  95, UNK, UNK, UNK,
        122, 123,  90,  91,  85, UNK, UNK, UNK, UNK, UNK, UNK, UNK, 111, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, 179, 180, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, 111, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
         29,  42,  56, 125,  97,  54, 100, 126, 164, 166, 165, 163, 161, 115, 114, 113,
        150, 158, 159, 128, 136, 177, 178, 176, 142, 152, 173, 140, UNK, UNK, UNK, UNK,
    ];
    // https://cgit.freedesktop.org/xorg/driver/xf86-input-evdev/tree/src/evdev.c#n73
    const MIN_KEYCODE: u8 = 8;

    // TODO: The combination of these two operations should be precomputed
    // in a const expr into a single look-up table.
    HID_KEYBOARD[rmk_keycode as usize] + MIN_KEYCODE
    // xkb::Keycode::new((HID_KEYBOARD[rmk_keycode as usize] + MIN_KEYCODE) as u32)
}

pub const fn hid_to_xkb_keycode(hid_keycode: u8) -> xkb::Keycode {
    const HID_TO_XKB: [u8; 256] = {
        let mut hid_to_xkb = [0_u8; _];
        let mut hid: u8 = 0;

        loop {
            hid_to_xkb[hid as usize] = hid_to_xkb_keycode_inner(hid);
            let overflow;
            (hid, overflow) = hid.overflowing_add(1);
            if overflow {
                break;
            }
        }

        hid_to_xkb
    };

    xkb::Keycode::new(HID_TO_XKB[hid_keycode as usize] as u32)
}

pub const fn xkb_to_hid_keycode(xkb_keycode: xkb::Keycode) -> Option<u8> {
    const XKB_TO_HID: [u8; 256] = {
        let mut xkb_to_hid = [0_u8; _];
        let mut hid: u8 = 0;

        loop {
            xkb_to_hid[hid_to_xkb_keycode_inner(hid) as usize] = hid;
            let overflow;
            (hid, overflow) = hid.overflowing_add(1);
            if overflow {
                break;
            }
        }

        xkb_to_hid
    };

    if (xkb_keycode.raw() as usize) < XKB_TO_HID.len() {
        Some(XKB_TO_HID[xkb_keycode.raw() as usize])
    } else {
        None
    }
}

pub trait TryFromKeysym: Sized {
    fn try_from_keysym(k: Keysym) -> Option<Self>;
}

macro_rules! declare_consts_for_special_keys {
    ($($_char:literal # $name:ident # $($_qt:ident)|* # $($_winit:ident $(($_pos:ident))?)|*    # $($xkb:ident)|*;)*) => {
        impl TryFromKeysym for Key {
            fn try_from_keysym(k: Keysym) -> Option<Self> {
                match k {
                    $(
                        $(Keysym::$xkb => Some(Key::$name),)*
                    )*
                    _ => None
                }
            }
        }
    };
}

i_slint_common::for_each_special_keys!(declare_consts_for_special_keys);
