"""
TENGWAR Binary AST Protocol (TBAP)

The killer AI feature: instead of generating text like
    (filter {> _ 5} (map {* _ 2} data))
and parsing it, an AI emits compact binary opcodes that map
directly to AST nodes. 

Benefits:
    - ZERO syntax errors possible (structurally valid by construction)
    - ZERO parsing overhead (binary → AST in O(n), no lexer/parser)
    - ~60% fewer bytes than text representation
    - Can be embedded in tool-call JSON as base64
    - Streamable: AI can emit opcodes as it generates

Wire format:
    Each node is: [opcode: u8] [payload...]
    Opcodes map 1:1 to AST node types.
    Strings are length-prefixed (u16 length + utf8 bytes).
    Integers are varint-encoded.
    Lists use a count prefix.

This is something NO other language offers to AI.
"""

import struct
from typing import List, Tuple
from .ast_nodes import *
from .interpreter import (
    Interpreter, TengwarValue, TengwarInt, TengwarFloat, TengwarStr,
    TengwarBool, TengwarUnit, TengwarVector, TengwarTuple, TengwarDict,
    TengwarClosure, TengwarBuiltin
)

# === OPCODES ===
# One byte per AST node type

class Op:
    # Literals (0x00-0x0F)
    INT         = 0x01  # + varint
    FLOAT       = 0x02  # + f64
    STRING      = 0x03  # + u16 len + utf8
    TRUE        = 0x04
    FALSE       = 0x05
    UNIT        = 0x06

    # Identifiers (0x10-0x1F)
    SYMBOL      = 0x10  # + u16 len + utf8
    HASH_ID     = 0x11  # + u16 len + utf8
    ADDR_REF    = 0x12  # + u16 len + utf8

    # Core forms (0x20-0x3F)
    LAMBDA      = 0x20  # + u8 param_count + params + body
    APPLY       = 0x21  # + func + u8 arg_count + args
    COND        = 0x22  # + condition + then + else
    MATCH       = 0x23  # + expr + u8 case_count + (pattern, body)*
    SEQ         = 0x24  # + u8 count + exprs
    PARALLEL    = 0x25  # + u8 count + exprs
    BIND        = 0x26  # + expr + target
    DEFINE      = 0x27  # + name + value
    RECURSE     = 0x28  # + name + body
    MODULE      = 0x29  # + u16 name_len + name + u8 count + body
    LET         = 0x2A  # + u8 binding_count + (name, value)* + body
    PIPE        = 0x2B  # + value + u8 func_count + funcs
    THROW       = 0x2C  # + expr
    CATCH       = 0x2D  # + expr + handler
    PY_IMPORT   = 0x2E  # + module_name + u16 alias_len + alias
    PROOF       = 0x2F  # + assertion
    MUTATE      = 0x30  # + name + value

    # Collections (0x40-0x4F)
    TUPLE       = 0x40  # + u8 count + elements
    VECTOR      = 0x41  # + u16 count + elements

    # Operators (0x50-0x6F)
    BINOP       = 0x50  # + u8 op_id + left + right
    UNOP        = 0x51  # + u8 op_id + operand

    # Short lambda (0x70-0x7F)
    SHORT_LAMBDA = 0x70  # + body (implicit _ param)

    # Program
    PROGRAM     = 0xFF  # + u16 count + exprs


# Operator IDs for binary encoding
OP_IDS = {
    '+': 0, '-': 1, '*': 2, '/': 3, '%': 4,
    '=': 5, '!=': 6, '<': 7, '>': 8, '<=': 9, '>=': 10,
    '&': 11, '|': 12, '!': 13,
}
ID_OPS = {v: k for k, v in OP_IDS.items()}


# === ENCODER: AST → Binary ===

class Encoder:
    """Encode a Tengwar AST to binary format"""

    def __init__(self):
        self.buf = bytearray()

    def _write_u8(self, v: int):
        self.buf.append(v & 0xFF)

    def _write_u16(self, v: int):
        self.buf.extend(struct.pack('>H', v))

    def _write_varint(self, v: int):
        """Variable-length integer encoding"""
        negative = v < 0
        if negative:
            v = -v
            # Set high bit of first byte to indicate negative
        while v >= 0x80:
            self.buf.append((v & 0x7F) | 0x80)
            v >>= 7
        if negative:
            self.buf.append(v | 0x40)  # second-highest bit = negative
            self.buf.append(0x00)  # terminator
        else:
            self.buf.append(v)

    def _write_f64(self, v: float):
        self.buf.extend(struct.pack('>d', v))

    def _write_string(self, s: str):
        encoded = s.encode('utf-8')
        self._write_u16(len(encoded))
        self.buf.extend(encoded)

    def encode(self, node: ASTNode) -> bytes:
        """Encode an AST node to binary"""
        self._encode_node(node)
        return bytes(self.buf)

    def _encode_node(self, node):
        if isinstance(node, Program):
            self._write_u8(Op.PROGRAM)
            self._write_u16(len(node.body))
            for expr in node.body:
                self._encode_node(expr)

        elif isinstance(node, IntLit):
            self._write_u8(Op.INT)
            self._write_varint(node.value)

        elif isinstance(node, FloatLit):
            self._write_u8(Op.FLOAT)
            self._write_f64(node.value)

        elif isinstance(node, StrLit):
            self._write_u8(Op.STRING)
            self._write_string(node.value)

        elif isinstance(node, BoolLit):
            self._write_u8(Op.TRUE if node.value else Op.FALSE)

        elif isinstance(node, UnitLit):
            self._write_u8(Op.UNIT)

        elif isinstance(node, Symbol):
            self._write_u8(Op.SYMBOL)
            self._write_string(node.name)

        elif isinstance(node, HashId):
            self._write_u8(Op.HASH_ID)
            self._write_string(node.hash)

        elif isinstance(node, AddrRef):
            self._write_u8(Op.ADDR_REF)
            self._write_string(node.addr)

        elif isinstance(node, Lambda):
            self._write_u8(Op.LAMBDA)
            self._write_u8(len(node.params))
            for p in node.params:
                self._encode_node(p)
            self._encode_node(node.body)

        elif isinstance(node, Apply):
            self._write_u8(Op.APPLY)
            self._encode_node(node.func)
            self._write_u8(len(node.args))
            for a in node.args:
                self._encode_node(a)

        elif isinstance(node, Cond):
            self._write_u8(Op.COND)
            self._encode_node(node.condition)
            self._encode_node(node.then_branch)
            self._encode_node(node.else_branch)

        elif isinstance(node, Match):
            self._write_u8(Op.MATCH)
            self._encode_node(node.expr)
            self._write_u8(len(node.cases))
            for pattern, body in node.cases:
                self._encode_node(pattern)
                self._encode_node(body)

        elif isinstance(node, Seq):
            self._write_u8(Op.SEQ)
            self._write_u8(len(node.exprs))
            for e in node.exprs:
                self._encode_node(e)

        elif isinstance(node, Parallel):
            self._write_u8(Op.PARALLEL)
            self._write_u8(len(node.exprs))
            for e in node.exprs:
                self._encode_node(e)

        elif isinstance(node, Bind):
            self._write_u8(Op.BIND)
            self._encode_node(node.expr)
            self._encode_node(node.target)

        elif isinstance(node, Define):
            self._write_u8(Op.DEFINE)
            self._encode_node(node.name)
            self._encode_node(node.value)

        elif isinstance(node, Recurse):
            self._write_u8(Op.RECURSE)
            self._encode_node(node.name)
            self._encode_node(node.body)

        elif isinstance(node, Module):
            self._write_u8(Op.MODULE)
            self._write_string(node.name)
            self._write_u8(len(node.body))
            for e in node.body:
                self._encode_node(e)

        elif isinstance(node, Let):
            self._write_u8(Op.LET)
            self._write_u8(len(node.bindings))
            for name_node, val_node in node.bindings:
                self._encode_node(name_node)
                self._encode_node(val_node)
            self._encode_node(node.body)

        elif isinstance(node, Pipe):
            self._write_u8(Op.PIPE)
            self._encode_node(node.value)
            self._write_u8(len(node.funcs))
            for f in node.funcs:
                self._encode_node(f)

        elif isinstance(node, Throw):
            self._write_u8(Op.THROW)
            self._encode_node(node.expr)

        elif isinstance(node, Catch):
            self._write_u8(Op.CATCH)
            self._encode_node(node.expr)
            self._encode_node(node.handler)

        elif isinstance(node, PyImportNode):
            self._write_u8(Op.PY_IMPORT)
            self._encode_node(node.module_name)
            self._write_string(node.alias)

        elif isinstance(node, Proof):
            self._write_u8(Op.PROOF)
            self._encode_node(node.assertion)

        elif isinstance(node, Mutate):
            self._write_u8(Op.MUTATE)
            self._encode_node(node.name)
            self._encode_node(node.value)

        elif isinstance(node, Tuple):
            self._write_u8(Op.TUPLE)
            self._write_u8(len(node.elements))
            for e in node.elements:
                self._encode_node(e)

        elif isinstance(node, Vector):
            self._write_u8(Op.VECTOR)
            self._write_u16(len(node.elements))
            for e in node.elements:
                self._encode_node(e)

        elif isinstance(node, BinOp):
            self._write_u8(Op.BINOP)
            self._write_u8(OP_IDS.get(node.op, 0))
            self._encode_node(node.left)
            self._encode_node(node.right)

        elif isinstance(node, UnOp):
            self._write_u8(Op.UNOP)
            self._write_u8(OP_IDS.get(node.op, 0))
            self._encode_node(node.operand)

        else:
            raise ValueError(f"Cannot encode: {type(node).__name__}")


# === DECODER: Binary → AST ===

class Decoder:
    """Decode binary format back to Tengwar AST"""

    def __init__(self, data: bytes):
        self.data = data
        self.pos = 0

    def _read_u8(self) -> int:
        v = self.data[self.pos]
        self.pos += 1
        return v

    def _read_u16(self) -> int:
        v = struct.unpack('>H', self.data[self.pos:self.pos+2])[0]
        self.pos += 2
        return v

    def _read_varint(self) -> int:
        result = 0
        shift = 0
        while True:
            b = self.data[self.pos]
            self.pos += 1
            if b & 0x80:
                result |= (b & 0x7F) << shift
                shift += 7
            else:
                if b & 0x40:
                    # Negative: next byte is terminator
                    result |= (b & 0x3F) << shift
                    self.pos += 1  # skip terminator
                    return -result
                result |= b << shift
                return result

    def _read_f64(self) -> float:
        v = struct.unpack('>d', self.data[self.pos:self.pos+8])[0]
        self.pos += 8
        return v

    def _read_string(self) -> str:
        length = self._read_u16()
        s = self.data[self.pos:self.pos+length].decode('utf-8')
        self.pos += length
        return s

    def decode(self) -> ASTNode:
        """Decode binary data to AST"""
        return self._decode_node()

    def _decode_node(self) -> ASTNode:
        op = self._read_u8()

        if op == Op.PROGRAM:
            count = self._read_u16()
            body = [self._decode_node() for _ in range(count)]
            return Program(body=body)

        elif op == Op.INT:
            return IntLit(value=self._read_varint())

        elif op == Op.FLOAT:
            return FloatLit(value=self._read_f64())

        elif op == Op.STRING:
            return StrLit(value=self._read_string())

        elif op == Op.TRUE:
            return BoolLit(value=True)

        elif op == Op.FALSE:
            return BoolLit(value=False)

        elif op == Op.UNIT:
            return UnitLit()

        elif op == Op.SYMBOL:
            return Symbol(name=self._read_string())

        elif op == Op.HASH_ID:
            return HashId(hash=self._read_string())

        elif op == Op.ADDR_REF:
            return AddrRef(addr=self._read_string())

        elif op == Op.LAMBDA:
            param_count = self._read_u8()
            params = [self._decode_node() for _ in range(param_count)]
            body = self._decode_node()
            return Lambda(params=params, body=body)

        elif op == Op.APPLY:
            func = self._decode_node()
            arg_count = self._read_u8()
            args = [self._decode_node() for _ in range(arg_count)]
            return Apply(func=func, args=args)

        elif op == Op.COND:
            condition = self._decode_node()
            then_branch = self._decode_node()
            else_branch = self._decode_node()
            return Cond(condition=condition, then_branch=then_branch, else_branch=else_branch)

        elif op == Op.MATCH:
            expr = self._decode_node()
            case_count = self._read_u8()
            cases = []
            for _ in range(case_count):
                pattern = self._decode_node()
                body = self._decode_node()
                cases.append((pattern, body))
            return Match(expr=expr, cases=cases)

        elif op == Op.SEQ:
            count = self._read_u8()
            exprs = [self._decode_node() for _ in range(count)]
            return Seq(exprs=exprs)

        elif op == Op.PARALLEL:
            count = self._read_u8()
            exprs = [self._decode_node() for _ in range(count)]
            return Parallel(exprs=exprs)

        elif op == Op.BIND:
            expr = self._decode_node()
            target = self._decode_node()
            return Bind(expr=expr, target=target)

        elif op == Op.DEFINE:
            name = self._decode_node()
            value = self._decode_node()
            return Define(name=name, value=value)

        elif op == Op.RECURSE:
            name = self._decode_node()
            body = self._decode_node()
            return Recurse(name=name, body=body)

        elif op == Op.MODULE:
            name = self._read_string()
            count = self._read_u8()
            body = [self._decode_node() for _ in range(count)]
            return Module(name=name, body=body)

        elif op == Op.LET:
            binding_count = self._read_u8()
            bindings = []
            for _ in range(binding_count):
                name = self._decode_node()
                value = self._decode_node()
                bindings.append((name, value))
            body = self._decode_node()
            return Let(bindings=bindings, body=body)

        elif op == Op.PIPE:
            value = self._decode_node()
            func_count = self._read_u8()
            funcs = [self._decode_node() for _ in range(func_count)]
            return Pipe(value=value, funcs=funcs)

        elif op == Op.THROW:
            expr = self._decode_node()
            return Throw(expr=expr)

        elif op == Op.CATCH:
            expr = self._decode_node()
            handler = self._decode_node()
            return Catch(expr=expr, handler=handler)

        elif op == Op.PY_IMPORT:
            module_name = self._decode_node()
            alias = self._read_string()
            return PyImportNode(module_name=module_name, alias=alias)

        elif op == Op.PROOF:
            assertion = self._decode_node()
            return Proof(assertion=assertion)

        elif op == Op.MUTATE:
            name = self._decode_node()
            value = self._decode_node()
            return Mutate(name=name, value=value)

        elif op == Op.TUPLE:
            count = self._read_u8()
            elements = [self._decode_node() for _ in range(count)]
            return Tuple(elements=elements)

        elif op == Op.VECTOR:
            count = self._read_u16()
            elements = [self._decode_node() for _ in range(count)]
            return Vector(elements=elements)

        elif op == Op.BINOP:
            op_id = self._read_u8()
            left = self._decode_node()
            right = self._decode_node()
            return BinOp(op=ID_OPS.get(op_id, '+'), left=left, right=right)

        elif op == Op.UNOP:
            op_id = self._read_u8()
            operand = self._decode_node()
            return UnOp(op=ID_OPS.get(op_id, '!'), operand=operand)

        elif op == Op.SHORT_LAMBDA:
            body = self._decode_node()
            param = Symbol(name='_')
            return Lambda(params=[param], body=body)

        else:
            raise ValueError(f"Unknown opcode: 0x{op:02x} at position {self.pos-1}")


# === CONVENIENCE FUNCTIONS ===

def encode(source: str) -> bytes:
    """Parse Tengwar source and encode to binary"""
    from .lexer import tokenize
    from .parser import parse
    tokens = tokenize(source)
    ast = parse(tokens)
    return Encoder().encode(ast)


def decode(data: bytes) -> ASTNode:
    """Decode binary data to AST"""
    return Decoder(data).decode()


def run_binary(data: bytes, interpreter: Interpreter = None) -> TengwarValue:
    """Decode binary AST and execute it"""
    ast = decode(data)
    if interpreter is None:
        interpreter = Interpreter()
    return interpreter.eval(ast)


def encode_b64(source: str) -> str:
    """Encode to base64 (for embedding in JSON/tool calls)"""
    import base64
    return base64.b64encode(encode(source)).decode('ascii')


def run_b64(b64_data: str, interpreter: Interpreter = None) -> TengwarValue:
    """Decode base64 binary AST and execute"""
    import base64
    data = base64.b64decode(b64_data)
    return run_binary(data, interpreter)


# === AST BUILDER (for AI to construct ASTs programmatically) ===

class ASTBuilder:
    """
    Fluent API for building Tengwar ASTs without text.
    
    An AI can use this to construct programs directly:
    
        b = ASTBuilder()
        program = b.program(
            b.define("double", b.lam(["x"], b.binop("*", b.sym("x"), b.int(2)))),
            b.apply(b.sym("double"), b.int(21))
        )
        result = b.run(program)  # → 42
    """

    def int(self, v: int) -> IntLit:
        return IntLit(value=v)

    def float(self, v: float) -> FloatLit:
        return FloatLit(value=v)

    def str(self, v: str) -> StrLit:
        return StrLit(value=v)

    def bool(self, v: bool) -> BoolLit:
        return BoolLit(value=v)

    def unit(self) -> UnitLit:
        return UnitLit()

    def sym(self, name: str) -> Symbol:
        return Symbol(name=name)

    def lam(self, params: list, body: ASTNode) -> Lambda:
        return Lambda(
            params=[Symbol(name=p) if isinstance(p, __builtins__['str'] if isinstance(__builtins__, dict) else str) else p for p in params],
            body=body
        )

    def apply(self, func: ASTNode, *args: ASTNode) -> Apply:
        return Apply(func=func, args=list(args))

    def binop(self, op: str, left: ASTNode, right: ASTNode) -> BinOp:
        return BinOp(op=op, left=left, right=right)

    def unop(self, op: str, operand: ASTNode) -> UnOp:
        return UnOp(op=op, operand=operand)

    def cond(self, condition: ASTNode, then_b: ASTNode, else_b: ASTNode) -> Cond:
        return Cond(condition=condition, then_branch=then_b, else_branch=else_b)

    def seq(self, *exprs: ASTNode) -> Seq:
        return Seq(exprs=list(exprs))

    def bind(self, expr: ASTNode, target: ASTNode) -> Bind:
        return Bind(expr=expr, target=target)

    def define(self, name: str, value: ASTNode) -> Define:
        return Define(name=Symbol(name=name), value=value)

    def recurse(self, name: str, body: ASTNode) -> Recurse:
        return Recurse(name=Symbol(name=name), body=body)

    def let(self, bindings: list, body: ASTNode) -> Let:
        """bindings: [("name", value_node), ...]"""
        return Let(
            bindings=[(Symbol(name=n), v) for n, v in bindings],
            body=body
        )

    def pipe(self, value: ASTNode, *funcs: ASTNode) -> Pipe:
        return Pipe(value=value, funcs=list(funcs))

    def vec(self, *elements: ASTNode) -> Vector:
        return Vector(elements=list(elements))

    def tup(self, *elements: ASTNode) -> Tuple:
        return Tuple(elements=list(elements))

    def match(self, expr: ASTNode, *cases) -> Match:
        """cases: [(pattern, body), ...]"""
        return Match(expr=expr, cases=list(cases))

    def throw(self, expr: ASTNode) -> Throw:
        return Throw(expr=expr)

    def catch(self, expr: ASTNode, handler: ASTNode) -> Catch:
        return Catch(expr=expr, handler=handler)

    def program(self, *exprs: ASTNode) -> Program:
        return Program(body=list(exprs))

    def encode(self, node: ASTNode) -> bytes:
        """Encode AST to binary"""
        return Encoder().encode(node)

    def run(self, node: ASTNode, interpreter: Interpreter = None) -> TengwarValue:
        """Execute an AST directly"""
        if interpreter is None:
            interpreter = Interpreter()
        return interpreter.eval(node)
