from enum import Enum
from typing import Any


class TokenType(Enum):
	# Special tokens
	EOF = "EOF"
	ILLEGAL = "ILLEGAL"

	# Data types
	IDENT = "IDENT"
	INT = "INT"
	FLOAT = "FLOAT"
	STRING = "STRING"

	# Arithmetic symbols
	PLUS = "PLUS"
	MINUS = "MINUS"
	ASTERISK = "ASTERISK"
	SLASH = "SLASH"
	POW = "POW"
	MODULUS = "MODULUS"

	# Assignment symbols
	EQ = "EQ"

	# Comparison symbols
	LT = "<"
	GT = ">"
	EQ_EQ = "=="
	LT_EQ = "<="
	GT_EQ = ">="
	NOT_EQ = "!="

	# Symbols
	LPAREN = "LPAREN"
	RPAREN = "RPAREN"
	LBRACKET = "LBRACKET"
	RBRACKET = "RBRACKET"
	LBRACE = "LBRACE"
	RBRACE = "RBRACE"
	COLON = "COLON"
	SEMICOLON = "SEMICOLON"
	COMMA = "COMMA"
	DOLLARSIGN = "DOLLARSIGN"

	# Keywords
	RETURN = "RETURN"
	IF = "IF"
	UNLESS = "UNLESS"
	TRUE = "TRUE"
	FALSE = "FALSE"
	WHILE = "WHILE"

	# Typing
	TYPE = "TYPE"

class Token:
	def __init__(self, type: TokenType, literal: Any, line_no: int, position: int) -> None:
		self.type = type
		self.literal = literal
		self.line_no = line_no
		self.position = position

	def __str__(self) -> str:
		return f"token[{self.type} : {self.literal} : Line {self.line_no} : Position {self.position}]"
	
	def __repr__(self) -> str:
		return str(self)
	

KEYWORDS: dict[str, TokenType] = {
	"return": TokenType.RETURN,
	"if": TokenType.IF,
	"unless": TokenType.UNLESS,
	"true": TokenType.TRUE,
	"false": TokenType.FALSE,
	"while": TokenType.WHILE
}

ALT_KEYWORDS: dict[str, TokenType] = {
	"else": TokenType.UNLESS,
	"ret": TokenType.RETURN
}

TYPE_KEYWORDS: list[str] = ["Int", "Float", "String", "Bool", "List", "Nil", "Func"]

def lookup_ident(ident: str) -> TokenType:
	tt: TokenType | None = KEYWORDS.get(ident)
	if tt is not None:
		return tt
	
	tt: TokenType | None = ALT_KEYWORDS.get(ident)
	if tt is not None:
		return tt
	
	if ident in TYPE_KEYWORDS:
		return TokenType.TYPE
	
	return TokenType.IDENT