"""Lexer for MicroBasic source (.mb files).

Handles:
  - comments: ' to end of line
  - numbers: decimal (120), hex (0xA1), binary (0b101)
  - strings: "..." with \\-escapes and \\xHH hex escapes
  - identifiers / keywords (case-insensitive)
  - labels: identifier at start of line followed by ':'
  - operators, per tokens.MULTI_CHAR_OPERATORS / SINGLE_CHAR_OPERATORS
  - significant newlines (MicroBasic statements are one-per-line)

This does not yet resolve #define macro substitution -- that belongs in a
preprocessing pass ahead of the parser. Left as a TODO.
"""

from __future__ import annotations

from .tokens import (
    KEYWORDS,
    MULTI_CHAR_OPERATORS,
    SINGLE_CHAR_OPERATORS,
    Token,
    TokenType,
)


class LexError(Exception):
    def __init__(self, message: str, line: int, col: int):
        super().__init__(f"{message} (line {line}, col {col})")
        self.line = line
        self.col = col


_SIMPLE_ESCAPES = {
    "'": "'",
    '"': '"',
    "\\": "\0",  # per manual: \\ -> Null
    "a": "\a",
    "b": "\b",
    "f": "\f",
    "n": "\n",
    "r": "\r",
    "t": "\t",
    "v": "\v",
}


class Lexer:
    def __init__(self, source: str):
        self.src = source
        self.pos = 0
        self.line = 1
        self.col = 1
        self._at_line_start = True

    def tokenize(self) -> list[Token]:
        tokens: list[Token] = []
        while True:
            tok = self._next_token()
            tokens.append(tok)
            if tok.type == TokenType.EOF:
                break
        return tokens

    # -- internals ---------------------------------------------------

    def _peek(self, offset: int = 0) -> str:
        i = self.pos + offset
        return self.src[i] if i < len(self.src) else ""

    def _advance(self) -> str:
        ch = self.src[self.pos]
        self.pos += 1
        if ch == "\n":
            self.line += 1
            self.col = 1
        else:
            self.col += 1
        return ch

    def _next_token(self) -> Token:
        while True:
            # Skip spaces/tabs (not newlines -- those are tokens).
            while self._peek() in (" ", "\t", "\r"):
                self._advance()

            if self._peek() == "'":
                while self._peek() and self._peek() != "\n":
                    self._advance()
                continue  # loop back around to handle newline / EOF

            break

        if not self._peek():
            return Token(TokenType.EOF, "", self.line, self.col)

        line, col = self.line, self.col
        ch = self._peek()

        if ch == "\n":
            self._advance()
            self._at_line_start = True
            return Token(TokenType.NEWLINE, "\n", line, col)

        at_line_start = self._at_line_start
        self._at_line_start = False

        if ch.isdigit():
            return self._read_number(line, col)

        if ch == '"':
            return self._read_string(line, col)

        if ch.isalpha() or ch == "_":
            return self._read_identifier(line, col, at_line_start)

        if ch == "#":
            return self._read_identifier(line, col, at_line_start)

        if ch == "(":
            self._advance()
            return Token(TokenType.LPAREN, "(", line, col)
        if ch == ")":
            self._advance()
            return Token(TokenType.RPAREN, ")", line, col)
        if ch == "[":
            self._advance()
            return Token(TokenType.LBRACKET, "[", line, col)
        if ch == "]":
            self._advance()
            return Token(TokenType.RBRACKET, "]", line, col)
        if ch == ",":
            self._advance()
            return Token(TokenType.COMMA, ",", line, col)
        if ch == ":":
            self._advance()
            return Token(TokenType.COLON, ":", line, col)

        for op in MULTI_CHAR_OPERATORS:
            if self.src.startswith(op, self.pos):
                for _ in op:
                    self._advance()
                return Token(TokenType.OP, op, line, col)

        if ch in SINGLE_CHAR_OPERATORS:
            self._advance()
            return Token(TokenType.OP, ch, line, col)

        raise LexError(f"Unexpected character {ch!r}", line, col)

    def _read_number(self, line: int, col: int) -> Token:
        start = self.pos
        if self._peek() == "0" and self._peek(1) in ("x", "X"):
            self._advance()
            self._advance()
            while self._peek() and (self._peek() in "0123456789abcdefABCDEF"):
                self._advance()
        elif self._peek() == "0" and self._peek(1) in ("b", "B"):
            self._advance()
            self._advance()
            while self._peek() in ("0", "1"):
                self._advance()
        else:
            while self._peek().isdigit():
                self._advance()
        text = self.src[start:self.pos]
        return Token(TokenType.INTEGER, text, line, col)

    def _read_string(self, line: int, col: int) -> Token:
        self._advance()  # opening quote
        chars: list[str] = []
        while True:
            ch = self._peek()
            if not ch or ch == "\n":
                raise LexError("Unterminated string literal", line, col)
            if ch == '"':
                self._advance()
                break
            if ch == "\\":
                self._advance()
                esc = self._peek()
                if esc == "x":
                    self._advance()
                    hex_digits = ""
                    while self._peek() and self._peek() in "0123456789abcdefABCDEF":
                        hex_digits += self._advance()
                    if not hex_digits:
                        raise LexError("Empty hex escape", self.line, self.col)
                    chars.append(chr(int(hex_digits, 16) & 0xFF))
                elif esc in _SIMPLE_ESCAPES:
                    self._advance()
                    chars.append(_SIMPLE_ESCAPES[esc])
                else:
                    raise LexError(f"Unknown escape sequence '\\{esc}'", self.line, self.col)
            else:
                chars.append(self._advance())
        return Token(TokenType.STRING, "".join(chars), line, col)

    def _read_identifier(self, line: int, col: int, at_line_start: bool) -> Token:
        start = self.pos
        if self._peek() == "#":
            self._advance()
        while self._peek() and (self._peek().isalnum() or self._peek() == "_"):
            self._advance()
        text = self.src[start:self.pos]

        # Label: identifier at start of line immediately followed by ':'
        if at_line_start and self._peek() == ":":
            self._advance()
            return Token(TokenType.LABEL, text, line, col)

        if text.lower() in KEYWORDS:
            return Token(TokenType.KEYWORD, text.lower(), line, col)

        return Token(TokenType.IDENT, text, line, col)


def tokenize(source: str) -> list[Token]:
    return Lexer(source).tokenize()