from __future__ import annotations

import re
from dataclasses import dataclass
from fractions import Fraction
from decimal import Decimal
from typing import Any, Callable

###############################################################################
# Runtime values
###############################################################################

@dataclass(frozen=True)
class Sym:
    name: str

Value = Any  # Fraction | Sym | str (meak's string results mainly from ? and for/o rules)

def is_num(v: Value) -> bool:
    return isinstance(v, Fraction)

def to_str(v: Value) -> str:
    if isinstance(v, Fraction):
        if v.denominator == 1:
            return str(v.numerator)
        return f"{v.numerator}/{v.denominator}"
    if isinstance(v, Sym):
        return v.name
    return str(v)

def atom_to_value(tok: str) -> Value:
    # exact numeric parsing
    if re.fullmatch(r"[+-]?\d+/\d+", tok):
        a, b = tok.split("/")
        return Fraction(int(a), int(b))
    if re.fullmatch(r"[+-]?\d+\.\d+", tok):
        return Fraction(Decimal(tok))  # exact rational from decimal string
    if re.fullmatch(r"[+-]?\d+", tok):
        return Fraction(int(tok), 1)
    return Sym(tok)

###############################################################################
# Lexer
###############################################################################

TOKEN_RE = re.compile(
    r"""
    \s*(
        \#.*$                              |  # comment
        [A-Za-z_][A-Za-z0-9_]*             |  # identifiers / keywords
        [+-]?\d+/\d+                       |  # rationals
        [+-]?\d+\.\d+                      |  # decimals
        [+-]?\d+                           |  # ints
        \?                                 |  # ?
        \^                                 |  # ^
        [()+\-*/=,;]                          # ops / parens / separators
    )\s*
    """,
    re.VERBOSE,
)

OPS = {"o", "for", "+", "-", "*", "/", "^"}

def tokenize(src: str) -> list[str]:
    src = src.strip()
    if not src or src.startswith("#"):
        return []
    pos = 0
    toks: list[str] = []
    while pos < len(src):
        m = TOKEN_RE.match(src, pos)
        if not m:
            raise SyntaxError(f"Unexpected character at pos {pos}: {src[pos:pos+20]!r}")
        tok = m.group(1)
        pos = m.end()
        if tok.startswith("#"):
            break
        toks.append(tok)
    # normalize keywords (o/for) so they behave like operators
    return toks

###############################################################################
# Pratt parser
###############################################################################

class ParseError(Exception):
    pass

@dataclass
class TokStream:
    toks: list[str]
    i: int = 0

    def peek(self) -> str | None:
        return self.toks[self.i] if self.i < len(self.toks) else None

    def pop(self) -> str | None:
        t = self.peek()
        if t is None:
            return None
        self.i += 1
        return t

# AST:
# ("assign", name, expr)
# expr:
# ("num", Fraction) | ("sym", name) | ("call", name, [args])
# ("un", op, expr) | ("bin", op, left, right)

BP = {
    "o": 10,     # lowest
    "for": 20,
    "+": 30, "-": 30,
    "*": 40, "/": 40,
    "^": 60,     # right-assoc
}

RIGHT_ASSOC = {"^"}

def lbp(op: str) -> int:
    return BP.get(op, 0)

def parse_stmt(ts: TokStream):
    # assignment: IDENT "=" expr
    if ts.peek() and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", ts.peek() or ""):
        name = ts.peek()
        if ts.i + 1 < len(ts.toks) and ts.toks[ts.i + 1] == "=":
            ts.pop()  # name
            ts.pop()  # =
            expr = parse_expr(ts, 0)
            if ts.peek() is not None:
                raise ParseError(f"Extra tokens after assignment: {ts.toks[ts.i:]}")
            return ("assign", name, expr)

    expr = parse_expr(ts, 0)
    if ts.peek() is not None:
        raise ParseError(f"Extra tokens: {ts.toks[ts.i:]}")
    return expr

def parse_expr(ts: TokStream, min_bp: int):
    left = parse_prefix(ts)

    while True:
        op = ts.peek()
        if op is None:
            break

        # keyword operators
        if op not in BP:
            # allow o/for as identifiers if user writes them weirdly? we treat them as ops only
            break

        bp = lbp(op)
        if bp < min_bp:
            break

        ts.pop()  # consume op
        next_min = bp if op in RIGHT_ASSOC else bp + 1
        right = parse_expr(ts, next_min)
        left = ("bin", op, left, right)

    return left

def parse_prefix(ts: TokStream):
    tok = ts.pop()
    if tok is None:
        raise ParseError("Unexpected end of input")

    # unary ops
    if tok in ("?", "+", "-"):
        expr = parse_expr(ts, 70)
        return ("un", tok, expr)

    if tok == "(":
        expr = parse_expr(ts, 0)
        if ts.pop() != ")":
            raise ParseError("Expected ')'")
        return expr

    # number
    if re.fullmatch(r"[+-]?\d+/\d+|[+-]?\d+\.\d+|[+-]?\d+", tok):
        v = atom_to_value(tok)
        if not isinstance(v, Fraction):
            raise ParseError("Internal numeric parse error")
        return ("num", v)

    # identifier / call / keyword-op tokens
    if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", tok):
        # treat "o" and "for" as operators only; if they appear here, it's user error
        if tok in ("o", "for"):
            raise ParseError(f"'{tok}' cannot appear as an atom; use it as an operator")

        if ts.peek() == "(":
            ts.pop()  # (
            args = []
            if ts.peek() != ")":
                while True:
                    args.append(parse_expr(ts, 0))
                    if ts.peek() == ",":
                        ts.pop()
                        continue
                    break
            if ts.pop() != ")":
                raise ParseError("Expected ')' after arguments")
            return ("call", tok, args)

        return ("sym", tok)

    raise ParseError(f"Unexpected token: {tok}")

###############################################################################
# Interpreter
###############################################################################

def _expect_num(x: Value, fname="value") -> Fraction:
    if not isinstance(x, Fraction):
        raise RuntimeError(f"{fname} expects a number, got {to_str(x)}")
    return x

def builtin_abs(x: Value) -> Fraction:
    x = _expect_num(x, "abs")
    return abs(x)

def builtin_pow(a: Value, b: Value) -> Fraction:
    a = _expect_num(a, "pow")
    b = _expect_num(b, "pow")
    if b.denominator != 1:
        raise RuntimeError("pow exponent must be an integer")
    return a ** b.numerator

def builtin_gcd(a: Value, b: Value) -> Fraction:
    a = _expect_num(a, "gcd")
    b = _expect_num(b, "gcd")
    if a.denominator != 1 or b.denominator != 1:
        raise RuntimeError("gcd expects integers")
    import math
    return Fraction(math.gcd(a.numerator, b.numerator), 1)

def builtin_lcm(a: Value, b: Value) -> Fraction:
    a = _expect_num(a, "lcm")
    b = _expect_num(b, "lcm")
    if a.denominator != 1 or b.denominator != 1:
        raise RuntimeError("lcm expects integers")
    import math
    return Fraction(abs(a.numerator * b.numerator) // math.gcd(a.numerator, b.numerator), 1)

def format_value(v: Value) -> str:
    if isinstance(v, Sym):
        return v.name
    if isinstance(v, Fraction):
        return to_str(v)
    return str(v)

class MeakInterpreter:
    """
    meak v0.1:
      - exact rational arithmetic with Fraction
      - operators:
          ?x        => to_str(value(x)) + xname
          X for Y   => if both numbers: add; else concat strings
          A o B     => if B is bound variable: value(B)
                     else if A is (?x) and B is symbol: x + B + "o" + B
                     else value(A)
      - + - * / ^ (math), ^ exponent must be integer
      - ; separates statements on the same line
    """

    def __init__(self):
        self.env: dict[str, Value] = {}
        # builtins
        # 'get'/'how' mini-state: accumulate numeric inputs via `get <expr>` and
        # report with `how`.
        self._how_total: Fraction = Fraction(0, 1)
        self.env["abs"] = builtin_abs
        self.env["pow"] = builtin_pow
        self.env["gcd"] = builtin_gcd
        self.env["lcm"] = builtin_lcm

    def _sym_value(self, name: str) -> Value:
        v = self.env.get(name, Sym(name))
        return v

    def _eval_expr_src(self, src: str) -> Value:
        """Evaluate a pure expression string (no assignments)."""
        toks = tokenize(src)
        if not toks:
            raise ParseError("Expected an expression after 'get'")
        ts = TokStream(toks)
        expr = parse_expr(ts, 0)
        if ts.peek() is not None:
            raise ParseError(f"Extra tokens in expression: {ts.toks[ts.i:]}")
        return self.eval_expr(expr)

    def eval_expr(self, node) -> Value:
        kind = node[0]

        if kind == "num":
            return node[1]

        if kind == "sym":
            return self._sym_value(node[1])

        if kind == "call":
            name, args_ast = node[1], node[2]
            fn = self.env.get(name)
            if fn is None or not callable(fn):
                raise RuntimeError(f"Unknown function: {name}")
            args = [self.eval_expr(a) for a in args_ast]
            return fn(*args)

        if kind == "un":
            op, x_ast = node[1], node[2]
            x = self.eval_expr(x_ast)

            if op == "+":
                return _expect_num(x, "unary +")
            if op == "-":
                return -_expect_num(x, "unary -")

            if op == "?":
                # Prefer variable name if operand is a symbol node
                if isinstance(x_ast, tuple) and x_ast[0] == "sym":
                    xname = x_ast[1]
                elif isinstance(x, Sym):
                    xname = x.name
                else:
                    xname = to_str(x)
                return to_str(x) + xname

            raise RuntimeError(f"Unknown unary op: {op}")

        if kind == "bin":
            op, a_ast, b_ast = node[1], node[2], node[3]

            # meak: o
            if op == "o":
                # if right is a symbol bound in env (and not a function): return its value
                if b_ast[0] == "sym":
                    bname = b_ast[1]
                    if bname in self.env and not callable(self.env[bname]):
                        return self.env[bname]

                # special: (?x) o b  =>  x + b + "o" + b
                if (
                    a_ast[0] == "un" and a_ast[1] == "?"
                    and a_ast[2][0] == "sym"
                    and b_ast[0] == "sym"
                ):
                    xname = a_ast[2][1]
                    bname = b_ast[1]
                    return xname + bname + "o" + bname

                return self.eval_expr(a_ast)

            # meak: for
            if op == "for":
                va = self.eval_expr(a_ast)
                vb = self.eval_expr(b_ast)
                if is_num(va) and is_num(vb):
                    return va + vb
                return to_str(va) + to_str(vb)

            # math ops
            va = self.eval_expr(a_ast)
            vb = self.eval_expr(b_ast)

            if op in {"+", "-", "*", "/", "^"}:
                va = _expect_num(va, op)
                vb = _expect_num(vb, op)

                if op == "+":
                    return va + vb
                if op == "-":
                    return va - vb
                if op == "*":
                    return va * vb
                if op == "/":
                    if vb == 0:
                        raise RuntimeError("Division by zero")
                    return va / vb
                if op == "^":
                    if vb.denominator != 1:
                        raise RuntimeError("Exponent must be an integer")
                    return va ** vb.numerator

            raise RuntimeError(f"Unknown binary op: {op}")

        raise RuntimeError(f"Unknown AST node: {node}")

    def run_stmt(self, src: str) -> Value | None:
        toks = tokenize(src)
        if not toks:
            return None

        # Convert keyword operators in token stream:
        # We keep them as strings but parser expects them in BP table.
        # So we rewrite identifiers "o"/"for" into operators by leaving them as is,
        # but we must ensure they're not parsed as atoms. parse_prefix already blocks.
        # Good.

        ts = TokStream(toks)
        ast = parse_stmt(ts)

        if isinstance(ast, tuple) and ast and ast[0] == "assign":
            _, name, expr = ast
            val = self.eval_expr(expr)
            self.env[name] = val
            return None

        return self.eval_expr(ast)
    def run_line(self, line: str) -> list[str]:
        # support multiple statements with ';'
        outs: list[str] = []
        for part in [p.strip() for p in line.split(";")]:
            if not part or part.startswith("#"):
                continue

            # mini-commands (outside the meak expression grammar)
            # cycle: repeat a command N times. Syntax: cycle <cmd> for <N>
            # Example: cycle get 100 for 50
            if part.startswith("cycle "):
                rest = part[len("cycle "):].strip()
                cmd_src, sep, n_src = rest.rpartition(" for ")
                if not sep:
                    raise ParseError("Usage: cycle <cmd> for <N>")
                cmd_src = cmd_src.strip()
                n_src = n_src.strip()
                if not cmd_src or not n_src:
                    raise ParseError("Usage: cycle <cmd> for <N>")
                try:
                    n = int(n_src, 10)
                except ValueError:
                    raise ParseError("cycle: <N> must be an integer")
                if n < 0:
                    raise ParseError("cycle: <N> must be >= 0")
                if n > 100000:
                    raise ParseError("cycle: <N> too large (max 100000)")
                for _ in range(n):
                    outs.extend(self.run_line(cmd_src))
                continue

            if part == "how":
                outs.append(format_value(self._how_total))
                continue

            if part == "clear":
                self._how_total = Fraction(0, 1)
                continue

            if part == "get" or part.startswith("get "):
                # allow: "get 100", "get   1+2*3"
                # we only accept a pure expression after get (no assignment)
                _, _, expr_src = part.partition(" ")
                expr_src = expr_src.strip()
                if not expr_src:
                    raise ParseError("Usage: get <expr>")
                v = self._eval_expr_src(expr_src)
                v_num = _expect_num(v, "get")
                self._how_total += v_num
                continue

            v = self.run_stmt(part)
            if v is not None:
                outs.append(format_value(v))
        return outs

    def run_file(self, path: str) -> list[str]:
        outs: list[str] = []
        with open(path, "r", encoding="utf-8") as f:
            for raw in f:
                raw = raw.strip()
                if not raw or raw.startswith("#"):
                    continue
                for out in self.run_line(raw):
                    outs.append(out)
        return outs

    def repl(self):
        print("meak v0.1 (exact rational math).  :q to quit, :env to list vars.  Also: get <expr>, how, clear, cycle <cmd> for <N>.")
        while True:
            try:
                line = input("meak> ").strip()
            except EOFError:
                print()
                break

            if not line:
                continue
            if line in (":q", ":quit", "exit", "quit"):
                break
            if line == ":env":
                for k, v in sorted(self.env.items()):
                    if callable(v):
                        continue
                    print(f"{k} = {format_value(v)}")
                continue

            try:
                outs = self.run_line(line)
                for o in outs:
                    print(o)
            except Exception as e:
                print(f"Error: {e}")


def main():
    """Entry point for CLI usage."""
    MeakInterpreter().repl()


if __name__ == "__main__":
    main()
