use std::collections::HashMap;

use automata::{
    automatan::{fa::Fa, pda::Pda, tm::Tm}, delta_lower, epsilon, gamma_upper, loader::{self, Context, Machine, Span, Spanned, lexer::Lexer}, sigma_upper
};

use wasm_bindgen::prelude::wasm_bindgen;

#[wasm_bindgen]
pub fn init() {
    console_error_panic_hook::set_once();
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub enum Kind {
    Ident = "ident",
    Keyword = "keyword",
    String = "string",
    Error = "error",
    Comment = "comment",
    Punc = "punc",

    LPar = "lpar",
    LBrace = "lbrace",
    LBracket = "lbracket",

    RPar = "rpar",
    RBrace = "rbrace",
    RBracket = "rbracket",
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub struct Tok {
    pub start: usize,
    pub end: usize,
    pub scope_level: usize,
    pub kind: Kind,
}

#[wasm_bindgen]
pub fn lex(input: &str) -> Vec<Tok> {
    let mut scope_level = 0;
    let mut index_utf16 = 0;
    let mut index_utf8 = 0;
    Lexer::new(input)
        .map(|Spanned(tok, Span(start_utf8, end_utf8))| {
            let since_last = &input[index_utf8..start_utf8];
            let since_start = &input[start_utf8..end_utf8];

            index_utf8 = end_utf8;
            let start = index_utf16 + since_last.chars().map(char::len_utf16).sum::<usize>();
            let end = start + since_start.chars().map(char::len_utf16).sum::<usize>();
            index_utf16 = end;

            let Ok(tok) = tok else {
                return Tok {
                    start,
                    end,
                    kind: Kind::Error,
                    scope_level,
                };
            };
            use automata::loader::lexer::Token;
            let kind = match tok {
                Token::LPar => Kind::LPar,
                Token::RPar => Kind::RPar,
                Token::LBrace => Kind::LBrace,
                Token::RBrace => Kind::RBrace,
                Token::LBracket => Kind::LBracket,
                Token::RBracket => Kind::RBracket,
                Token::Tilde => Kind::Keyword,
                Token::Eq => Kind::Punc,
                Token::Comma => Kind::Punc,
                Token::Or => Kind::Punc,
                Token::Plus => Kind::Punc,
                Token::Dash => Kind::Punc,
                Token::Star => Kind::Punc,
                Token::And => Kind::Punc,
                Token::LSmallArrow => Kind::Punc,
                Token::LBigArrow => Kind::Punc,
                Token::Comment(_) => Kind::Comment,
                Token::String(_, _, _) => Kind::String,
                Token::Ident(_)
                    if input[..start_utf8]
                        .split("\n")
                        .last()
                        .unwrap_or_default()
                        .trim()
                        .is_empty() =>
                {
                    Kind::Keyword
                }

                // ugly hack to keep single ascii letters non keyworded for user
                Token::Ident(ident) if ident.is_ascii() && ident.len() == 1 => Kind::Ident,
                Token::Ident(
                    epsilon!(pat) | delta_lower!(pat) | sigma_upper!(pat) | gamma_upper!(pat),
                ) => Kind::Keyword,
                Token::Ident(_) => Kind::Ident,
                Token::LineEnd => Kind::Punc,
            };

            let scope_level = match kind {
                Kind::LPar | Kind::LBrace | Kind::LBracket => {
                    scope_level = scope_level.saturating_add(1);
                    scope_level.saturating_sub(1)
                }
                Kind::RPar | Kind::RBrace | Kind::RBracket => {
                    scope_level = scope_level.saturating_sub(1);
                    scope_level
                }
                _ => scope_level,
            };
            Tok {
                start,
                end,
                kind,
                scope_level,
            }
        })
        .collect()
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub enum LogLevel {
    Help = "help",
    Info = "info",
    Warning = "warning",
    Error = "error",
}

#[wasm_bindgen(getter_with_clone)]
#[derive(Clone)]
pub struct CompileLog {
    pub level: LogLevel,
    pub message: String,
    pub start: Option<usize>,
    pub end: Option<usize>,
}

#[wasm_bindgen(getter_with_clone)]
pub struct CompileResult {
    pub log: Vec<CompileLog>,
    pub ansi_log: String,
    pub machine: Option<String>,
}

trait FixupSpan{
    fn fixup(&mut self, func: impl FnMut(Span) -> Span);
}

impl<'a> FixupSpan for Machine<'a>{
    fn fixup(&mut self, func: impl FnMut(Span) -> Span) {
        match self{
            Machine::Fa(fa) => fa.fixup(func),
            Machine::Pda(pda) => pda.fixup(func),
            Machine::Tm(tm) => tm.fixup(func),
        }
    }
}
impl<'a> FixupSpan for Fa<'a>{
    fn fixup(&mut self, mut func: impl FnMut(Span) -> Span) {
        self.alphabet.values_mut().for_each(|v| v.definition = func(v.definition));
        self.states.values_mut().for_each(|v| v.definition = func(v.definition));
        self.final_states.values_mut().for_each(|v| v.definition = func(v.definition));
        self.transitions.values_mut().flat_map(|v|v.iter_mut()).for_each(|e|{
            e.transition = func(e.transition);
            e.function = func(e.function);
        });
    }
}

impl<'a> FixupSpan for Pda<'a>{
    fn fixup(&mut self, mut func: impl FnMut(Span) -> Span) {
        self.alphabet.values_mut().for_each(|v| v.definition = func(v.definition));
        self.states.values_mut().for_each(|v| v.definition = func(v.definition));
        self.symbols.values_mut().for_each(|v| v.definition = func(v.definition));
        self.final_states.as_mut().unwrap_or(&mut HashMap::new()).values_mut().for_each(|v| v.definition = func(v.definition));
        self.transitions.values_mut().flat_map(|v|v.iter_mut()).for_each(|e|{
            e.transition = func(e.transition);
            e.function = func(e.function);
        });
    }
}

impl<'a> FixupSpan for Tm<'a>{
    fn fixup(&mut self, mut func: impl FnMut(Span) -> Span) {
        self.states.values_mut().for_each(|v| v.definition = func(v.definition));
        self.symbols.values_mut().for_each(|v| v.definition = func(v.definition));
        self.final_states.values_mut().for_each(|v| v.definition = func(v.definition));
        self.transitions.values_mut().flat_map(|v|v.iter_mut()).for_each(|e|{
            e.transition = func(e.transition);
            e.function = func(e.function);
        });
    }
}



#[wasm_bindgen]
pub fn compile(input: &str) -> CompileResult {
    let mut ctx = Context::new(input);
    let result = automata::loader::parse_universal(&mut ctx);

    let machine = result.map(|mut result| {
        result.fixup(|span|Span(input[..span.0].chars().map(char::len_utf16).sum(), input[..span.1].chars().map(char::len_utf16).sum()));
        serde_json::to_string(&result).unwrap()
    });

    use std::fmt::Write;
    let ansi_log = ctx.logs_display().fold(String::new(), |mut s, e| {
        writeln!(&mut s, "{e}").unwrap();
        s
    });

    let log = ctx
        .into_logs()
        .into_entries()
        .map(|e| CompileLog {
            level: match e.level {
                loader::log::LogLevel::Help => LogLevel::Help,
                loader::log::LogLevel::Info => LogLevel::Info,
                loader::log::LogLevel::Warning => LogLevel::Warning,
                loader::log::LogLevel::Error => LogLevel::Error,
            },
            message: e.message,
            start: e
                .span
                .map(|span| input[..span.0].chars().map(char::len_utf16).sum()),
            end: e
                .span
                .map(|span| input[..span.1].chars().map(char::len_utf16).sum()),
        })
        .collect();

    CompileResult {
        log,
        ansi_log,
        machine,
    }
}
