from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional


class LexError(Exception):
    def __init__(self, message: str, line: int, col: int):
        super().__init__(f"LexError at {line}:{col} - {message}")
        self.line = line
        self.col = col


@dataclass(frozen=True)
class Token:
    kind: str
    value: Optional[str]
    line: int
    col: int


KEYWORDS = {
    "Resource",
    "State",
    "Rule",
    "Proof",
    "verify",
    "external",
    "true",
    "false",
    "sophiResource",
}


def _is_ident_start(ch: str) -> bool:
    return ch.isalpha() or ch == "_"


def _is_ident_part(ch: str) -> bool:
    return ch.isalnum() or ch == "_"


def tokenise(source: str) -> List[Token]:
    tokens: List[Token] = []
    indent_stack = [0]

    lines = source.splitlines()
    for line_idx, raw_line in enumerate(lines, start=1):
        line = raw_line.rstrip("\r\n")
        # Strip comments
        if "#" in line:
            line = line.split("#", 1)[0]

        if not line.strip():
            continue

        # Indentation
        indent = 0
        while indent < len(line) and line[indent] == " ":
            indent += 1
        if indent < len(line) and line[indent] == "\t":
            raise LexError("Tabs are not allowed for indentation", line_idx, indent + 1)

        if indent > indent_stack[-1]:
            indent_stack.append(indent)
            tokens.append(Token("INDENT", None, line_idx, 1))
        while indent < indent_stack[-1]:
            indent_stack.pop()
            tokens.append(Token("DEDENT", None, line_idx, 1))
        if indent != indent_stack[-1]:
            raise LexError("Inconsistent indentation", line_idx, 1)

        i = indent
        while i < len(line):
            ch = line[i]
            if ch.isspace():
                i += 1
                continue

            # Two-char operators
            if i + 1 < len(line):
                two = line[i : i + 2]
                if two in ("==", "!=", "<=", ">="):
                    tokens.append(Token("OP", two, line_idx, i + 1))
                    i += 2
                    continue

            # Single-char operators / punctuation
            if ch in ":.=()+-*/<>":
                tokens.append(Token("OP", ch, line_idx, i + 1))
                i += 1
                continue

            # Number literal
            if ch.isdigit():
                start = i
                while i < len(line) and line[i].isdigit():
                    i += 1
                if i < len(line) and line[i] == ".":
                    i += 1
                    if i >= len(line) or not line[i].isdigit():
                        raise LexError("Invalid number literal", line_idx, i + 1)
                    while i < len(line) and line[i].isdigit():
                        i += 1
                value = line[start:i]
                tokens.append(Token("NUMBER", value, line_idx, start + 1))
                continue

            # Identifier / keyword
            if _is_ident_start(ch):
                start = i
                i += 1
                while i < len(line) and _is_ident_part(line[i]):
                    i += 1
                value = line[start:i]
                kind = "KEYWORD" if value in KEYWORDS else "IDENT"
                tokens.append(Token(kind, value, line_idx, start + 1))
                continue

            raise LexError(f"Unexpected character '{ch}'", line_idx, i + 1)

        tokens.append(Token("NEWLINE", None, line_idx, len(line) + 1))

    while len(indent_stack) > 1:
        indent_stack.pop()
        tokens.append(Token("DEDENT", None, len(lines) + 1, 1))

    tokens.append(Token("EOF", None, len(lines) + 1, 1))
    return tokens
