"""
Nexa Language — Lexer v27
New: PIPE_ERR (|?), TEMPLATE string, NEWTYPE, SPAWN, CHANNEL tokens
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from enum import Enum, auto
from typing import List, Optional


class TT(Enum):
    # Literals
    NUMBER = auto(); STRING = auto(); FSTRING = auto(); TRIPLE_STRING = auto(); BOOL = auto(); NULL = auto()
    # Identifiers & keywords
    IDENT = auto(); DEF = auto(); LET = auto(); RETURN = auto()
    IF = auto(); ELSE = auto(); FOR = auto(); IN = auto(); WHILE = auto()
    BREAK = auto(); CONTINUE = auto()
    IMPORT = auto(); FROM = auto(); AS = auto()
    CLASS = auto(); SELF = auto(); SUPER = auto(); NEW = auto()
    TRY = auto(); CATCH = auto(); FINALLY = auto(); THROW = auto()
    MATCH = auto(); CASE = auto(); DEFAULT = auto()
    TEST = auto(); ASSERT = auto()
    BENCH = auto()
    YIELD = auto(); STRUCT = auto(); ENUM = auto()
    PRIVATE = auto(); PUBLIC = auto()
    INTERFACE = auto(); IMPLEMENTS = auto()
    ABSTRACT = auto(); EXPORT = auto(); LOOP = auto()
    CONST = auto()
    SEALED = auto(); MACRO = auto(); TYPE = auto()   # v24
    FN = auto(); VAR = auto()                        # v25: fn = lambda kw, var = mutable let
    TRAIT = auto(); IMPL = auto(); USE = auto()      # v25: trait/impl system + use imports
    NEWTYPE = auto(); TEMPLATE = auto(); SPAWN = auto()  # v27: promoted to real tokens
    # AI keywords
    MODEL = auto(); TRAIN = auto(); PREDICT = auto(); LAYER = auto()
    AUTOML = auto(); ASYNC = auto(); AWAIT = auto()
    # Operators
    PLUS = auto(); MINUS = auto(); STAR = auto(); SLASH = auto(); MODULO = auto()
    POW = auto()   # **
    EQ = auto(); NEQ = auto(); LT = auto(); GT = auto(); LTE = auto(); GTE = auto()
    AND = auto(); OR = auto(); NOT = auto()
    AMPERSAND = auto(); VBAR = auto(); CARET = auto(); TILDE = auto()
    LSHIFT = auto(); RSHIFT = auto()
    ASSIGN = auto(); PLUS_EQ = auto(); MINUS_EQ = auto(); STAR_EQ = auto(); SLASH_EQ = auto()
    AMPERSAND_EQ = auto(); VBAR_EQ = auto(); CARET_EQ = auto()
    LSHIFT_EQ = auto(); RSHIFT_EQ = auto()
    DOT = auto(); COMMA = auto(); COLON = auto(); ARROW = auto()
    FAT_ARROW = auto() # =>
    PIPE = auto()      # |>
    PIPE_ERR = auto()  # |?  — error pipe (v27: runs only on Err)
    AT = auto()        # @
    QUESTION = auto(); QUESTION_QUESTION = auto(); QUESTION_DOT = auto()
    ELLIPSIS = auto() # ...
    DOTDOT = auto()   # ..
    DOTDOT_EQ = auto() # ..=
    BANG = auto()     # ! (non-null assertion)
    # Delimiters
    LPAREN = auto(); RPAREN = auto()
    LBRACE = auto(); RBRACE = auto()
    LBRACKET = auto(); RBRACKET = auto()
    NEWLINE = auto(); SEMICOLON = auto(); EOF = auto()


KEYWORDS = {
    "def": TT.DEF, "let": TT.LET, "return": TT.RETURN,
    "if": TT.IF, "else": TT.ELSE, "for": TT.FOR, "in": TT.IN, "while": TT.WHILE,
    "break": TT.BREAK, "continue": TT.CONTINUE,
    "import": TT.IMPORT, "from": TT.FROM, "as": TT.AS,
    "class": TT.CLASS, "self": TT.SELF, "super": TT.SUPER, "new": TT.NEW,
    "try": TT.TRY, "catch": TT.CATCH, "finally": TT.FINALLY, "throw": TT.THROW,
    "match": TT.MATCH, "case": TT.CASE, "default": TT.DEFAULT,
    "test": TT.TEST, "assert": TT.ASSERT, "bench": TT.BENCH,
    "yield": TT.YIELD, "struct": TT.STRUCT, "enum": TT.ENUM,
    "private": TT.PRIVATE, "public": TT.PUBLIC,
    "interface": TT.INTERFACE, "implements": TT.IMPLEMENTS,
    "abstract": TT.ABSTRACT, "export": TT.EXPORT, "loop": TT.LOOP,
    "const": TT.CONST,
    "sealed": TT.SEALED, "macro": TT.MACRO, "type": TT.TYPE,  # v24
    "fn": TT.FN, "var": TT.VAR,                               # v25
    "trait": TT.TRAIT, "impl": TT.IMPL, "use": TT.USE,        # v25
    "model": TT.MODEL, "train": TT.TRAIN, "predict": TT.PREDICT,
    "layer": TT.LAYER, "automl": TT.AUTOML,
    "async": TT.ASYNC, "await": TT.AWAIT,
    "true": TT.BOOL, "false": TT.BOOL, "null": TT.NULL,
    "and": TT.AND, "or": TT.OR, "not": TT.NOT,
    # v27 keywords — promoted to real token types (not IDENT)
    "newtype": TT.NEWTYPE, "template": TT.TEMPLATE, "spawn": TT.SPAWN,
}


@dataclass
class Token:
    type: TT
    value: object
    line: int

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


class LexError(Exception):
    pass


class Lexer:
    """Tokenise Nexa v3 source code."""

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

    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
        return ch

    def _match(self, expected: str) -> bool:
        if self.pos < len(self.src) and self.src[self.pos] == expected:
            self.pos += 1
            return True
        return False

    def _add(self, tt: TT, value: object = None):
        self.tokens.append(Token(tt, value, self.line))

    def tokenize(self) -> List[Token]:
        while self.pos < len(self.src):
            self._scan_token()
        self._add(TT.EOF)
        return self.tokens

    def _scan_token(self):
        ch = self._advance()

        # Whitespace
        if ch in (" ", "\t", "\r"):
            return

        # Multi-line comments /* ... */
        if ch == "/" and self._peek() == "*":
            self._advance()
            while self.pos < len(self.src) - 1:
                if self._peek() == "*" and self._peek(1) == "/":
                    self._advance(); self._advance(); break
                self._advance()
            return

        # Single-line comment
        if ch == "/" and self._peek() == "/":
            while self._peek() and self._peek() != "\n":
                self._advance()
            return

        # Newline / semicolon as statement separator
        if ch == "\n":
            self._add(TT.NEWLINE); return
        if ch == ";":
            self._add(TT.SEMICOLON); return

        # Numbers (int, float, scientific, hex, binary)
        if ch.isdigit() or (ch == "." and self._peek().isdigit()):
            # Hexadecimal
            if ch == "0" and self._peek() in ("x", "X"):
                self._advance() # consume x
                hex_val = ""
                while self._peek() and self._peek().lower() in "0123456789abcdef":
                    hex_val += self._advance()
                self._add(TT.NUMBER, int(hex_val, 16))
                return
            # Binary
            if ch == "0" and self._peek() in ("b", "B"):
                self._advance() # consume b
                bin_val = ""
                while self._peek() in ("0", "1"):
                    bin_val += self._advance()
                self._add(TT.NUMBER, int(bin_val, 2))
                return
            
            num = ch
            while self._peek().isdigit() or (
                self._peek() == "." and self._peek(1) != "."  # don't eat .. range op
            ):
                num += self._advance()
            if self._peek() in ("e", "E"):
                num += self._advance()
                if self._peek() in ("+", "-"):
                    num += self._advance()
                while self._peek().isdigit():
                    num += self._advance()
            self._add(TT.NUMBER, float(num) if ("." in num or "e" in num.lower()) else int(num))
            return

        # F-strings: f"..." or f'...'
        if ch == "f" and self._peek() in ('"', "'"):
            q = self._advance()
            parts = []
            raw = ""
            while self._peek() and self._peek() != q:
                c = self._advance()
                # v26: {{ → literal '{',  }} → literal '}'
                if c == "{" and self._peek() == "{":
                    self._advance()   # consume second {
                    raw += "{"
                elif c == "}" and self._peek() == "}":
                    self._advance()   # consume second }
                    raw += "}"
                elif c == "{":
                    if raw:
                        parts.append(("lit", raw)); raw = ""
                    expr = ""
                    depth = 1
                    while self._peek() and depth > 0:
                        nc = self._advance()
                        if nc == "{": depth += 1
                        elif nc == "}": depth -= 1
                        if depth > 0: expr += nc
                    parts.append(("expr", expr))
                elif c == "\\" and self._peek() in ('"', "'", "n", "t", "\\"):
                    esc = {"n": "\n", "t": "\t", "\\": "\\", '"': '"', "'": "'"}
                    raw += esc.get(self._advance(), "")
                else:
                    raw += c
            self._advance()  # closing quote
            if raw:
                parts.append(("lit", raw))
            self._add(TT.FSTRING, parts)
            return

        # Triple-quoted strings: """...""" or '''...'''
        if ch in ('"', "'") and self._peek() == ch and self._peek(1) == ch:
            self._advance(); self._advance()  # consume remaining two quotes
            q3 = ch * 3
            s = ""
            while self.pos < len(self.src) - 2:
                if self.src[self.pos:self.pos+3] == q3:
                    self.pos += 3; break
                s += self._advance()
            self._add(TT.TRIPLE_STRING, s)
            return

        # Regular strings
        if ch in ('"', "'"):
            s = ""
            while self._peek() and self._peek() != ch:
                c = self._advance()
                if c == "\\" and self._peek() in ('"', "'", "n", "t", "\\"):
                    esc = {"n": "\n", "t": "\t", "\\": "\\", '"': '"', "'": "'"}
                    s += esc.get(self._advance(), "")
                else:
                    s += c
            self._advance()
            self._add(TT.STRING, s)
            return

        # Identifiers / keywords
        if ch.isalpha() or ch == "_":
            ident = ch
            while self._peek().isalnum() or self._peek() == "_":
                ident += self._advance()
            tt = KEYWORDS.get(ident, TT.IDENT)
            value = {"true": True, "false": False, "null": None}.get(ident, ident)
            self._add(tt, value)
            return

        # Single-char operators (no multi-char variants)
        MAP = {
            "%": TT.MODULO, "@": TT.AT, "~": TT.TILDE,
            ",": TT.COMMA, ":": TT.COLON,
            "(": TT.LPAREN, ")": TT.RPAREN,
            "{": TT.LBRACE, "}": TT.RBRACE,
            "[": TT.LBRACKET, "]": TT.RBRACKET,
        }
        if ch in MAP:
            self._add(MAP[ch], ch); return

        # Multi-char operators
        if ch == ".":
            if self._peek() == ".":
                self._advance()   # consume second dot
                if self._peek() == "=":
                    self._advance(); self._add(TT.DOTDOT_EQ, "..=")
                elif self._peek() == ".":
                    self._advance(); self._add(TT.ELLIPSIS, "...")
                else:
                    self._add(TT.DOTDOT, "..")
            else:
                self._add(TT.DOT, ".")
            return
        if ch == "?":
            if self._match("?"): self._add(TT.QUESTION_QUESTION, "??")
            elif self._match("."): self._add(TT.QUESTION_DOT, "?.")
            else: self._add(TT.QUESTION, "?")
            return
        if ch == "-":
            if self._match(">"):  self._add(TT.ARROW)
            elif self._match("="): self._add(TT.MINUS_EQ)
            else:                  self._add(TT.MINUS, ch)
            return
        if ch == "+":
            self._add(TT.PLUS_EQ if self._match("=") else TT.PLUS, ch); return
        if ch == "*":
            if self._match("*"):  self._add(TT.POW)
            elif self._match("="): self._add(TT.STAR_EQ)
            else:                  self._add(TT.STAR, ch)
            return
        if ch == "/":
            if self._match("="): self._add(TT.SLASH_EQ)
            else:                self._add(TT.SLASH, ch)
            return
        if ch == "=":
            if self._match(">"): self._add(TT.FAT_ARROW)
            elif self._match("="): self._add(TT.EQ)
            else: self._add(TT.ASSIGN)
            return
        if ch == "!":
            if self._match("="): self._add(TT.NEQ)
            else: self._add(TT.BANG)  # non-null assertion / logical not handled by NOT keyword
            return
        if ch == "<":
            if self._match("="): self._add(TT.LTE)
            elif self._match("<"):
                if self._match("="): self._add(TT.LSHIFT_EQ)
                else: self._add(TT.LSHIFT)
            else: self._add(TT.LT)
            return
        if ch == ">":
            if self._match("="): self._add(TT.GTE)
            elif self._match(">"):
                if self._match("="): self._add(TT.RSHIFT_EQ)
                else: self._add(TT.RSHIFT)
            else: self._add(TT.GT)
            return
        if ch == "&":
            if self._match("&"): self._add(TT.AND)
            elif self._match("="): self._add(TT.AMPERSAND_EQ)
            else: self._add(TT.AMPERSAND)
            return
        if ch == "|":
            if self._match(">"):   self._add(TT.PIPE)       # |>
            elif self._match("?"): self._add(TT.PIPE_ERR)   # |?  v27 error pipe
            elif self._match("|"): self._add(TT.OR)
            elif self._match("="): self._add(TT.VBAR_EQ)
            else: self._add(TT.VBAR)
            return
        if ch == "^":
            if self._match("="): self._add(TT.CARET_EQ)
            else: self._add(TT.CARET)
            return

        raise LexError(f"Unexpected character {ch!r} at line {self.line}")
