uxn-tal 0.1.2

//! Lexer for TAL assembly language

use crate::error::{AssemblerError, Result};

/// Token types in TAL assembly
#[derive(Debug, Clone, PartialEq)]
pub enum Token {
    /// Hexadecimal literal (e.g., #1234, #ab)
    HexLiteral(String),
    /// Raw hexadecimal byte (e.g., 20, ff)
    RawHex(String),
    /// Decimal literal (e.g., 42)
    DecLiteral(String),
    /// Binary literal (e.g., #b10101010)
    BinLiteral(String),
    /// Character literal (e.g., 'A')
    CharLiteral(char),
    /// Instruction/opcode (e.g., ADD, LDA2k)
    Instruction(String),
    /// Label definition (e.g., @main)
    LabelDef(String),
    /// Label reference (e.g., ;main)
    LabelRef(String),
    /// Sublabel definition (e.g., &loop)
    SublabelDef(String),
    /// Sublabel reference (e.g., ,loop)
    SublabelRef(String),
    /// Relative address reference (e.g., /loop)
    RelativeRef(String),
    /// Conditional jump reference (e.g., ?loop)
    ConditionalRef(String),
    /// Conditional operator (e.g., ?)
    ConditionalOperator,
    /// Conditional block start (e.g., ?{)
    ConditionalBlockStart,
    /// Raw address reference (e.g., =label)
    RawAddressRef(String),
    /// JSR call reference (e.g., !label)
    JSRRef(String),
    /// Hyphen address reference (e.g., -Screen/auto)
    HyphenRef(String),
    /// Padding directive (e.g., |0100)
    Padding(u16),
    PaddingLabel(String),
    /// Skip bytes directive (e.g., $2)
    Skip(u16),
    /// Device access (e.g., .Screen/width)
    DeviceAccess(String, String), // device, field

    /// Macro definition (e.g., MACRO)
    MacroDef(String),

    /// Macro call (e.g., MACRO)
    MacroCall(String),
    /// Dot reference - generates LIT + 8-bit address (like uxnasm's '.' rune)
    DotRef(String),
    /// Semicolon reference - generates LIT2 + 16-bit address (like uxnasm's ';' rune)  
    SemicolonRef(String),
    /// Equals reference - generates 16-bit address directly (like uxnasm's '=' rune)
    EqualsRef(String),
    /// Comma reference - generates LIT + relative 8-bit address (like uxnasm's ',' rune)
    CommaRef(String),
    /// Underscore reference - generates relative 8-bit address (like uxnasm's '_' rune)
    UnderscoreRef(String),
    /// Question reference - generates conditional jump (like uxnasm's '?' rune)
    QuestionRef(String),
    /// Exclamation reference - generates JSR call (like uxnasm's '!' rune)
    ExclamationRef(String),
    /// Brace open
    BraceOpen,
    /// Brace close
    BraceClose,
    /// Bracket open (inline assembly)
    BracketOpen,
    /// Bracket close (inline assembly)
    BracketClose,
    /// Include directive (e.g., ~filename.tal)
    Include(String),
    /// Raw string literal
    RawString(String),
    /// Comment
    Comment(String),
    /// Newline
    Newline,
    /// End of file
    Eof,
}

/// Token with position information
#[derive(Debug, Clone, PartialEq)]
pub struct TokenWithPos {
    pub token: Token,
    pub line: usize,
    pub start_pos: usize,
    pub end_pos: usize,
}

/// Lexer for TAL assembly language
pub struct Lexer {
    input: String,
    position: usize,
    line: usize,
    path: Option<String>,
    position_on_line: usize,
    /// Next sublabel for conditional operator, if any
    next_conditional_sublabel: Option<String>,
}

impl Lexer {
    // get_column removed; use position_on_line directly
    pub fn new(input: String, path: Option<String>) -> Self {
        Self {
            input,
            position: 0,
            line: 1,
            path,
            position_on_line: 1,
            next_conditional_sublabel: None,
        }
    }

    /// Get the current line content for error reporting
    fn get_current_line(&self) -> String {
        let lines: Vec<&str> = self.input.lines().collect();
        if self.line > 0 && self.line <= lines.len() {
            lines[self.line - 1].to_string()
        } else {
            String::new()
        }
    }

    /// Create a syntax error with current line information
    fn syntax_error(&self, message: String) -> AssemblerError {
        AssemblerError::SyntaxError {
            path: self.path.clone().unwrap_or_default(),
            line: self.line,
            position: self.position_on_line,
            message,
            source_line: self.get_current_line(),
        }
    }

    pub fn line(&self) -> usize {
        self.line
    }

    /// Tokenize the entire input
    pub fn tokenize(&mut self) -> Result<Vec<TokenWithPos>> {
        let mut tokens = Vec::new();

        while self.position < self.input.len() {
            let start_line = self.line;
            let start_pos = self.position_on_line;
            let token_start_position = self.position;

            // Handle newlines and whitespace
            let ch = self.current_char();
            if ch == '\n' {
                tokens.push(TokenWithPos {
                    token: Token::Newline,
                    line: start_line,
                    start_pos,
                    end_pos: start_pos,
                });
                self.advance();
                continue;
            }
            if ch.is_whitespace() && ch != '\n' {
                self.advance();
                continue;
            }

            // Handle comments
            if ch == '(' {
                let comment_start_line = self.line;
                let comment_start_pos = self.position_on_line;
                self.advance();
                let comment = self.read_comment()?;
                let comment_end_pos = self.position_on_line;
                tokens.push(TokenWithPos {
                    token: Token::Comment(comment),
                    line: comment_start_line,
                    start_pos: comment_start_pos,
                    end_pos: comment_end_pos,
                });
                continue;
            }

            // Get next token
            match self.next_token()? {
                Token::Eof => break,
                token => {
                    let token_end_position = self.position;
                    let token_end_pos =
                        start_pos + (token_end_position - token_start_position).max(1) - 1;
                    // Debug print for position tracking
                    println!(
                        "DEBUG: token={:?} line={} start_pos={} end_pos={} position={} position_on_line={}",
                        token, start_line, start_pos, token_end_pos, self.position, self.position_on_line
                    );
                    tokens.push(TokenWithPos {
                        token,
                        line: start_line,
                        start_pos,
                        end_pos: token_end_pos,
                    });
                }
            }
        }

        Ok(tokens)
    }

    /// Move to the next character
    fn advance(&mut self) {
        if self.position < self.input.len() {
            if self.input.chars().nth(self.position) == Some('\n') {
                self.line += 1;
                self.position_on_line = 1;
            } else {
                self.position_on_line += 1;
            }
            self.position += 1;
        }
    }

    /// Skip whitespace characters except newlines
    fn skip_whitespace(&mut self) {
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch.is_whitespace() && ch != '\n' {
                self.advance();
            } else {
                break;
            }
        }
    }

    /// Read characters until the specified delimiter is found
    #[allow(dead_code)]
    fn read_until(&mut self, delimiter: char) -> Result<String> {
        let mut result = String::new();
        while self.position < self.input.len() && self.current_char() != delimiter {
            result.push(self.current_char());
            self.advance();
        }
        Ok(result)
    }

    /// Read comment with proper nested parentheses handling
    fn read_comment(&mut self) -> Result<String> {
        let mut result = String::new();
        let mut depth = 1; // We've already consumed the opening (

        while self.position < self.input.len() && depth > 0 {
            let ch = self.current_char();

            match ch {
                '(' => {
                    depth += 1;
                    result.push(ch);
                    self.advance();
                }
                ')' => {
                    depth -= 1;
                    if depth > 0 {
                        result.push(ch);
                    }
                    self.advance();
                }
                _ => {
                    result.push(ch);
                    self.advance();
                }
            }
        }

        Ok(result)
    }

    /// Read hexadecimal digits for #hex literals
    fn read_hex(&mut self) -> Result<String> {
        let mut result = String::new();
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch.is_ascii_hexdigit() {
                result.push(ch);
                self.advance();
            } else {
                break;
            }
        }
        if result.is_empty() {
            return Err(AssemblerError::SyntaxError {
                path: self.path.clone().unwrap_or_default(),
                line: self.line,
                position: self.position_on_line,
                message: "Expected hexadecimal digits".to_string(),
                source_line: self.get_current_line(),
            });
        }
        Ok(result)
    }

    /// Read binary digits for #b binary literals
    fn read_binary(&mut self) -> Result<String> {
        let mut result = String::new();
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch == '0' || ch == '1' {
                result.push(ch);
                self.advance();
            } else {
                break;
            }
        }
        if result.is_empty() {
            return Err(AssemblerError::SyntaxError {
                path: self.path.clone().unwrap_or_default(),
                line: self.line,
                position: self.position_on_line,
                message: "Expected binary digits".to_string(),
                source_line: self.get_current_line(),
            });
        }
        Ok(result)
    }

    /// Read a decimal number (digits only)
    #[allow(dead_code)]
    fn read_number(&mut self) -> Result<String> {
        let mut result = String::new();
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch.is_ascii_digit() {
                result.push(ch);
                self.advance();
            } else {
                break;
            }
        }
        Ok(result)
    }

    /// Read a hexadecimal number (hex digits only)
    /// For asset data, reads up to 4 hex digits for RawHex tokens
    fn read_hex_number(&mut self) -> Result<String> {
        let mut result = String::new();
        let mut count = 0;
        while self.position < self.input.len() && count < 4 {
            let ch = self.current_char();
            if ch.is_ascii_hexdigit() {
                result.push(ch);
                self.advance();
                count += 1;
            } else {
                break;
            }
        }
        // Fix position_on_line drift: always update position_on_line based on chars consumed
        // This is already handled by advance(), so do NOT manually update position_on_line here.
        Ok(result)
    }

    /// Read an identifier (labels, instructions, etc.)
    /// Allows alphanumeric characters, underscores, hyphens, and forward slashes
    fn read_identifier(&mut self) -> Result<String> {
        let mut result = String::new();
        let mut first = true;
        while self.position < self.input.len() {
            let ch = self.current_char();
            // Accept alphanumeric, underscore, hyphen, slash, angle brackets, allow '@', '{' , '&', and '"' as first char
            if ch.is_ascii_alphanumeric()
                || ch == '_'
                || ch == '-'
                || ch == '/'
                || ch == '<'
                || ch == '>'
                || ch == '"'
                || ch == '#'
                || (first && (ch == '@' || ch == '{' || ch == '&'))
            {
                // Special case: treat "" as a CharLiteral('"')
                if first && ch == '"' && self.input.chars().nth(self.position + 1) == Some('"') {
                    self.advance(); // skip first "
                    self.advance(); // skip second "
                                    // Return as CharLiteral token, not identifier
                    return Ok(String::from("\"")); // This will be handled as a CharLiteral by the parser
                }
                result.push(ch);
                self.advance();
                first = false;
            } else {
                break;
            }
        }
        // Debug: print what we read to see if it's being truncated
        eprintln!("DEBUG: read_identifier result: '{}'", result);
        // If we returned "" above, result will be empty, so skip error
        if result.is_empty() {
            // If we just parsed "", treat as CharLiteral('"')
            if self.position > 1
                && self.input.chars().nth(self.position - 2) == Some('"')
                && self.input.chars().nth(self.position - 1) == Some('"')
            {
                return Ok(String::from("\""));
            }
            return Err(AssemblerError::SyntaxError {
                path: self.path.clone().unwrap_or_default(),
                line: self.line,
                position: self.position_on_line,
                message: "Expected identifier".to_string(),
                source_line: self.get_current_line(),
            });
        }
        Ok(result)
    }

    /// Read a string literal, handling both quoted and unquoted TAL strings
    /// Quoted strings end at the closing quote, unquoted strings end at whitespace
    fn read_string(&mut self) -> Result<String> {
        let mut result = String::new();

        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch == '"' {
                self.advance(); // Skip closing quote
                break;
            } else if ch.is_whitespace() {
                // In TAL, unclosed strings end at whitespace
                break;
            } else {
                result.push(ch);
                self.advance();
            }
        }

        Ok(result)
    }

    /// Read device access syntax like Console/char or Screen/width
    fn read_device_access(&mut self) -> Result<String> {
        let mut result = String::new();
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' || ch == '/' {
                result.push(ch);
                self.advance();
            } else {
                break;
            }
        }
        if result.is_empty() {
            return Err(AssemblerError::SyntaxError {
                path: self.path.clone().unwrap_or_default(),
                line: self.line,
                position: self.position_on_line,
                message: "Expected device access".to_string(),
                source_line: self.get_current_line(),
            });
        }
        Ok(result)
    }

    /// Read include path after ~ token
    fn read_include_path(&mut self) -> Result<String> {
        let mut path = String::new();

        // Skip any whitespace after ~
        self.skip_whitespace();

        while self.position < self.input.len() {
            let ch = self.current_char();

            // Include path ends at whitespace or newline
            if ch.is_whitespace() {
                break;
            }

            path.push(ch);
            self.advance();
        }

        if path.is_empty() {
            return Err(AssemblerError::SyntaxError {
                path: self.path.clone().unwrap_or_default(),
                line: self.line,
                position: self.position_on_line,
                message: "Empty include path".to_string(),
                source_line: self.get_current_line(),
            });
        }

        Ok(path)
    }

    /// Read macro call name between < and >
    fn read_macro_call(&mut self) -> Result<String> {
        let mut name = String::new();

        while self.position < self.input.len() {
            let ch = self.current_char();

            if ch == '>' {
                self.advance(); // consume closing >
                break;
            }

            // Allow alphanumeric, hyphens, underscores, and slashes in macro names
            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '/' {
                name.push(ch);
                self.advance();
            } else {
                return Err(AssemblerError::SyntaxError {
                    path: self.path.clone().unwrap_or_default(),
                    line: self.line,
                    position: self.position_on_line,
                    message: format!("Invalid character in macro call: {}", ch),
                    source_line: self.get_current_line(),
                });
            }
        }

        if name.is_empty() {
            return Err(AssemblerError::SyntaxError {
                path: self.path.clone().unwrap_or_default(),
                line: self.line,
                position: self.position_on_line,
                message: "Empty macro call".to_string(),
                source_line: self.get_current_line(),
            });
        }

        Ok(name)
    }

    /// Get the current character at the position
    fn current_char(&self) -> char {
        self.input.chars().nth(self.position).unwrap_or('\0')
    }

    /// Get the next token from the input
    pub fn next_token(&mut self) -> Result<Token> {
        // Check for pending sublabel/label from conditional operator
        if let Some(label) = self.next_conditional_sublabel.take() {
            if label.starts_with('&') {
                return Ok(Token::SublabelRef(
                    label.trim_start_matches('&').to_string(),
                ));
            }
            if label.starts_with('<') {
                return Ok(Token::LabelRef(label.trim_start_matches('<').to_string()));
            }
            return Ok(Token::LabelRef(label));
        }

        // Robustly skip all whitespace except newlines before tokenizing
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch.is_whitespace() && ch != '\n' {
                self.advance();
            } else {
                break;
            }
        }

        // Skip comments before tokenizing
        while self.position < self.input.len() {
            let ch = self.current_char();
            if ch.is_whitespace() && ch != '\n' {
                self.advance();
            } else if ch == '(' {
                self.advance();
                let mut depth = 1;
                while self.position < self.input.len() && depth > 0 {
                    let ch = self.current_char();
                    if ch == '(' {
                        depth += 1;
                    } else if ch == ')' {
                        depth -= 1;
                    }
                    self.advance();
                }
                continue;
            } else {
                break;
            }
        }

        if self.position >= self.input.len() {
            return Ok(Token::Eof);
        }

        let ch = self.current_char();

        match ch {
            '\n' => {
                self.advance();
                Ok(Token::Newline)
            }
            '(' => {
                self.advance();
                let comment = self.read_comment()?;
                Ok(Token::Comment(comment))
            }
            '"' => {
                self.advance();
                if self.current_char() == '"' {
                    self.advance();
                    return Ok(Token::RawString(String::new()));
                }
                let ch = self.current_char();
                if ch != '\0' && !ch.is_whitespace() {
                    let next_pos = self.position + 1;
                    let next_ch = if next_pos < self.input.len() {
                        self.input.chars().nth(next_pos).unwrap_or('\0')
                    } else {
                        '\0'
                    };
                    if next_ch.is_whitespace()
                        || next_ch == '\0'
                        || next_ch == ']'
                        || next_ch == ')'
                    {
                        self.advance();
                        Ok(Token::CharLiteral(ch))
                    } else {
                        let string = self.read_string()?;
                        Ok(Token::RawString(string))
                    }
                } else {
                    let string = self.read_string()?;
                    Ok(Token::RawString(string))
                }
            }
            '\'' => {
                self.advance();
                let ch = self.current_char();
                if ch == '\0' {
                    return Err(AssemblerError::SyntaxError {
                        path: self.path.clone().unwrap_or_default(),
                        line: self.line,
                        position: self.position_on_line,
                        message: "Unexpected end of file in character literal".to_string(),
                        source_line: self.get_current_line(),
                    });
                }
                self.advance();
                Ok(Token::CharLiteral(ch))
            }
            '#' => {
                self.advance();
                loop {
                    let ch = self.current_char();
                    if ch.is_whitespace() && ch != '\n' {
                        self.advance();
                    } else if ch == '(' {
                        self.advance();
                        let comment = self.read_comment()?;
                        let comment_newlines = comment.chars().filter(|c| *c == '\n').count();
                        self.line += comment_newlines;
                    } else {
                        break;
                    }
                }
                if self.current_char() == '"' {
                    self.advance();
                    let ch = self.current_char();
                    if ch != '\0' {
                        self.advance();
                        return Ok(Token::CharLiteral(ch));
                    }
                }
                if self.current_char() == '\'' {
                    self.advance();
                    let ch = self.current_char();
                    if ch != '\0' {
                        self.advance();
                        return Ok(Token::CharLiteral(ch));
                    }
                }
                if self.current_char() == '[' {
                    self.advance();
                    let ch = self.current_char();
                    if ch != '\0' && self.input.chars().nth(self.position + 1) == Some(']') {
                        self.advance();
                        self.advance();
                        return Ok(Token::CharLiteral(ch));
                    }
                }
                if self.current_char().is_ascii_alphabetic()
                    && self
                        .input
                        .chars()
                        .nth(self.position + 1)
                        .map(|c| c.is_whitespace() || c == ']')
                        == Some(true)
                {
                    let ch = self.current_char();
                    self.advance();
                    return Ok(Token::CharLiteral(ch));
                }
                if !self.current_char().is_ascii_hexdigit() {
                    return Err(AssemblerError::SyntaxError {
                        path: self.path.clone().unwrap_or_default(),
                        line: self.line,
                        position: self.position_on_line,
                        message: "Expected hexadecimal digits after '#'".to_string(),
                        source_line: self.get_current_line(),
                    });
                }
                let hex = self.read_hex()?;
                Ok(Token::HexLiteral(hex))
            }
            '@' => {
                self.advance();
                let label = self.read_identifier()?;
                Ok(Token::LabelDef(label))
            }
            ';' => {
                self.advance();
                let label = self.read_identifier()?;
                Ok(Token::SemicolonRef(label))
            }
            '.' => {
                self.advance();
                let ident = self.read_identifier()?;
                Ok(Token::DotRef(ident))
            }
            '=' => {
                self.advance();
                let label = self.read_identifier()?;
                Ok(Token::EqualsRef(label))
            }
            ',' => {
                self.advance();
                let label = self.read_identifier()?;
                Ok(Token::CommaRef(label))
            }
            '_' => {
                self.advance();
                let label = self.read_identifier()?;
                Ok(Token::UnderscoreRef(label))
            }
            '-' => {
                self.advance();
                let identifier = self.read_identifier()?;
                Ok(Token::HyphenRef(identifier))
            }
            '/' => {
                self.advance();
                let label = self.read_identifier()?;
                // Don't include the leading slash in the label name - it's just syntax
                Ok(Token::RelativeRef(label))
            }
            '?' => {
                self.advance();
                if self.current_char() == '{' {
                    return Ok(Token::ConditionalOperator);
                }
                let label = self.read_identifier()?;
                Ok(Token::QuestionRef(label))
            }
            '!' => {
                self.advance();
                let label = self.read_identifier()?;
                Ok(Token::ExclamationRef(label))
            }
            '{' => {
                self.advance();
                Ok(Token::BraceOpen)
            }
            '}' => {
                self.advance();
                Ok(Token::BraceClose)
            }
            '[' => {
                self.advance();
                if self.current_char() == '"'
                    && self.input.chars().nth(self.position + 1) == Some('"')
                {
                    self.advance();
                    self.advance();
                    return Ok(Token::CharLiteral('"'));
                }
                Ok(Token::BracketOpen)
            }
            ']' => {
                self.advance();
                Ok(Token::BracketClose)
            }
            '~' => {
                self.advance();
                let filename = self.read_include_path()?;
                Ok(Token::Include(filename))
            }
            '<' => {
                self.advance();
                let macro_name = self.read_macro_call()?;
                Ok(Token::MacroCall(macro_name))
            }
            _ if ch.is_ascii_digit() => {
                let number = self.read_hex_number()?;
                Ok(Token::RawHex(number))
            }
            _ if ch.is_ascii_alphabetic() || ch == '_' => {
                let identifier = self.read_identifier()?;
                // Only treat as hex if it's a simple pattern like "ff", "1234", etc.
                // Don't treat words with hyphens or complex patterns as hex
                // Also don't treat known instruction names as hex
                if identifier.len() >= 2
                    && identifier.len() <= 4
                    && identifier.len() % 2 == 0
                    && identifier.chars().all(|c| c.is_ascii_hexdigit())
                    && !identifier.contains('-')  // No hyphens in hex values
                    && !identifier.contains('/')  // No slashes in hex values
                    && !identifier.contains('_')  // No underscores in hex values
                    && !is_instruction_name(&identifier)
                // Don't treat instruction names as hex
                {
                    Ok(Token::RawHex(identifier))
                } else {
                    Ok(Token::Instruction(identifier))
                }
            }
            '|' => {
                self.advance();
                let padding_value = self.read_identifier()?;

                // Check if it's a hex value or a label
                if padding_value.chars().all(|c| c.is_ascii_hexdigit()) {
                    // Parse as hex address
                    let addr = u16::from_str_radix(&padding_value, 16)
                        .map_err(|_| self.syntax_error("Invalid padding address".to_string()))?;
                    Ok(Token::Padding(addr))
                } else {
                    // It's a label reference for padding
                    Ok(Token::PaddingLabel(padding_value))
                }
            }
            '$' => {
                self.advance();
                let skip_value = self.read_identifier()?;
                // Parse as hex value for skip count
                if skip_value.chars().all(|c| c.is_ascii_hexdigit()) {
                    let count = u16::from_str_radix(&skip_value, 16)
                        .map_err(|_| self.syntax_error("Invalid skip count".to_string()))?;
                    Ok(Token::Skip(count))
                } else {
                    // Could be a label-based skip, but for now treat as error
                    Err(self.syntax_error("Skip directive requires hex value".to_string()))
                }
            }
            '%' => {
                self.advance();
                let macro_name = self.read_identifier()?;
                Ok(Token::MacroDef(macro_name))
            }
            '&' => {
                self.advance();
                let sublabel = self.read_identifier()?;
                // In uxnasm.c, '&' always creates a sublabel definition (makelabel)
                // Comma references use ',' prefix, not '&'
                Ok(Token::SublabelDef(sublabel))
            }
            _ => Err(self.syntax_error(format!("Unexpected character: '{}'", ch))),
        }
    }
}

/// Check if an identifier is a known instruction name
fn is_instruction_name(identifier: &str) -> bool {
    // List of UXN instruction names that could be confused with hex
    const INSTRUCTIONS: &[&str] = &[
        "ADD", "ADD2", "SUB", "SUB2", "MUL", "MUL2", "DIV", "DIV2", "AND", "AND2", "ORA", "ORA2",
        "EOR", "EOR2", "SFT", "SFT2", "LDZ", "LDZ2", "STZ", "STZ2", "LDR", "LDR2", "STR", "STR2",
        "LDA", "LDA2", "STA", "STA2", "DEI", "DEI2", "DEO", "DEO2", "INC", "INC2", "POP", "POP2",
        "NIP", "NIP2", "SWP", "SWP2", "ROT", "ROT2", "DUP", "DUP2", "OVR", "OVR2", "EQU", "EQU2",
        "NEQ", "NEQ2", "GTH", "GTH2", "LTH", "LTH2", "JMP", "JMP2", "JCN", "JCN2", "JSR", "JSR2",
        "STH", "STH2", "BRK", "LIT", "LIT2",
    ];

    // Check the base instruction name (without mode flags)
    let base_name = if identifier.len() > 3 {
        let mut base = identifier.to_string();
        // Remove mode flags (k, r, 2)
        while base.ends_with('k') || base.ends_with('r') || base.ends_with('2') {
            base.pop();
        }
        base
    } else {
        identifier.to_string()
    };

    INSTRUCTIONS.iter().any(|&inst| {
        inst == identifier
            || inst == base_name
            || (inst.ends_with('2') && inst[..inst.len() - 1] == base_name)
    })
}