asa 1.0.0

//! Converts a string into a vector of tokens, resolving includes along the way

use crate::{args, terminate};
use crate::{
    asm_error, asm_error_no_terminate, asm_hint,
    tokens::{Info, LabelOffset, Token, TokenVariant},
};
use colored::Colorize;
use std::cell::RefCell;
use std::path::{self, Path, PathBuf};
use std::{env, fs};

use log::LevelFilter;
use unescape::unescape;

thread_local! {
    // Array of all included files
    pub static FILES: RefCell<Vec<PathBuf>> = const { RefCell::new(vec![]) };
}

#[derive(Debug, PartialEq, Eq)]
enum Context {
    None,
    LineComment,
    String,
    Char,
    EscapedChar,
    BlockComment,

    SubleqOrNegativeOrLabelArrow,
    LabelArrow,
    HexOrDec,
    Hex,
    Dec,

    Label,
    DontConsume,

    MacroDeclaration,
    MacroCall,
    Relative,

    Inclusion,

    AsteriskOrBlockComment,
    PossibleBlockCommentEnd,
}

/// Does some basic and safe sanitisation. It's fine to apply it multiple times.
pub fn generic_sanitisation(text: &str) -> String {
    text.replace("\r\n", "\n").replace("\t", "    ")
}

fn is_valid_macro_name(c: char) -> bool {
    !c.is_whitespace() && !matches!(c, '(' | ')' | '{' | '}' | '[' | ']')
}
fn is_valid_label_name(c: char) -> bool {
    c.is_alphanumeric() || c == '_' || c == ':' || c == '?' || c == '.'
}

/*
    This code really should be updated to use utils::IterVec, would make is significantly cleaner
    but it works, is readable enough and I can't really be bothered
*/
/// Update the state machine. Returns a new context, possibly the character that should be added
/// to the buffer and possibly a token if enough information has been gathered
fn updated_context(
    context: &Context,
    buffer: &str,
    cur_char: char,
    info: &Info,
) -> (Context, Option<char>, Option<TokenVariant>) {
    match context {
        Context::DontConsume => unreachable!(),

        Context::None => match cur_char {
            '\n' => (Context::None, None, Some(TokenVariant::Linebreak)),
            ' ' => (Context::None, None, None),
            ';' => (Context::LineComment, None, None),
            '-' => (Context::SubleqOrNegativeOrLabelArrow, Some(cur_char), None),
            '0' => (Context::HexOrDec, Some(cur_char), None), // This will cause hex numbers to have leading zeros
            c if c.is_ascii_digit() => (Context::Dec, Some(c), None),
            c if c.is_ascii_alphabetic() || c == '_' || c == '.' => (Context::Label, Some(c), None),

            '*' => (Context::AsteriskOrBlockComment, None, None),
            '=' => (Context::None, None, Some(TokenVariant::Equals)),

            '{' => (Context::None, None, Some(TokenVariant::Scope)),
            '}' => (Context::None, None, Some(TokenVariant::Unscope)),
            '[' => (Context::None, None, Some(TokenVariant::MacroBodyStart)),
            ']' => (Context::None, None, Some(TokenVariant::MacroBodyEnd)),

            '(' => (Context::None, None, Some(TokenVariant::BraceOpen)),
            ')' => (Context::None, None, Some(TokenVariant::BraceClose)),
            '$' => (Context::Relative, None, None),
            '&' => (
                Context::None,
                None,
                Some(TokenVariant::Relative { offset: 1 }),
            ),
            '\'' => (Context::Char, None, None),
            '"' => (Context::String, None, None),

            '@' => (Context::MacroDeclaration, None, None),
            '!' => (Context::MacroCall, None, None),
            '#' => (Context::Inclusion, None, None),

            '?' => {
                asm_error_no_terminate!(info, "Unexpected character");
                asm_hint!("Labels may not start with a '?'");
                terminate!();
            }
            _ => asm_error!(info, "Unexpected character"),
        },

        Context::AsteriskOrBlockComment => match cur_char {
            '*' => (Context::BlockComment, None, None),
            _ => (Context::DontConsume, None, Some(TokenVariant::Asterisk)),
        },

        Context::BlockComment => match cur_char {
            '*' => (Context::PossibleBlockCommentEnd, None, None),
            '\n' => (Context::BlockComment, None, Some(TokenVariant::Linebreak)),
            _ => (Context::BlockComment, None, None),
        },
        Context::PossibleBlockCommentEnd => match cur_char {
            '*' => (Context::None, None, None),
            _ => (Context::BlockComment, None, None),
        },

        Context::LineComment => match cur_char {
            '\n' => (Context::DontConsume, None, None),
            _ => (Context::LineComment, None, None),
        },
        Context::Inclusion => match cur_char {
            '\n' => (
                Context::DontConsume,
                None,
                Some(TokenVariant::Inclusion {
                    path: buffer.to_owned(),
                }),
            ),
            _ => (Context::Inclusion, Some(cur_char), None),
        },

        Context::SubleqOrNegativeOrLabelArrow => match cur_char {
            '=' => (Context::None, None, Some(TokenVariant::Subleq)),
            '0' => (Context::HexOrDec, Some(cur_char), None),
            c if c.is_ascii_digit() => (Context::Dec, Some(cur_char), None),
            'a' | 'b' | 'c' => (Context::LabelArrow, Some(cur_char), None),
            '>' => (
                Context::None,
                Some(cur_char),
                Some(TokenVariant::LabelArrow {
                    offset: LabelOffset::Int(0),
                }),
            ),
            _ => {
                asm_error!(
                    info,
                    "Unexpected character, for Subleq use '-=', for label use ->"
                )
            }
        },

        Context::LabelArrow => match cur_char {
            '>' => {
                let ch: char = buffer.chars().nth(1).unwrap();
                if ch == 'a' || ch == 'b' || ch == 'c' {
                    return (
                        Context::None,
                        Some(cur_char),
                        Some(TokenVariant::LabelArrow {
                            offset: LabelOffset::Char(ch),
                        }),
                    );
                }
                asm_error!(info, "Unexpected character");
            }
            _ => asm_error!(info, "Unexpected character"),
        },

        Context::HexOrDec => match cur_char {
            'x' => (Context::Hex, None, None),
            c if c.is_ascii_digit() => (Context::Dec, Some(c), None),
            c if c.is_ascii_alphabetic() => {
                asm_error_no_terminate!(
                    info,
                    "Unexpected character when defining Hex or Dec literal",
                );
                asm_hint!("Labels may not start with a number");
                terminate!();
            }
            _ => (
                Context::DontConsume,
                None,
                Some(TokenVariant::DecLiteral {
                    value: buffer.parse::<i32>().unwrap(),
                }),
            ),
        },

        Context::Hex => match cur_char {
            c if c.is_ascii_hexdigit() => (Context::Hex, Some(c), None),
            c if c.is_ascii_alphabetic() => {
                asm_error!(info, "Unexpected character when defining Hex literal")
            }
            _ => (
                Context::DontConsume,
                None,
                Some(TokenVariant::HexLiteral {
                    value: buffer.to_owned(),
                }),
            ), // may cause issues
        },

        Context::Dec => match cur_char {
            c if c.is_ascii_digit() => (Context::Dec, Some(c), None),
            c if c.is_ascii_alphabetic() => {
                asm_error!(info, "Unexpected character when defining Dec literal")
            }
            _ => (
                Context::DontConsume,
                None,
                Some(TokenVariant::DecLiteral {
                    value: buffer
                        .parse::<i32>()
                        .unwrap_or_else(|_| asm_error!(&info, "Invalid decimal literal")),
                }),
            ),
        },
        Context::Label => match cur_char {
            c if is_valid_label_name(c) => (Context::Label, Some(cur_char), None),
            _ => (
                Context::DontConsume,
                None,
                Some(TokenVariant::Label {
                    name: buffer.to_owned(),
                }),
            ),
        },

        Context::MacroDeclaration => match cur_char {
            c if is_valid_macro_name(c) => (Context::MacroDeclaration, Some(cur_char), None),

            _ => (
                Context::DontConsume,
                None,
                Some(TokenVariant::MacroDeclaration {
                    name: buffer.to_owned(),
                }),
            ),
        },
        Context::MacroCall => match cur_char {
            c if is_valid_macro_name(c) => (Context::MacroCall, Some(cur_char), None),
            _ => (
                Context::DontConsume,
                None,
                Some(TokenVariant::MacroCall {
                    name: buffer.to_owned(),
                }),
            ),
        },
        Context::Char => match cur_char {
            // Not great
            '\\' => (Context::EscapedChar, None, None),
            '\'' => (
                Context::None,
                None,
                Some(TokenVariant::CharLiteral {
                    value: buffer.chars().next().unwrap_or_else(|| {
                        asm_error!(&info, "Char literal does not contain a character")
                    }),
                }),
            ),
            _ => (Context::Char, Some(cur_char), None),
        },
        Context::EscapedChar => match cur_char {
            '\'' => (
                Context::None,
                None,
                Some(TokenVariant::CharLiteral {
                    value: unescape(&format!("\\{}", buffer.chars().next().unwrap()))
                        .unwrap_or_else(|| asm_error!(&info, "Invalid escape sequence"))
                        .chars()
                        .next()
                        .unwrap(),
                }),
            ),
            _ => (Context::EscapedChar, Some(cur_char), None),
        },
        Context::String => match cur_char {
            '"' => (
                Context::None,
                None,
                Some(TokenVariant::StrLiteral {
                    value: buffer.to_owned(),
                }),
            ),
            _ => (Context::String, Some(cur_char), None),
        },

        Context::Relative => match cur_char {
            c if c.is_ascii_digit() => (Context::Relative, Some(c), None),
            _ => {
                let offset = if !buffer.is_empty() {
                    buffer
                        .parse::<i32>()
                        .unwrap_or_else(|_| crate::error!("Relative offsets must be decimal"))
                } else {
                    asm_error!(info, "Expected an offset",);
                };

                (
                    Context::DontConsume,
                    None,
                    Some(TokenVariant::Relative { offset }),
                )
            }
        },
    }
}

/// Takes a string, recursively solves includes and tokenises everything whilst
/// keeping track of the metadata of each token, like the file it originates from
pub fn tokenise(text: String, path: String) -> Vec<Token> {
    FILES.set(vec![Path::new(&path).to_path_buf()]);

    // Includes are resolved relative to the file being assembled
    let base_dir = path::absolute(Path::new(&path).to_path_buf().parent().unwrap())
        .unwrap_or_else(|_| env::current_dir().unwrap());

    let result = recursive_tokenisation(
        text,
        0,
        &mut vec![Path::new(&path).to_path_buf()],
        &base_dir,
    );
    log::debug!("Included files");
    if log::max_level() >= LevelFilter::Debug {
        FILES.with_borrow(|files| println!("{files:#?}"));
    }
    result
}

fn fix_include_path(path: &mut PathBuf) {
    // When trying to import a folder, it looks for a file named Lib.sbl in the folder
    if path.is_dir() {
        path.push("Lib");
    }

    // Imports in sublang don't need an extension but the filesystem does
    if path.extension().is_none() {
        path.set_extension("sbl");
    }
}

fn include(
    name: &String,
    currently_imported: &mut Vec<PathBuf>,
    base_dir: &Path,
) -> Option<Vec<Token>> {
    let mut path = base_dir.to_path_buf().join(name);
    fix_include_path(&mut path);
    if !path.is_file() {
        let libs_path = if args::exist() {
            &args::get().libs_path
        } else {
            "./subleq/libs"
        };
        let mut new_path = Path::new(libs_path).to_path_buf().join(name);
        fix_include_path(&mut new_path);
        if !new_path.is_file() {
            log::error!(
                "File to include not found locally: {} or in the libraries folder: {}",
                path.display(),
                new_path.display()
            );
            log::info!("Make sure the library path is correctly set using the '-l' argument");
            terminate!();
        }
        path = new_path;
    }

    let exists = FILES.with_borrow_mut(|files| {
        if !files.contains(&path) {
            files.push(path.clone());
            return false;
        }
        true
    });

    if !exists {
        let contents = fs::read_to_string(&path).unwrap_or_else(|_| {
            crate::error!("Couldn't include the file: '{path:?}'");
        });
        Some(recursive_tokenisation(
            contents,
            FILES.with_borrow_mut(|files| files.len() - 1),
            currently_imported,
            base_dir,
        ))
    } else {
        None
    }
}

/// Recursion occurs on file includes
fn recursive_tokenisation(
    mut text: String,
    file_idx: usize,
    currently_imported: &mut Vec<PathBuf>,
    base_dir: &Path,
) -> Vec<Token> {
    let mut result_tokens: Vec<Token> = Vec::new();

    text = generic_sanitisation(&text);
    text.push('\n'); // Little hack

    let mut context: Context = Context::None;
    let mut buffer: String = String::new();
    let mut info: Info = Info {
        start_char: 0,
        length: 0,
        line_number: 1,
        file: file_idx,
        sourceline_suffix: None,
    };
    let mut column = 1;
    for c in text.chars() {
        loop {
            let (new_context, add_to_buffer, variant_to_add) =
                updated_context(&context, &buffer, c, &info);

            context = new_context;
            info.start_char = column - info.length;

            if let Some(ch) = add_to_buffer {
                buffer.push(ch)
            }

            if let Some(var) = &variant_to_add {
                if let TokenVariant::Inclusion { path } = var {
                    if let Some(mut toks) = include(path, currently_imported, base_dir) {
                        result_tokens.append(&mut toks);
                    }
                }

                if context == Context::None {
                    info.length += 1;
                }
                let token = Token {
                    info: info.clone(),
                    variant: var.clone(),
                    origin_info: vec![],
                };

                if let TokenVariant::Inclusion { .. } = token.variant {
                } else {
                    // Consecutive newlines do not carry any information

                    if let Some(x) = result_tokens.last() {
                        if token.variant != TokenVariant::Linebreak
                            || x.variant != TokenVariant::Linebreak
                        {
                            result_tokens.push(token);
                        }
                    } else {
                        result_tokens.push(token);
                    }
                }

                // Reset
                buffer.clear();
                if let TokenVariant::Linebreak = var {
                    info.line_number += 1;
                    info.start_char = 0;
                    info.length = 0;
                    column = 0;
                    break;
                }
            }

            if let Context::DontConsume = context {
                context = Context::None;
                continue;
            }

            break;
        }

        if context == Context::None {
            info.length = 0;
        } else {
            info.length += 1;
        }
        column += 1;
    }

    result_tokens
}