"""
This module provides the lexer which is used to parse templates. The members of this module are not required in most templating workflows.
"""

from dataclasses import dataclass
from enum import Enum
import re
import string
from typing import Optional


class TokenType(Enum):
    """
    An enumeration of all the token types a scanner can emit.
    """

    # Single character tokens.
    LEFT_DOUBLE_CURLY_PAREN = 1
    RIGHT_DOUBLE_CURLY_PAREN = 2
    PLUS = 3
    MINUS = 4
    MULTIPLY = 5
    DIVIDE = 6
    LEFT_PAREN = 7
    RIGHT_PAREN = 8
    DOT = 9
    COMMA = 10

    # Literals.
    STRING = 11
    NUMBER = 12

    # Keywords.
    FOREACH = 13
    ENDFOREACH = 14
    IN = 15

    # Line delimiters.
    NEWLINE = 16
    EOF = 17


@dataclass
class Token:
    """
    A token generated by a scanner. Has an optional literal value depending on the token's type
    for example, a token of TokenType.STRING will have the string's literal value in the literal
    field.
    """

    type: TokenType
    lexeme: str
    literal: Optional[str | int | float]
    line_number: int


allowed_whitespace = " \t\v\f"
printable_characters = list(
    string.digits + string.ascii_letters + string.punctuation + string.whitespace
)
string_chars_without_nl = [c for c in printable_characters if c not in ["\n"]]
identifer_chars = [
    c
    for c in list(string.digits + string.ascii_letters + string.punctuation)
    if c not in ["(", ")", ","]
]


class Scanner:
    """
    Scanner class for tokenising documents.
    """

    def __init__(self, source):
        self.source = self._normalise_newlines(source)
        self.current = 0
        self.line_number = 1

        self.break_on_whitespace = False

    def _normalise_newlines(self, source: str):
        return source.replace("\r\n", "\n")

    def _add_token(self, type: TokenType, literal=None):
        return Token(type, None, literal, self.line_number)

    def _char_in_list(self, character: str, symbols: list[str]):
        if character in symbols:
            return character

    def _add_literal(self, value: str):
        if not self.break_on_whitespace:
            return self._add_token(TokenType.STRING, value)

        if value == "+":
            return self._add_token(TokenType.PLUS)
        elif value == "-":
            return self._add_token(TokenType.MINUS)
        elif value == "*":
            return self._add_token(TokenType.MULTIPLY)
        elif value == "/":
            return self._add_token(TokenType.DIVIDE)
        elif value == "(":
            return self._add_token(TokenType.LEFT_PAREN)
        elif value == ")":
            return self._add_token(TokenType.RIGHT_PAREN)
        elif value == ".":
            return self._add_token(TokenType.DOT)
        elif re.match(r"^-?\d+\.\d+$", value):
            return self._add_token(TokenType.NUMBER, float(value))
        elif value.isdigit():
            return self._add_token(TokenType.NUMBER, int(value))
        elif value == "foreach":
            return self._add_token(TokenType.FOREACH)
        elif value == "endforeach":
            return self._add_token(TokenType.ENDFOREACH)
        elif value == "in":
            return self._add_token(TokenType.IN)
        else:
            return self._add_token(TokenType.STRING, value)

    def trim_whitespace(self):
        """
        Sets the scanner to break on all whitespace.
        """
        self.break_on_whitespace = True

    def keep_whitespace(self):
        """
        Sets the scanner to parse all whitespace (except newlines) as a string.
        """
        self.break_on_whitespace = False

    def scan(self):
        """
        Scans a document, yielding tokens.
        """
        while self.current < len(self.source):
            c = self.source[self.current]
            if self.break_on_whitespace and (c in allowed_whitespace):
                self.current += 1
            elif c == "\n":
                self.current += 1
                yield self._add_token(TokenType.NEWLINE)
                self.line_number += 1
            elif c == "{" and self.source[self.current + 1] == "{":
                self.current += 2
                yield self._add_token(TokenType.LEFT_DOUBLE_CURLY_PAREN)
            elif c == "}" and self.source[self.current + 1] == "}":
                self.current += 2
                yield self._add_token(TokenType.RIGHT_DOUBLE_CURLY_PAREN)
            elif self.break_on_whitespace and c == "(":
                self.current += 1
                yield self._add_token(TokenType.LEFT_PAREN)
            elif self.break_on_whitespace and c == ")":
                self.current += 1
                yield self._add_token(TokenType.RIGHT_PAREN)
            elif self.break_on_whitespace and c == ",":
                self.current += 1
                yield self._add_token(TokenType.COMMA)
            elif c in string.printable:
                if self.break_on_whitespace:
                    # break on whitespace and emit string
                    allowed_chars = identifer_chars
                else:
                    # eat all the characters in printable stop on newline
                    allowed_chars = string_chars_without_nl

                literal_value = c
                self.current += 1
                while (self.current < len(self.source)) and (
                    s := self._char_in_list(self.source[self.current], allowed_chars)
                ):
                    self.current += 1

                    if s == "{" and self.source[self.current] == "{":
                        yield self._add_literal(literal_value)
                        self.current += 1
                        yield self._add_token(TokenType.LEFT_DOUBLE_CURLY_PAREN)
                        break
                    elif s == "}" and self.source[self.current] == "}":
                        self.current += 1
                        yield self._add_literal(literal_value)
                        yield self._add_token(TokenType.RIGHT_DOUBLE_CURLY_PAREN)
                        break
                    else:
                        literal_value += s
                else:
                    yield self._add_literal(literal_value)
            else:
                raise ValueError(f"Unknown token: {c} at {self.line_number}.")

        yield self._add_token(TokenType.EOF)
