"""
TENGWAR Parser

Transforms token stream into AST. The syntax is strict prefix notation
with explicit structure — no ambiguity, no operator precedence needed.

Grammar (simplified):
    program   := expr*
    expr      := atom | list | bind_expr | tuple | vector
    atom      := INT | FLOAT | STRING | TRUE | FALSE | SYMBOL | HASH_ID | ADDR_REF | UNIT
    list      := '(' head expr* ')'
    bind_expr := expr '→' target
    tuple     := '⟨' expr* '⟩'
    vector    := '⟦' expr* '⟧'
    head      := LAMBDA | COND | MATCH | SEQ | PARALLEL | MODULE | RECURSE
               | DEFINE | TYPE | PROOF | EFFECT | MUTATE | IMPORT | operator | expr
    operator  := + | - | * | / | % | = | != | < | > | <= | >= | & | |
"""
from typing import List, Optional, Tuple as PyTuple
from .lexer import Token, TokenType
from .ast_nodes import *
from .errors import ParseError


class Parser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0

    def peek(self) -> Token:
        return self.tokens[self.pos]

    def peek_type(self) -> TokenType:
        return self.tokens[self.pos].type

    def advance(self) -> Token:
        tok = self.tokens[self.pos]
        self.pos += 1
        return tok

    def expect(self, type: TokenType) -> Token:
        tok = self.advance()
        if tok.type != type:
            raise ParseError(
                f"Expected {type.name}, got {tok.type.name} ({tok.value!r})",
                tok.line, tok.col
            )
        return tok

    def at_end(self) -> bool:
        return self.peek_type() == TokenType.EOF

    def parse(self) -> Program:
        """Parse full program"""
        body = []
        while not self.at_end():
            if self.peek_type() == TokenType.COMMENT:
                self.advance()  # skip comments at top level
                continue
            expr = self.parse_expr()
            if expr is not None:
                # Check for bind: expr → target
                if not self.at_end() and self.peek_type() == TokenType.ARROW:
                    expr = self.parse_bind(expr)
                body.append(expr)
        return Program(body=body, type=None)

    def parse_expr(self) -> ASTNode:
        """Parse a single expression"""
        tok = self.peek()

        # Skip comments
        while tok.type == TokenType.COMMENT:
            self.advance()
            tok = self.peek()

        if tok.type == TokenType.LPAREN:
            return self.parse_list()
        elif tok.type == TokenType.LANGLE:
            return self.parse_tuple()
        elif tok.type == TokenType.LBRACKET:
            return self.parse_vector()
        elif tok.type == TokenType.LBRACE:
            return self.parse_short_lambda()
        elif tok.type == TokenType.UNIT:
            self.advance()
            return UnitLit(line=tok.line, col=tok.col)
        elif tok.type == TokenType.LAMBDA:
            return self.parse_bare_lambda()
        elif tok.type in self.OPERATOR_TOKENS:
            # Bare operator as value: + → Symbol("+")
            self.advance()
            return Symbol(name=tok.value, line=tok.line, col=tok.col)
        else:
            return self.parse_atom()

    # Operator tokens that can appear as values
    OPERATOR_TOKENS = {
        TokenType.PLUS, TokenType.MINUS, TokenType.STAR,
        TokenType.SLASH, TokenType.PERCENT, TokenType.EQ,
        TokenType.NEQ, TokenType.LT, TokenType.GT,
        TokenType.LTE, TokenType.GTE, TokenType.AND,
        TokenType.OR, TokenType.NOT,
    }

    def parse_atom(self) -> ASTNode:
        """Parse an atomic expression"""
        tok = self.advance()

        if tok.type == TokenType.INT:
            return IntLit(value=int(tok.value), line=tok.line, col=tok.col)
        elif tok.type == TokenType.FLOAT:
            return FloatLit(value=float(tok.value), line=tok.line, col=tok.col)
        elif tok.type == TokenType.STRING:
            return StrLit(value=tok.value, line=tok.line, col=tok.col)
        elif tok.type == TokenType.TEMPLATE:
            # Template string: desugar $"Hi {name}!" to (fmt "Hi {}!" name)
            parts = tok.value.split('\x00')
            fmt_str = parts[0]
            expr_sources = parts[1:]
            # Parse each expression
            from .lexer import tokenize as lex_tokenize
            parsed_exprs = []
            for src in expr_sources:
                src = src.strip()
                if src:
                    expr_tokens = lex_tokenize(src)
                    sub_parser = Parser(expr_tokens)
                    parsed_exprs.append(sub_parser.parse_expr())
            args = [StrLit(value=fmt_str, line=tok.line, col=tok.col)] + parsed_exprs
            return Apply(
                func=Symbol(name='fmt', line=tok.line, col=tok.col),
                args=args,
                line=tok.line, col=tok.col
            )
        elif tok.type == TokenType.TRUE:
            return BoolLit(value=True, line=tok.line, col=tok.col)
        elif tok.type == TokenType.FALSE:
            return BoolLit(value=False, line=tok.line, col=tok.col)
        elif tok.type == TokenType.SYMBOL:
            return Symbol(name=tok.value, line=tok.line, col=tok.col)
        elif tok.type == TokenType.HASH_ID:
            return HashId(hash=tok.value, line=tok.line, col=tok.col)
        elif tok.type == TokenType.ADDR_REF:
            return AddrRef(addr=tok.value, line=tok.line, col=tok.col)
        elif tok.type == TokenType.UNDERSCORE:
            return Symbol(name='_', line=tok.line, col=tok.col)
        elif tok.type == TokenType.UNIT:
            return UnitLit(line=tok.line, col=tok.col)
        else:
            raise ParseError(
                f"Unexpected token: {tok.type.name} ({tok.value!r})",
                tok.line, tok.col
            )

    def parse_list(self) -> ASTNode:
        """Parse a parenthesized expression: (head args...)"""
        lparen = self.expect(TokenType.LPAREN)
        line, col = lparen.line, lparen.col

        if self.peek_type() == TokenType.RPAREN:
            # () is unit
            self.advance()
            return UnitLit(line=line, col=col)

        head = self.peek()

        # Special forms
        if head.type == TokenType.LAMBDA:
            return self.parse_lambda(line, col)
        elif head.type == TokenType.COND:
            return self.parse_cond(line, col)
        elif head.type == TokenType.MATCH:
            return self.parse_match(line, col)
        elif head.type == TokenType.SEQ:
            return self.parse_seq(line, col)
        elif head.type == TokenType.PARALLEL:
            return self.parse_parallel(line, col)
        elif head.type == TokenType.MODULE:
            return self.parse_module(line, col)
        elif head.type == TokenType.DEFINE:
            return self.parse_define(line, col)
        elif head.type == TokenType.RECURSE:
            return self.parse_recurse(line, col)
        elif head.type == TokenType.PROOF:
            return self.parse_proof(line, col)
        elif head.type == TokenType.MUTATE:
            return self.parse_mutate(line, col)
        elif head.type == TokenType.IMPORT:
            return self.parse_import(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'let':
            return self.parse_let(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'pipe':
            return self.parse_pipe(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'throw':
            return self.parse_throw(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'catch':
            return self.parse_catch(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'cond':
            return self.parse_cond_multi(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'defn':
            return self.parse_defn(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'for':
            return self.parse_for(line, col)
        elif head.type == TokenType.SYMBOL and head.value == 'py-import':
            return self.parse_py_import(line, col)
        elif head.type in (TokenType.PLUS, TokenType.MINUS, TokenType.STAR,
                           TokenType.SLASH, TokenType.PERCENT, TokenType.EQ,
                           TokenType.NEQ, TokenType.LT, TokenType.GT,
                           TokenType.LTE, TokenType.GTE, TokenType.AND,
                           TokenType.OR):
            return self.parse_binop(line, col)
        elif head.type == TokenType.NOT:
            return self.parse_unop(line, col)
        else:
            return self.parse_apply(line, col)

    def parse_lambda(self, line: int, col: int) -> Lambda:
        """(λ param1 param2 ... body) or (λ param1 param2 ... τ(...) body)"""
        self.advance()  # consume λ

        params = []
        type_ann = None

        # Collect parameters until we find the body
        # Params are hash IDs or symbols, body is everything else
        while self.peek_type() != TokenType.RPAREN:
            if self.peek_type() in (TokenType.HASH_ID, TokenType.SYMBOL):
                # Check if this is the last thing before ) — if so, it's the body
                saved_pos = self.pos
                candidate = self.parse_expr()

                if self.peek_type() == TokenType.RPAREN:
                    # This was the body
                    self.expect(TokenType.RPAREN)
                    return Lambda(params=params, body=candidate, type_ann=type_ann,
                                 line=line, col=col)
                else:
                    # This was a parameter, rewind isn't needed since it was simple
                    params.append(candidate)
            elif self.peek_type() == TokenType.TYPE:
                # τ type annotation
                self.advance()
                type_ann = self.parse_expr()
            else:
                # This must be the body
                body = self.parse_expr()
                # Check for bind after body
                if self.peek_type() == TokenType.ARROW:
                    body = self.parse_bind(body)
                self.expect(TokenType.RPAREN)
                return Lambda(params=params, body=body, type_ann=type_ann,
                             line=line, col=col)

        # If we get here, lambda with no body
        self.expect(TokenType.RPAREN)
        return Lambda(params=params, body=UnitLit(line=line, col=col),
                     type_ann=type_ann, line=line, col=col)

    def parse_bare_lambda(self) -> Lambda:
        """λ without parens — must be inside a list already"""
        tok = self.advance()  # consume λ
        # Read params then body
        params = []
        while self.peek_type() in (TokenType.HASH_ID, TokenType.SYMBOL):
            params.append(self.parse_atom())
        body = self.parse_expr()
        return Lambda(params=params, body=body, line=tok.line, col=tok.col)

    def parse_short_lambda(self) -> Lambda:
        """{expr} — short lambda with _ as implicit parameter
        {< _ 5} becomes (λ _ (< _ 5))
        {+ _ 1} becomes (λ _ (+ _ 1))
        {f _}   becomes (λ _ (f _))
        """
        tok = self.expect(TokenType.LBRACE)
        line, col = tok.line, tok.col

        # Parse contents like a list expression: head args...
        head = self.peek()

        if head.type in (TokenType.PLUS, TokenType.MINUS, TokenType.STAR,
                         TokenType.SLASH, TokenType.PERCENT, TokenType.EQ,
                         TokenType.NEQ, TokenType.LT, TokenType.GT,
                         TokenType.LTE, TokenType.GTE, TokenType.AND,
                         TokenType.OR):
            # Operator expression: {+ _ 1} → (+ _ 1)
            op = self.advance()
            args = []
            while self.peek_type() != TokenType.RBRACE:
                args.append(self.parse_expr())
            if len(args) == 2:
                body = BinOp(op=op.value, left=args[0], right=args[1],
                            line=op.line, col=op.col)
            elif len(args) == 1:
                body = UnOp(op=op.value, operand=args[0],
                           line=op.line, col=op.col)
            else:
                raise ParseError(f"Short lambda operator expects 1-2 args, got {len(args)}",
                               line, col)
        elif head.type == TokenType.NOT:
            op = self.advance()
            operand = self.parse_expr()
            body = UnOp(op=op.value, operand=operand, line=op.line, col=op.col)
        else:
            # Could be a function call {f _ args} or a single expression {_.1}
            func = self.parse_expr()
            if self.peek_type() == TokenType.RBRACE:
                # Single expression: {_.1} → (λ _ _.1)
                body = func
            else:
                # Function call: {f _ args} → (f _ args)
                args = []
                while self.peek_type() != TokenType.RBRACE:
                    args.append(self.parse_expr())
                body = Apply(func=func, args=args, line=line, col=col)

        self.expect(TokenType.RBRACE)

        # Collect all _ references to determine params
        # For now, always use single _ param
        param = Symbol(name='_', line=line, col=col)
        return Lambda(params=[param], body=body, line=line, col=col)

    def parse_cond(self, line: int, col: int) -> Cond:
        """(? condition then else)"""
        self.advance()  # consume ?
        condition = self.parse_expr()
        then_branch = self.parse_expr()
        else_branch = self.parse_expr() if self.peek_type() != TokenType.RPAREN else UnitLit(line=line, col=col)
        self.expect(TokenType.RPAREN)
        return Cond(condition=condition, then_branch=then_branch,
                   else_branch=else_branch, line=line, col=col)

    def parse_match(self, line: int, col: int) -> Match:
        """(~ expr (pattern result) (pattern result) ...)"""
        self.advance()  # consume ~
        expr = self.parse_expr()
        cases = []
        while self.peek_type() != TokenType.RPAREN:
            pattern = self.parse_expr()
            body = self.parse_expr()
            cases.append((pattern, body))
        self.expect(TokenType.RPAREN)
        return Match(expr=expr, cases=cases, line=line, col=col)

    def parse_cond_multi(self, line: int, col: int) -> Cond:
        """Multi-branch conditional. Desugars to nested Cond nodes.
        Flat syntax: (cond test1 val1 test2 val2 ... _ default)
        Last test can be 'else' or '_' as catch-all."""
        self.advance()  # consume 'cond'
        branches = []
        while self.peek_type() != TokenType.RPAREN:
            test = self.parse_expr()
            if self.peek_type() == TokenType.RPAREN:
                # Odd — last is default value with implicit _ test
                branches.append((Symbol(name='_', line=line, col=col), test))
                break
            val = self.parse_expr()
            branches.append((test, val))
        self.expect(TokenType.RPAREN)
        if not branches:
            return UnitLit(line=line, col=col)
        # Build nested Cond from right to left
        result = UnitLit(line=line, col=col)
        for test, val in reversed(branches):
            if isinstance(test, Symbol) and test.name in ('else', '_'):
                result = val
            else:
                result = Cond(condition=test, then_branch=val,
                             else_branch=result, line=line, col=col)
        return result

    def parse_seq(self, line: int, col: int) -> Seq:
        """(>> expr1 expr2 ...)"""
        self.advance()  # consume >>
        exprs = []
        while self.peek_type() != TokenType.RPAREN:
            expr = self.parse_expr()
            # Check for bind within seq
            if self.peek_type() == TokenType.ARROW:
                expr = self.parse_bind(expr)
            exprs.append(expr)
        self.expect(TokenType.RPAREN)
        return Seq(exprs=exprs, line=line, col=col)

    def parse_parallel(self, line: int, col: int) -> Parallel:
        """(∥ expr1 expr2 ...)"""
        self.advance()  # consume ∥
        exprs = []
        while self.peek_type() != TokenType.RPAREN:
            exprs.append(self.parse_expr())
        self.expect(TokenType.RPAREN)
        return Parallel(exprs=exprs, line=line, col=col)

    def parse_module(self, line: int, col: int) -> Module:
        """(□ name? ...definitions)"""
        self.advance()  # consume □
        name = ""
        if self.peek_type() == TokenType.SYMBOL:
            name = self.advance().value

        body = []
        while self.peek_type() != TokenType.RPAREN:
            expr = self.parse_expr()
            if self.peek_type() == TokenType.ARROW:
                expr = self.parse_bind(expr)
            body.append(expr)
        self.expect(TokenType.RPAREN)
        return Module(name=name, body=body, line=line, col=col)

    def parse_define(self, line: int, col: int) -> Define:
        """:= name value"""
        self.advance()  # consume :=
        name = self.parse_expr()
        value = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Define(name=name, value=value, line=line, col=col)

    def parse_recurse(self, line: int, col: int) -> Recurse:
        """(↺ name body)"""
        self.advance()  # consume ↺
        name = self.parse_expr()
        body = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Recurse(name=name, body=body, line=line, col=col)

    def parse_proof(self, line: int, col: int) -> Proof:
        """(⊢ assertion)"""
        self.advance()  # consume ⊢
        assertion = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Proof(assertion=assertion, line=line, col=col)

    def parse_mutate(self, line: int, col: int) -> Mutate:
        """(μ name value)"""
        self.advance()  # consume μ
        name = self.parse_expr()
        value = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Mutate(name=name, value=value, line=line, col=col)

    def parse_import(self, line: int, col: int) -> Import:
        """(⇐ @addr)"""
        self.advance()  # consume ⇐
        addr = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Import(addr=addr, line=line, col=col)

    def parse_let(self, line: int, col: int) -> Let:
        """(let x 1 y 2 ... body)"""
        self.advance()  # consume 'let'
        bindings = []
        # Parse pairs until we hit the last expression (the body)
        exprs = []
        while self.peek_type() != TokenType.RPAREN:
            exprs.append(self.parse_expr())
        self.expect(TokenType.RPAREN)
        # Last expr is body, everything before is name/value pairs
        if len(exprs) < 1:
            raise ParseError("let requires at least a body", line, col)
        body = exprs[-1]
        pairs = exprs[:-1]
        if len(pairs) % 2 != 0:
            raise ParseError("let requires even number of binding terms", line, col)
        for i in range(0, len(pairs), 2):
            bindings.append((pairs[i], pairs[i+1]))
        return Let(bindings=bindings, body=body, line=line, col=col)

    def parse_pipe(self, line: int, col: int) -> Pipe:
        """(pipe value f1 f2 f3)"""
        self.advance()  # consume 'pipe'
        value = self.parse_expr()
        funcs = []
        while self.peek_type() != TokenType.RPAREN:
            funcs.append(self.parse_expr())
        self.expect(TokenType.RPAREN)
        return Pipe(value=value, funcs=funcs, line=line, col=col)

    def parse_throw(self, line: int, col: int) -> Throw:
        """(throw expr)"""
        self.advance()  # consume 'throw'
        expr = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Throw(expr=expr, line=line, col=col)

    def parse_catch(self, line: int, col: int) -> Catch:
        """(catch expr handler)"""
        self.advance()  # consume 'catch'
        expr = self.parse_expr()
        handler = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return Catch(expr=expr, handler=handler, line=line, col=col)

    def parse_py_import(self, line: int, col: int) -> PyImportNode:
        """(py-import \"module\") or (py-import \"module\" alias)"""
        self.advance()  # consume 'py-import'
        module_name = self.parse_expr()
        alias = ""
        if self.peek_type() != TokenType.RPAREN:
            alias_node = self.parse_expr()
            if isinstance(alias_node, Symbol):
                alias = alias_node.name
        self.expect(TokenType.RPAREN)
        return PyImportNode(module_name=module_name, alias=alias, line=line, col=col)

    def parse_defn(self, line: int, col: int) -> Seq:
        """(defn name params... body) → (>> (def name (fn params... body)))
        Short form for defining named functions with auto-recursion."""
        self.advance()  # consume 'defn'
        name = self.parse_expr()  # function name
        # Collect params until we hit the body (last expr before RPAREN)
        exprs = []
        while self.peek_type() != TokenType.RPAREN:
            exprs.append(self.parse_expr())
        self.expect(TokenType.RPAREN)
        # Last expr is body, rest are params
        if not exprs:
            raise ParseError("defn requires at least a body")
        body = exprs[-1]
        params = exprs[:-1]
        # Build: (def name (fn params body))
        lam = Lambda(params=params, body=body, line=line, col=col)
        return Define(name=name, value=lam, line=line, col=col)

    def parse_for(self, line: int, col: int):
        """(for x coll body) → (map (fn x body) coll)
        (for x coll when pred body) → (map (fn x body) (filter (fn x pred) coll))
        Comprehension-like syntax."""
        self.advance()  # consume 'for'
        var = self.parse_expr()  # binding variable
        coll = self.parse_expr()  # collection
        # Check for 'when' filter
        when_pred = None
        if (self.peek_type() == TokenType.SYMBOL and
            self.tokens[self.pos].value == 'when'):
            self.advance()  # consume 'when'
            when_pred = self.parse_expr()
        body = self.parse_expr()
        self.expect(TokenType.RPAREN)
        # Build: (map (fn var body) coll) or (map (fn var body) (filter (fn var pred) coll))
        map_fn = Lambda(params=[var], body=body, line=line, col=col)
        if when_pred:
            filter_fn = Lambda(params=[var], body=when_pred, line=line, col=col)
            filtered = Apply(func=Symbol(name='filter', line=line, col=col),
                           args=[filter_fn, coll], line=line, col=col)
            return Apply(func=Symbol(name='map', line=line, col=col),
                       args=[map_fn, filtered], line=line, col=col)
        return Apply(func=Symbol(name='map', line=line, col=col),
                   args=[map_fn, coll], line=line, col=col)

    def parse_binop(self, line: int, col: int) -> BinOp:
        """(op left right)"""
        op = self.advance()
        left = self.parse_expr()
        right = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return BinOp(op=op.value, left=left, right=right, line=line, col=col)

    def parse_unop(self, line: int, col: int) -> UnOp:
        """(! operand)"""
        op = self.advance()
        operand = self.parse_expr()
        self.expect(TokenType.RPAREN)
        return UnOp(op=op.value, operand=operand, line=line, col=col)

    def parse_apply(self, line: int, col: int) -> Apply:
        """(func arg1 arg2 ...)"""
        func = self.parse_expr()
        args = []
        while self.peek_type() != TokenType.RPAREN:
            args.append(self.parse_expr())
        self.expect(TokenType.RPAREN)
        return Apply(func=func, args=args, line=line, col=col)

    def parse_bind(self, expr: ASTNode) -> Bind:
        """expr → target"""
        arrow = self.expect(TokenType.ARROW)
        target = self.parse_expr()
        return Bind(expr=expr, target=target, line=arrow.line, col=arrow.col)

    def parse_tuple(self) -> Tuple:
        """⟨expr1 expr2 ...⟩"""
        tok = self.expect(TokenType.LANGLE)
        elements = []
        while self.peek_type() != TokenType.RANGLE:
            elements.append(self.parse_expr())
        self.expect(TokenType.RANGLE)
        return Tuple(elements=elements, line=tok.line, col=tok.col)

    def parse_vector(self) -> Vector:
        """⟦expr1 expr2 ...⟧"""
        tok = self.expect(TokenType.LBRACKET)
        elements = []
        while self.peek_type() != TokenType.RBRACKET:
            elements.append(self.parse_expr())
        self.expect(TokenType.RBRACKET)
        return Vector(elements=elements, line=tok.line, col=tok.col)


def parse(tokens: List[Token]) -> Program:
    """Convenience function to parse a token list"""
    return Parser(tokens).parse()
