from __future__ import annotations
from dataclasses import dataclass
from enum import Enum, auto
from typing import Callable, Dict, List, Optional

from .lexer import (
    Token,
    TokenType,
    TokenString,
    IntegerLiteral,
    IntegerBuiltInType,
)

from .interpreter_types import StackType
from .builtin import load_builtins


@dataclass
class StackEntry:
    type: StackType
    value: object  # Identifier | int | float | bool | TokenString | FunctionItem etc.


class FunctionType(Enum):
    TOKEN_STRING = auto()
    BUILTIN = auto()

type FunctionData = TokenString | Callable[["InterpreterState"], bool]

@dataclass
class FunctionItem:
    type: FunctionType
    token_string: Optional[TokenString] = None
    builtin: Optional[Callable[["InterpreterState"], bool]] = None

    @classmethod
    def from_item(cls, item: FunctionData) -> "FunctionItem":
        if isinstance(item, TokenString):
            return cls.from_token_string(item)
        else:
            return cls.from_builtin(item)

    @classmethod
    def from_token_string(cls, ts: TokenString) -> "FunctionItem":
        return cls(type=FunctionType.TOKEN_STRING, token_string=ts)

    @classmethod
    def from_builtin(cls, fn: Callable[["InterpreterState"], bool]) -> "FunctionItem":
        return cls(type=FunctionType.BUILTIN, builtin=fn)


class InterpreterState:
    """
    Interpreter state holds:
        - stack: list of StackEntry (top of stack is stack[-1])
        - functions: dict mapping names to FunctionItem
    """

    def __init__(self):
        self.stack: List[StackEntry] = []
        self.functions: Dict[str, FunctionItem] = {}

        if load_builtins is not None:
            ok = load_builtins(self)
            if not ok:
                raise RuntimeError("Failed to load builtins")

    # -------------------------
    # stack helpers
    # -------------------------
    def push(self, entry: StackEntry) -> None:
        self.stack.append(entry)

    def pop(self) -> Optional[StackEntry]:
        if not self.stack:
            return None
        return self.stack.pop()

    def top(self) -> Optional[StackEntry]:
        if not self.stack:
            return None
        return self.stack[-1]

    # -------------------------
    # function table
    # -------------------------
    def add_function(self, name: str, item: FunctionData) -> None:
        self.functions[name] = FunctionItem.from_item(item)

    def get_function(self, name: str) -> Optional[FunctionItem]:
        return self.functions.get(name)

    # -------------------------
    # execution primitives
    # -------------------------
    def push_token(self, token: Token) -> bool:
        """
        Push a token onto the stack. Returns True on success, False on failure.
        Mirrors the logic from the C implementation:
            - disallow TOKEN_STRING, TOKEN_ARRAY, TOKEN_TYPE_TUPLE
            - return True for TOKEN_EOF (no-op)
        """
        ttype = token.type

        if ttype == TokenType.EOF:
            return True

        # map token -> stack type + extract value
        if ttype == TokenType.IDENTIFIER:
            entry = StackEntry(StackType.IDENTIFIER, token.identifier)
            self.push(entry)
            return True

        if ttype == TokenType.INTEGER:
            ilit: IntegerLiteral = token.integer_literal  # type: ignore
            itype = ilit.type
            # choose stack type explicitly to mirror C widths
            if itype == IntegerBuiltInType.I64:
                st = StackType.I64
                val = int(ilit.value)
            elif itype == IntegerBuiltInType.I32:
                st = StackType.I32
                val = int(ilit.value) & 0xFFFFFFFF  # mimic truncation if desired
            elif itype == IntegerBuiltInType.I16:
                st = StackType.I16
                val = int(ilit.value) & 0xFFFF
            elif itype == IntegerBuiltInType.I8:
                st = StackType.I8
                val = int(ilit.value) & 0xFF
            elif itype == IntegerBuiltInType.U64:
                st = StackType.U64
                val = int(ilit.value)
            elif itype == IntegerBuiltInType.U32:
                st = StackType.U32
                val = int(ilit.value) & 0xFFFFFFFF
            elif itype == IntegerBuiltInType.U16:
                st = StackType.U16
                val = int(ilit.value) & 0xFFFF
            elif itype == IntegerBuiltInType.U8:
                st = StackType.U8
                val = int(ilit.value) & 0xFF
            else:
                return False

            self.push(StackEntry(st, val))
            return True

        if ttype == TokenType.FLOAT:
            self.push(StackEntry(StackType.FLOAT, token.float_literal))
            return True

        if ttype == TokenType.DOUBLE:
            self.push(StackEntry(StackType.DOUBLE, token.double_literal))
            return True

        if ttype == TokenType.CHARACTER:
            self.push(StackEntry(StackType.CHARACTER, token.character_literal))
            return True

        if ttype == TokenType.STRING:
            # C returned FALSE for TOKEN_STRING
            return False

        if ttype == TokenType.BOOLEAN:
            self.push(StackEntry(StackType.BOOLEAN, token.boolean_literal))
            return True

        if ttype == TokenType.ARRAY:
            # C returned FALSE for arrays
            return False

        if ttype == TokenType.TOKEN_STRING:
            # copy token string to mimic C copy semantics
            ts_copy = token.token_string.deep_copy()  # type: ignore
            self.push(StackEntry(StackType.TOKEN_STRING, ts_copy))
            return True

        if ttype == TokenType.TYPE_TUPLE:
            # C returned FALSE for type tuples
            return False

        # unknown token type
        return False

    def execute_func(self, key: str) -> bool:
        """
        Look up a function by name and execute it.
        Builtin functions are callables that accept InterpreterState and return bool.
        Token-string functions are executed via execute_token_string.
        """
        func = self.get_function(key)
        if func is None:
            return False

        if func.type == FunctionType.BUILTIN:
            if func.builtin is None:
                return False
            return func.builtin(self)

        if func.type == FunctionType.TOKEN_STRING:
            if func.token_string is None:
                return False
            return self.execute_token_string(func.token_string)

        return False

    def execute_token_string(self, token_string: TokenString) -> bool:
        """
        Execute each token in a TokenString.
        If token is an identifier and not `is_literal`, treat it as a function call.
        Otherwise push token onto the stack.
        """
        for tok in token_string.tokens:
            if tok.type == TokenType.IDENTIFIER and tok.identifier is not None and not tok.identifier.is_literal:
                rv = self.execute_func(tok.identifier.name)
            else:
                rv = self.push_token(tok)
            if not rv:
                return False
        return True

    def execute(self, token: Token) -> bool:
        """
        Execute a single Token (this corresponds to processing a single lexer result).
        If the token is an un-literal identifier -> function call, otherwise push it.
        """
        if token.type == TokenType.IDENTIFIER and token.identifier is not None and not token.identifier.is_literal:
            return self.execute_func(token.identifier.name)
        else:
            return self.push_token(token)
