from .tokens import KEYWORDS, Token, TokenType

OPERATORS = {"+", "-", "*", "/", "%", "=", "==", "!=", ">", "<", ">=", "<=", "!"}


class Lexer:

    def __init__(self, source):
        self.source = source
        self.position = 0
        self.tokens = []
        self.line = 1
        self.column = 1

    def peek(self):
        if self.position < len(self.source):
            return self.source[self.position]
        return None

    def advance(self):
        if self.peek() == "\n":
            self.line += 1
            self.column = 1
        else:
            self.column += 1
        self.position += 1

    def tokenize(self):

        while self.position < len(self.source):

            char = self.peek()

            # skip whitespace
            if char.isspace():
                self.advance()
                continue

            # skip comments
            if char == "#":
                while self.peek() and self.peek() != "\n":
                    self.advance()
                continue
            if char == "/" and self.position + 1 < len(self.source) and self.source[self.position + 1] == "/":
                while self.peek() and self.peek() != "\n":
                    self.advance()
                continue

            # numbers
            if char.isdigit():
                self.tokens.append(self.read_number())
                continue

            # identifiers / keywords
            if char.isalpha() or char == "_":
                self.tokens.append(self.read_identifier())
                continue

            # strings
            if char == '"' or char == "'":
                self.tokens.append(self.read_string())
                continue

            # operators
            if char in "+-*/%=<>!":
                self.tokens.append(self.read_operator())
                continue

            # punctuation
            if char == "(":
                self.tokens.append(Token(TokenType.LPAREN, "(", self.line, self.column))
                self.advance()
                continue

            if char == ")":
                self.tokens.append(Token(TokenType.RPAREN, ")", self.line, self.column))
                self.advance()
                continue

            if char == "{":
                self.tokens.append(Token(TokenType.LBRACE, "{", self.line, self.column))
                self.advance()
                continue

            if char == "}":
                self.tokens.append(Token(TokenType.RBRACE, "}", self.line, self.column))
                self.advance()
                continue

            if char == "[":
                self.tokens.append(
                    Token(TokenType.LBRACKET, "[", self.line, self.column)
                )
                self.advance()
                continue

            if char == "]":
                self.tokens.append(
                    Token(TokenType.RBRACKET, "]", self.line, self.column)
                )
                self.advance()
                continue

            if char == ",":
                self.tokens.append(Token(TokenType.COMMA, ",", self.line, self.column))
                self.advance()
                continue

            if char == ";":
                self.tokens.append(
                    Token(TokenType.SEMICOLON, ";", self.line, self.column)
                )
                self.advance()
                continue

            if char == ":":
                self.tokens.append(Token(TokenType.COLON, ":", self.line, self.column))
                self.advance()
                continue

            if char == ".":
                self.tokens.append(Token(TokenType.DOT, ".", self.line, self.column))
                self.advance()
                continue

            raise Exception(f"Unexpected character: {char}")

        self.tokens.append(Token(TokenType.EOF, None, self.line, self.column))
        return self.tokens

    def read_number(self):
        line, col = self.line, self.column
        start = self.position

        while self.peek() and self.peek().isdigit():
            self.advance()

        is_float = False

        if self.peek() == "." and (self.position + 1) < len(self.source):
            next_char = self.source[self.position + 1]
            if next_char.isdigit():
                is_float = True
                self.advance()  # consume '.'
                while self.peek() and self.peek().isdigit():
                    self.advance()

        number_text = self.source[start : self.position]
        value = float(number_text) if is_float else int(number_text)

        return Token(TokenType.NUMBER, value, line, col)

    def read_identifier(self):
        line, col = self.line, self.column
        start = self.position

        while self.peek() and (self.peek().isalnum() or self.peek() == "_"):
            self.advance()

        value = self.source[start : self.position]

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

        return Token(TokenType.IDENTIFIER, value, line, col)

    def read_string(self):
        line, col = self.line, self.column
        quote = self.peek()
        self.advance()

        chars = []

        while True:
            ch = self.peek()
            if ch is None:
                raise Exception("Unterminated string")
            if ch == quote:
                break

            if ch == "\\":
                self.advance()  # consume backslash
                esc = self.peek()
                if esc is None:
                    raise Exception("Unterminated escape sequence in string")

                if esc == "n":
                    chars.append("\n")
                elif esc == "t":
                    chars.append("\t")
                elif esc == "\\":
                    chars.append("\\")
                elif esc == '"':
                    chars.append('"')
                elif esc == "'":
                    chars.append("'")
                else:
                    raise Exception(f"Unknown escape sequence: \\{esc}")

                self.advance()  # consume escape char
                continue

            chars.append(ch)
            self.advance()

        self.advance()  # consume closing quote

        return Token(TokenType.STRING, "".join(chars), line, col)

    def read_operator(self):
        line, col = self.line, self.column
        first = self.peek()
        second = None
        if self.position + 1 < len(self.source):
            second = self.source[self.position + 1]

        if second is not None and (first + second) in OPERATORS:
            value = first + second
            self.advance()
            self.advance()
            return Token(TokenType.OPERATOR, value, line, col)

        if first in OPERATORS:
            value = first
            self.advance()
            return Token(TokenType.OPERATOR, value, line, col)

        raise Exception(f"Unknown operator: {first}")
