use std::{str::FromStr, thread, time::Duration};

use common::{asm::AsmOpcode, prelude::Register};

use crate::assembler::{AssembleError, Module, Symbol};
use crate::assemblerv2::Token;

pub struct Lexer {
    src: Vec<u8>,
    module_id: u64,
}

impl Lexer {
    pub fn new(src: String, module_id: u64) -> Self {
        Self {
            src: src.into_bytes().into_iter().rev().collect(),
            module_id,
        }
    }

    pub fn run(&mut self) -> Result<Vec<Token>, AssembleError> {
        let mut tokens = Vec::new();

        while !self.src.is_empty() {
            // thread::sleep(Duration::from_millis(10));
            // println!("{}", *self.src.last().unwrap() as char);
            // println!("{:?}", tokens);

            if self.src.ends_with(b"\n")
                || self.src.ends_with(b"\r")
                || self.src.ends_with(b" ")
                || self.src.ends_with(b",")
                || self.src.ends_with(b".")
                || self.src.ends_with(b":")
            {
                self.src.pop();
                continue;
            }

            if self.src.ends_with(b"//") {
                self.src.pop();
                self.src.pop();
                self.parse_comment();
                continue;
            }

            if self.src.ends_with(b"*/") {
                self.src.pop();
                self.src.pop();
                self.parse_comment_multiline();
                continue;
            }

            if self.src.ends_with(b"\"") {
                self.src.pop();
                tokens.push(Token::StringLit(self.parse_string_literal()));
                continue;
            }

            if self.src.ends_with(b"'") {
                self.src.pop();
                tokens.push(Token::CharLit(self.parse_char_literal()));
                continue;
            }

            if self.src.last().unwrap().is_ascii_digit() {
                tokens.push(Token::Immediate(self.parse_number()));
                continue;
            }

            if matches!(self.src.last().unwrap(), b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_') {
                let mut buffer = String::new();

                while matches!(self.src.last().unwrap(), b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_')
                {
                    buffer.push(self.src.pop().unwrap() as char);
                }

                if let Ok(opcode) = AsmOpcode::from_str(&buffer) {
                    tokens.push(Token::Opcode(opcode));
                    continue;
                }

                if let Ok(register) = Register::from_str(&buffer) {
                    tokens.push(Token::Register(register));
                    continue;
                }

                // Check for qualified symbol: identifier::identifier
                if self.src.ends_with(b"::") {
                    self.src.pop();
                    self.src.pop();
                    let mut rhs = String::new();
                    while matches!(
                        self.src.last(),
                        Some(b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_')
                    ) {
                        rhs.push(self.src.pop().unwrap() as char);
                    }
                    if rhs.is_empty() {
                        return Err(AssembleError::Generic);
                    }
                    tokens.push(Token::Symbol(Symbol {
                        name: rhs,
                        module: Module::Unresolved(buffer),
                    }));
                    continue;
                }

                tokens.push(Token::Symbol(Symbol {
                    name: buffer,
                    module: Module::Resolved(self.module_id),
                }));
                continue;
            }
        }

        Ok(tokens)
    }

    fn parse_number(&mut self) -> u32 {
        match self.src.last_chunk::<2>().map(|[x, y]| [*y, *x]).as_ref() {
            Some(b"0x") => {
                self.src.pop();
                self.src.pop();
                let mut n = 0u32;
                while let Some(&b) = self.src.last() {
                    if let Some(digit) = (b as char).to_digit(16) {
                        self.src.pop();
                        n = n * 16 + digit;
                    } else {
                        break;
                    }
                }
                n
            }
            Some(b"0b") => {
                self.src.pop();
                self.src.pop();
                let mut n = 0u32;
                while let Some(&b) = self.src.last() {
                    if b == b'0' || b == b'1' {
                        self.src.pop();
                        n = n * 2 + (b - b'0') as u32;
                    } else {
                        break;
                    }
                }
                n
            }
            Some(b"0o") => {
                self.src.pop();
                self.src.pop();
                let mut n = 0u32;
                while let Some(&b) = self.src.last() {
                    if matches!(b, b'0'..=b'7') {
                        self.src.pop();
                        n = n * 8 + (b - b'0') as u32;
                    } else {
                        break;
                    }
                }
                n
            }
            // decimal number
            _ if let Some(x) = self.src.last()
                && x.is_ascii_digit() =>
            {
                let mut n = 0u32;
                while let Some(&b) = self.src.last() {
                    if matches!(b, b'0'..=b'9') {
                        self.src.pop();
                        n = n * 10 + (b - b'0') as u32;
                    } else {
                        break;
                    }
                }
                n
            }
            _ => panic!("Invalid syntax"),
        }
    }

    // returns nothing as the assembler discards comments
    fn parse_comment(&mut self) {
        while !self.src.is_empty() && !self.src.ends_with(b"\n") {
            self.src.pop();
        }
    }

    fn parse_comment_multiline(&mut self) {
        while !self.src.is_empty() && !self.src.ends_with(b"*/") {
            self.src.pop();
        }
    }

    fn parse_string_literal(&mut self) -> String {
        let mut result = String::new();
        while self.src.last().is_some() && !self.src.ends_with(b"\"") {
            let ch = self.src.pop().unwrap();
            if ch == b'\\' {
                let escaped = self
                    .src
                    .pop()
                    .expect("A file should never end with a backslash!");
                result.push(match escaped {
                    b'n' => '\n',
                    b't' => '\t',
                    b'r' => '\r',
                    b'"' => '\"',
                    b'\\' => '\\',
                    _ => escaped as char,
                });

                continue;
            }

            result.push(ch as char);
        }

        self.src.pop().unwrap();

        result
    }

    fn parse_char_literal(&mut self) -> char {
        let ch = self
            .src
            .pop()
            .expect("Unexpected EOF while parsing char literal");

        if ch == b'\\' {
            let escaped = self
                .src
                .pop()
                .expect("A file should never end with a backslash!");

            assert!(
                self.src
                    .pop()
                    .expect("Unexpected EOF while parsing char literal")
                    != b'\'',
                "unterminated char literal"
            );

            match escaped {
                b'n' => '\n',
                b't' => '\t',
                b'r' => '\r',
                b'"' => '\"',
                b'\\' => '\\',
                _ => escaped as char,
            }
        } else {
            assert!(
                self.src
                    .pop()
                    .expect("Unexpected EOF while parsing char literal")
                    != b'\'',
                "unterminated char literal"
            );
            ch as char
        }
    }
}
