use std::{collections::BTreeMap, io::Read, sync::RwLock};

fn read1() -> Option<u8> {
    let mut buf = [0];
    match std::io::stdin().lock().read_exact(&mut buf) {
        Ok(_) => Some(buf[0]),
        Err(_) => None,
    }
}

pub fn read(debug: bool) -> Option<KbInput<'static>> {
    let mut reader = EscapingStdinReader::new(et());
    loop {
        let Some(b) = read1() else {
            break None;
        };
        match reader.process_byte(b) {
            ByteProcessingResult::Done(kb) => {
                if debug {
                    if let KbInput::Escape(e) = &kb {
                        for k in e.keys.iter() {
                            print!("{k} ")
                        }
                    }
                    println!("{}\r", kb.as_bytes().escape_ascii());
                }
                if let KbInput::InvalidEscape(x) = &kb
                    && x.len() == 1
                {
                    break Some(KbInput::Key([x[0]]));
                }
                break Some(kb);
            }
            ByteProcessingResult::Continue(r) => reader = r,
        }
    }
}

struct EscapingStdinReader<'a> {
    buf: Vec<u8>,
    trie: &'a EscapeTrie,
}

enum ByteProcessingResult<'a> {
    Done(KbInput<'a>),
    Continue(EscapingStdinReader<'a>),
}

impl<'a> EscapingStdinReader<'a> {
    pub fn new(trie: &'a EscapeTrie) -> Self {
        Self {
            buf: Vec::new(),
            trie,
        }
    }

    pub fn process_byte(mut self, byte: u8) -> ByteProcessingResult<'a> {
        match self.trie {
            EscapeTrie::Done(_) => ByteProcessingResult::Done(KbInput::Key([byte])),
            EscapeTrie::More(trie) => {
                self.buf.push(byte);
                match trie.get(&byte) {
                    Some(EscapeTrie::Done(keys)) => {
                        ByteProcessingResult::Done(KbInput::Escape(Escape {
                            keys: &keys[..],
                            value: self.buf,
                        }))
                    }
                    Some(trie) => {
                        self.trie = trie;
                        ByteProcessingResult::Continue(self)
                    }
                    None => ByteProcessingResult::Done(KbInput::InvalidEscape(self.buf)),
                }
            }
        }
    }
}

enum EscapeTrie {
    Done(Vec<&'static str>),
    More(BTreeMap<u8, EscapeTrie>),
}

#[derive(Clone, Debug)]
pub enum KbInput<'a> {
    Key([u8; 1]),
    Escape(Escape<'a>),
    InvalidEscape(Vec<u8>),
}

impl<'a> KbInput<'a> {
    pub fn as_bytes(&'a self) -> &'a [u8] {
        match self {
            KbInput::Key(x) => &x[..],
            KbInput::Escape(e) => &e.value[..],
            KbInput::InvalidEscape(e) => &e[..],
        }
    }
}

#[derive(Clone, Debug)]
pub struct Escape<'a> {
    pub keys: &'a [&'a str],
    pub value: Vec<u8>,
}

use terminfo_lean::parse::Terminfo;

static TERMINFO: RwLock<Option<&'static Terminfo<'static>>> = RwLock::new(None);
static ESCAPE_TRIE: RwLock<Option<&'static EscapeTrie>> = RwLock::new(None);

fn parse_terminfo() -> Result<Terminfo<'static>, ()> {
    let term = std::env::var_os("TERM").unwrap_or_else(|| "xterm".into());
    let terminfo_file_path = terminfo_lean::locate::locate(&term)
        .map_err(|e| println!("failed to locate terminfo file for terminal {term:?}: {e:?}",))?;
    let mut terminfo_file = std::fs::File::open(&terminfo_file_path).map_err(|e| {
        println!("failed to open terminfo file at location {terminfo_file_path:?}: {e:?}")
    })?;
    let mut buf = Vec::new();
    terminfo_file.read_to_end(&mut buf).map_err(|e| {
        println!("failed to read terminfo file at location {terminfo_file_path:?}: {e:?}")
    })?;
    buf.shrink_to_fit();
    let terminfo = terminfo_lean::parse::parse(buf.leak()).map_err(|e| {
        println!("failed to parse terminfo file at location {terminfo_file_path:?}: {e:?}")
    })?;
    Ok(terminfo)
}

fn parse_terminfo_backup() -> Terminfo<'static> {
    todo!("panic-safe backup terminfo")
}

pub fn setup() {
    let ti = parse_terminfo().unwrap_or_else(|_| {
        println!("using backup terminfo (might not be correct for this terminal)");
        parse_terminfo_backup()
    });
    let ti: &'static _ = Box::leak(Box::new(ti));

    TERMINFO.clear_poison();
    *TERMINFO.write().unwrap() = Some(ti);

    let et = Box::leak(Box::new(EscapeTrie::from(ti)));
    ESCAPE_TRIE.clear_poison();
    *ESCAPE_TRIE.write().unwrap() = Some(et);
}

pub fn ti() -> &'static Terminfo<'static> {
    TERMINFO.read().unwrap().unwrap()
}

fn et() -> &'static EscapeTrie {
    ESCAPE_TRIE.read().unwrap().unwrap()
}

fn is_parametrized(x: &[u8]) -> bool {
    let mut pct = false;
    for &b in x {
        if b == b'%' {
            pct = !pct;
        } else if pct {
            return true;
        }
    }
    false
}

fn trie_from_words(words: Vec<(&'static str, &[u8])>) -> EscapeTrie {
    let mut tree = BTreeMap::new();
    let mut all_empty = true;

    for (key, val) in words.iter() {
        if let Some(byte) = val.first() {
            all_empty = false;
            tree.entry(byte)
                .or_insert_with(Vec::new)
                .push((*key, &val[1..]));
        }
    }

    if all_empty {
        EscapeTrie::Done(words.into_iter().map(|x| x.0).collect())
    } else {
        let trie = tree
            .into_iter()
            .map(|(k, v)| (*k, trie_from_words(v)))
            .collect();
        EscapeTrie::More(trie)
    }
}

impl From<&Terminfo<'static>> for EscapeTrie {
    fn from(ti: &Terminfo<'static>) -> Self {
        let w: Vec<(&'static str, &'static [u8])> = ti
            .strings
            .iter()
            .filter(|(k, v)| k.starts_with("k") && !is_parametrized(v))
            .map(|(k, v)| (*k, *v))
            .collect();
        trie_from_words(w)
    }
}