"""
TENGWAR Lexer

Tokenizes TENGWAR source code. Handles both Unicode operators (primary form)
and ASCII fallbacks for environments with limited Unicode support.

Token efficiency is paramount: every token carries maximum semantic weight.
"""
from enum import Enum, auto
from dataclasses import dataclass
from typing import List
from .errors import LexError


class TokenType(Enum):
    # Delimiters
    LPAREN = auto()      # (
    RPAREN = auto()      # )
    LANGLE = auto()      # ⟨ or <|
    RANGLE = auto()      # ⟩ or |>
    LBRACKET = auto()    # ⟦ or [|
    RBRACKET = auto()    # ⟧ or |]
    LBRACE = auto()      # {
    RBRACE = auto()      # }

    # Keywords (single Unicode chars)
    LAMBDA = auto()      # λ or \
    ARROW = auto()       # → or ->
    COND = auto()        # ?
    MATCH = auto()       # ~
    SEQ = auto()         # >>
    PARALLEL = auto()    # ∥ or ||
    MODULE = auto()      # □ or []
    DEFINE = auto()      # :=
    RECURSE = auto()     # ↺ or @rec
    TYPE = auto()        # τ or :t
    PROOF = auto()       # ⊢ or |-
    EFFECT = auto()      # ⊕ or +>
    MUTATE = auto()      # μ or !
    UNIT = auto()        # ∅ or ()
    IMPORT = auto()      # ⇐ or <=

    # Operators
    PLUS = auto()        # +
    MINUS = auto()       # -
    STAR = auto()        # *
    SLASH = auto()       # /
    PERCENT = auto()     # %
    EQ = auto()          # =
    NEQ = auto()         # !=
    LT = auto()          # <
    GT = auto()          # >
    LTE = auto()         # <=
    GTE = auto()         # >=
    AND = auto()         # &
    OR = auto()          # |
    NOT = auto()         # !
    DOT = auto()         # .
    COLON = auto()       # :

    # Literals
    INT = auto()
    FLOAT = auto()
    STRING = auto()
    TEMPLATE = auto()     # $"...{expr}..." template string
    TRUE = auto()
    FALSE = auto()

    # Identifiers
    SYMBOL = auto()      # regular name
    HASH_ID = auto()     # #abc
    ADDR_REF = auto()    # @abc

    # Special
    COMMENT = auto()     # ;;
    UNDERSCORE = auto()  # _ (wildcard)
    EOF = auto()


@dataclass
class Token:
    type: TokenType
    value: str
    line: int
    col: int

    def __repr__(self):
        return f"Token({self.type.name}, {self.value!r}, {self.line}:{self.col})"


# Unicode → ASCII fallback mapping
UNICODE_MAP = {
    'λ': TokenType.LAMBDA,
    '≡': TokenType.DEFINE,
    '→': TokenType.ARROW,
    '↺': TokenType.RECURSE,
    'τ': TokenType.TYPE,
    '⊢': TokenType.PROOF,
    '⊕': TokenType.EFFECT,
    'μ': TokenType.MUTATE,
    '∅': TokenType.UNIT,
    '∥': TokenType.PARALLEL,
    '□': TokenType.MODULE,
    '⇐': TokenType.IMPORT,
    '⟨': TokenType.LANGLE,
    '⟩': TokenType.RANGLE,
    '⟦': TokenType.LBRACKET,
    '⟧': TokenType.RBRACKET,
}

KEYWORDS = {
    'true': TokenType.TRUE,
    'false': TokenType.FALSE,
    # ASCII aliases for Unicode operators
    'fn': TokenType.LAMBDA,
    'def': TokenType.DEFINE,
    'if': TokenType.COND,
    'match': TokenType.MATCH,
    'rec': TokenType.RECURSE,
    'do': TokenType.SEQ,
    'par': TokenType.PARALLEL,
    'nil': TokenType.UNIT,
    'mod': TokenType.MODULE,
    # Note: let, pipe, throw, catch, try are handled as symbols
    # by the parser since they need custom parse rules
}

SINGLE_CHAR = {
    '(': TokenType.LPAREN,
    ')': TokenType.RPAREN,
    '{': TokenType.LBRACE,
    '}': TokenType.RBRACE,
    '?': TokenType.COND,
    '~': TokenType.MATCH,
    '+': TokenType.PLUS,
    '*': TokenType.STAR,
    '/': TokenType.SLASH,
    '%': TokenType.PERCENT,
    '=': TokenType.EQ,
    '&': TokenType.AND,
    '.': TokenType.DOT,
    ':': TokenType.COLON,
    '_': TokenType.UNDERSCORE,
}


class Lexer:
    def __init__(self, source: str):
        self.source = source
        self.pos = 0
        self.line = 1
        self.col = 1
        self.tokens: List[Token] = []

    def peek(self) -> str:
        if self.pos >= len(self.source):
            return '\0'
        return self.source[self.pos]

    def peek_ahead(self, n: int = 1) -> str:
        pos = self.pos + n
        if pos >= len(self.source):
            return '\0'
        return self.source[pos]

    def advance(self) -> str:
        ch = self.source[self.pos]
        self.pos += 1
        if ch == '\n':
            self.line += 1
            self.col = 1
        else:
            self.col += 1
        return ch

    def add_token(self, type: TokenType, value: str, line: int, col: int):
        self.tokens.append(Token(type, value, line, col))

    def skip_whitespace(self):
        while self.pos < len(self.source) and self.source[self.pos] in ' \t\n\r':
            self.advance()

    def read_string(self) -> str:
        """Read a string literal (double-quoted)"""
        result = []
        self.advance()  # skip opening "
        while self.pos < len(self.source):
            ch = self.advance()
            if ch == '"':
                return ''.join(result)
            elif ch == '\\':
                if self.pos < len(self.source):
                    esc = self.advance()
                    escape_map = {'n': '\n', 't': '\t', 'r': '\r', '\\': '\\', '"': '"'}
                    result.append(escape_map.get(esc, esc))
            else:
                result.append(ch)
        raise LexError("Unterminated string literal", self.line, self.col)

    def read_number(self) -> Token:
        """Read integer or float literal"""
        start_line, start_col = self.line, self.col
        result = []
        is_float = False
        is_neg = False

        if self.peek() == '-':
            result.append(self.advance())
            is_neg = True

        while self.pos < len(self.source) and (self.source[self.pos].isdigit() or self.source[self.pos] == '.'):
            ch = self.advance()
            if ch == '.':
                if is_float:
                    raise LexError("Multiple decimal points in number", self.line, self.col)
                is_float = True
            result.append(ch)

        # Handle scientific notation
        if self.pos < len(self.source) and self.source[self.pos] in ('e', 'E'):
            is_float = True
            result.append(self.advance())
            if self.pos < len(self.source) and self.source[self.pos] in ('+', '-'):
                result.append(self.advance())
            while self.pos < len(self.source) and self.source[self.pos].isdigit():
                result.append(self.advance())

        value = ''.join(result)
        if is_float:
            return Token(TokenType.FLOAT, value, start_line, start_col)
        else:
            return Token(TokenType.INT, value, start_line, start_col)

    def read_symbol(self) -> str:
        """Read a symbol/identifier"""
        result = []
        valid_chars = set('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_.-')
        while self.pos < len(self.source) and self.source[self.pos] in valid_chars:
            result.append(self.advance())
        # Allow trailing ? or ! for predicates and mutators (e.g., empty?, set!)
        if self.pos < len(self.source) and self.source[self.pos] in ('?', '!'):
            result.append(self.advance())
        return ''.join(result)

    def read_comment(self) -> str:
        """Read comment until end of line"""
        result = []
        while self.pos < len(self.source) and self.source[self.pos] != '\n':
            result.append(self.advance())
        return ''.join(result).strip()

    def tokenize(self) -> List[Token]:
        """Tokenize the full source into a list of tokens"""
        while self.pos < len(self.source):
            self.skip_whitespace()
            if self.pos >= len(self.source):
                break

            ch = self.source[self.pos]
            start_line, start_col = self.line, self.col

            # Unicode operators (primary form)
            if ch in UNICODE_MAP:
                self.advance()
                self.add_token(UNICODE_MAP[ch], ch, start_line, start_col)
                continue

            # Comments: ;;
            if ch == ';' and self.peek_ahead() == ';':
                self.advance()
                self.advance()
                comment = self.read_comment()
                self.add_token(TokenType.COMMENT, comment, start_line, start_col)
                continue

            # Single ; is also a comment
            if ch == ';':
                self.advance()
                comment = self.read_comment()
                self.add_token(TokenType.COMMENT, comment, start_line, start_col)
                continue

            # Two-character operators (ASCII fallbacks and combos)
            if self.pos + 1 < len(self.source):
                two = self.source[self.pos:self.pos+2]
                if two == '->':
                    self.advance(); self.advance()
                    self.add_token(TokenType.ARROW, '→', start_line, start_col)
                    continue
                elif two == '>>':
                    self.advance(); self.advance()
                    self.add_token(TokenType.SEQ, '>>', start_line, start_col)
                    continue
                elif two == ':=':
                    self.advance(); self.advance()
                    self.add_token(TokenType.DEFINE, ':=', start_line, start_col)
                    continue
                elif two == '||':
                    self.advance(); self.advance()
                    self.add_token(TokenType.PARALLEL, '∥', start_line, start_col)
                    continue
                elif two == '|-':
                    self.advance(); self.advance()
                    self.add_token(TokenType.PROOF, '⊢', start_line, start_col)
                    continue
                elif two == '+>':
                    self.advance(); self.advance()
                    self.add_token(TokenType.EFFECT, '⊕', start_line, start_col)
                    continue
                elif two == '!=':
                    self.advance(); self.advance()
                    self.add_token(TokenType.NEQ, '!=', start_line, start_col)
                    continue
                elif two == '<=':
                    self.advance(); self.advance()
                    self.add_token(TokenType.LTE, '<=', start_line, start_col)
                    continue
                elif two == '>=':
                    self.advance(); self.advance()
                    self.add_token(TokenType.GTE, '>=', start_line, start_col)
                    continue
                elif two == '<|':
                    self.advance(); self.advance()
                    self.add_token(TokenType.LANGLE, '⟨', start_line, start_col)
                    continue
                elif two == '|>':
                    self.advance(); self.advance()
                    self.add_token(TokenType.RANGLE, '⟩', start_line, start_col)
                    continue
                elif two == '[|':
                    self.advance(); self.advance()
                    self.add_token(TokenType.LBRACKET, '⟦', start_line, start_col)
                    continue
                elif two == '|]':
                    self.advance(); self.advance()
                    self.add_token(TokenType.RBRACKET, '⟧', start_line, start_col)
                    continue

            # Backslash as lambda
            if ch == '\\':
                self.advance()
                self.add_token(TokenType.LAMBDA, 'λ', start_line, start_col)
                continue

            # Single-character tokens
            if ch in SINGLE_CHAR:
                # Special case: _ followed by . should be read as symbol for dot-indexing
                if ch == '_' and self.pos + 1 < len(self.source) and self.source[self.pos + 1] == '.':
                    name = self.read_symbol()
                    self.add_token(TokenType.SYMBOL, name, start_line, start_col)
                    continue
                self.advance()
                self.add_token(SINGLE_CHAR[ch], ch, start_line, start_col)
                continue

            # Less-than / Greater-than (single char, after checking two-char combos)
            if ch == '<':
                self.advance()
                self.add_token(TokenType.LT, '<', start_line, start_col)
                continue
            if ch == '>':
                self.advance()
                self.add_token(TokenType.GT, '>', start_line, start_col)
                continue

            # Pipe (single, after checking ||)
            if ch == '|':
                self.advance()
                self.add_token(TokenType.OR, '|', start_line, start_col)
                continue

            # Bang (single, after checking !=)
            if ch == '!':
                self.advance()
                self.add_token(TokenType.NOT, '!', start_line, start_col)
                continue

            # Hash identifiers: #abc
            if ch == '#':
                self.advance()
                name = self.read_symbol()
                if not name:
                    raise LexError("Expected identifier after #", start_line, start_col)
                self.add_token(TokenType.HASH_ID, '#' + name, start_line, start_col)
                continue

            # Address references: @abc
            if ch == '@':
                self.advance()
                name = self.read_symbol()
                if not name:
                    raise LexError("Expected identifier after @", start_line, start_col)
                self.add_token(TokenType.ADDR_REF, '@' + name, start_line, start_col)
                continue

            # Template strings: $"Hi {name}! Age: {(+ age 1)}"
            if ch == '$' and self.pos + 1 < len(self.source) and self.source[self.pos + 1] == '"':
                self.advance()  # consume $
                self.advance()  # consume "
                fmt_parts = []  # literal string parts
                expr_sources = []  # expression source strings
                current = []
                while self.pos < len(self.source) and self.source[self.pos] != '"':
                    c = self.source[self.pos]
                    if c == '\\':
                        self.advance()
                        escape_map = {'n': '\n', 't': '\t', 'r': '\r', '\\': '\\', '"': '"', '{': '{', '}': '}'}
                        if self.pos < len(self.source):
                            current.append(escape_map.get(self.source[self.pos], self.source[self.pos]))
                            self.advance()
                    elif c == '{':
                        # Start of expression
                        fmt_parts.append(''.join(current))
                        current = []
                        self.advance()  # consume {
                        depth = 1
                        expr_chars = []
                        while self.pos < len(self.source) and depth > 0:
                            ec = self.source[self.pos]
                            if ec == '{': depth += 1
                            elif ec == '}': depth -= 1
                            if depth > 0:
                                expr_chars.append(ec)
                            self.advance()
                        expr_sources.append(''.join(expr_chars))
                    else:
                        current.append(c)
                        self.advance()
                fmt_parts.append(''.join(current))
                if self.pos < len(self.source):
                    self.advance()  # consume closing "
                # Build format string with {} placeholders
                fmt_str = '{}' .join(fmt_parts)
                # Encode as: fmt_str\x00expr1\x00expr2\x00...
                value = '\x00'.join([fmt_str] + expr_sources)
                self.add_token(TokenType.TEMPLATE, value, start_line, start_col)
                continue

            # String literals
            if ch == '"':
                value = self.read_string()
                self.add_token(TokenType.STRING, value, start_line, start_col)
                continue

            # Numbers (including negative)
            if ch.isdigit():
                tok = self.read_number()
                self.tokens.append(tok)
                continue

            # Minus: could be negative number or operator
            if ch == '-':
                # Treat as negative number if followed by digit AND
                # not immediately after ( (where it would be the operator in (- a b))
                if (self.peek_ahead().isdigit() and
                    (not self.tokens or
                     self.tokens[-1].type != TokenType.LPAREN)):
                    tok = self.read_number()
                    self.tokens.append(tok)
                    continue
                else:
                    self.advance()
                    self.add_token(TokenType.MINUS, '-', start_line, start_col)
                    continue

            # Symbols / keywords
            if ch.isalpha() or ch == '_':
                name = self.read_symbol()
                if name in KEYWORDS:
                    self.add_token(KEYWORDS[name], name, start_line, start_col)
                else:
                    self.add_token(TokenType.SYMBOL, name, start_line, start_col)
                continue

            raise LexError(f"Unexpected character: {ch!r}", start_line, start_col)

        self.add_token(TokenType.EOF, '', self.line, self.col)
        return self.tokens


def tokenize(source: str) -> List[Token]:
    """Convenience function to tokenize a source string"""
    return Lexer(source).tokenize()
