"""Evaluate mathematical expressions on waveform data.

Supports expressions like "V(out) * I(R1)", "V(out) / V(in)",
"20*log10(abs(V(out)))", etc. Uses a safe recursive-descent parser
instead of Python's eval/exec.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from enum import Enum, auto

import numpy as np

from .raw_parser import RawFile

# -- Tokenizer ---------------------------------------------------------------


class _TokenType(Enum):
    NUMBER = auto()
    SIGNAL = auto()  # e.g., V(out), I(R1), time
    FUNC = auto()  # abs, sqrt, log10, dB
    PLUS = auto()
    MINUS = auto()
    STAR = auto()
    SLASH = auto()
    LPAREN = auto()
    RPAREN = auto()
    EOF = auto()


@dataclass
class _Token:
    type: _TokenType
    value: str


# Functions we support (case-insensitive lookup)
_FUNCTIONS = {"abs", "sqrt", "log10", "db"}

# Regex pieces for the tokenizer
_NUMBER_RE = re.compile(r"[0-9]*\.?[0-9]+(?:[eE][+-]?[0-9]+)?")

# Signal names: either V(...), I(...) style or bare identifiers like "time".
# Handles nested parens for things like V(N001) and I(R1).
_SIGNAL_RE = re.compile(r"[A-Za-z_]\w*\([^)]*\)")
_IDENT_RE = re.compile(r"[A-Za-z_]\w*")


def _tokenize(expr: str) -> list[_Token]:
    """Convert an expression string into a list of tokens."""
    tokens: list[_Token] = []
    i = 0
    s = expr.strip()

    while i < len(s):
        # Skip whitespace
        if s[i].isspace():
            i += 1
            continue

        # Operators / parens
        if s[i] == "+":
            tokens.append(_Token(_TokenType.PLUS, "+"))
            i += 1
        elif s[i] == "-":
            tokens.append(_Token(_TokenType.MINUS, "-"))
            i += 1
        elif s[i] == "*":
            tokens.append(_Token(_TokenType.STAR, "*"))
            i += 1
        elif s[i] == "/":
            tokens.append(_Token(_TokenType.SLASH, "/"))
            i += 1
        elif s[i] == "(":
            tokens.append(_Token(_TokenType.LPAREN, "("))
            i += 1
        elif s[i] == ")":
            tokens.append(_Token(_TokenType.RPAREN, ")"))
            i += 1

        # Number literal
        elif s[i].isdigit() or (s[i] == "." and i + 1 < len(s) and s[i + 1].isdigit()):
            m = _NUMBER_RE.match(s, i)
            if m:
                tokens.append(_Token(_TokenType.NUMBER, m.group()))
                i = m.end()
            else:
                raise ValueError(f"Invalid number at position {i}: {s[i:]!r}")

        # Identifier: could be a function, signal like V(out), or bare name
        elif s[i].isalpha() or s[i] == "_":
            # Try signal pattern first: name(...)
            m = _SIGNAL_RE.match(s, i)
            if m:
                full = m.group()
                # Check if the part before '(' is a known function
                paren_pos = full.index("(")
                prefix = full[:paren_pos].lower()
                if prefix in _FUNCTIONS:
                    # It's a function call -- emit FUNC token, then let
                    # the parser handle the parenthesized argument
                    tokens.append(_Token(_TokenType.FUNC, prefix))
                    i += paren_pos  # parser will see '(' next
                else:
                    # It's a signal name like V(out) or I(R1)
                    tokens.append(_Token(_TokenType.SIGNAL, full))
                    i = m.end()
            else:
                # Bare identifier (e.g., "time", or a function without parens)
                m = _IDENT_RE.match(s, i)
                if m:
                    name = m.group()
                    if name.lower() in _FUNCTIONS:
                        tokens.append(_Token(_TokenType.FUNC, name.lower()))
                    else:
                        tokens.append(_Token(_TokenType.SIGNAL, name))
                    i = m.end()
                else:
                    raise ValueError(f"Unexpected character at position {i}: {s[i:]!r}")
        else:
            raise ValueError(f"Unexpected character at position {i}: {s[i:]!r}")

    tokens.append(_Token(_TokenType.EOF, ""))
    return tokens


# -- Recursive-descent parser / evaluator ------------------------------------
#
# Grammar:
#   expr     -> term (('+' | '-') term)*
#   term     -> unary (('*' | '/') unary)*
#   unary    -> '-' unary | primary
#   primary  -> NUMBER | SIGNAL | FUNC '(' expr ')' | '(' expr ')'


class _Parser:
    """Recursive-descent expression evaluator over numpy arrays."""

    def __init__(self, tokens: list[_Token], variables: dict[str, np.ndarray]):
        self.tokens = tokens
        self.variables = variables
        self.pos = 0

    def _peek(self) -> _Token:
        return self.tokens[self.pos]

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

    def _expect(self, ttype: _TokenType) -> _Token:
        tok = self._advance()
        if tok.type != ttype:
            raise ValueError(f"Expected {ttype.name}, got {tok.type.name} ({tok.value!r})")
        return tok

    def parse(self) -> np.ndarray:
        result = self._expr()
        if self._peek().type != _TokenType.EOF:
            raise ValueError(f"Unexpected token after expression: {self._peek().value!r}")
        return np.real(result) if np.iscomplexobj(result) else result

    def _expr(self) -> np.ndarray:
        left = self._term()
        while self._peek().type in (_TokenType.PLUS, _TokenType.MINUS):
            op = self._advance()
            right = self._term()
            if op.type == _TokenType.PLUS:
                left = left + right
            else:
                left = left - right
        return left

    def _term(self) -> np.ndarray:
        left = self._unary()
        while self._peek().type in (_TokenType.STAR, _TokenType.SLASH):
            op = self._advance()
            right = self._unary()
            if op.type == _TokenType.STAR:
                left = left * right
            else:
                # Avoid division by zero -- replace zeros with tiny value
                safe = np.where(np.abs(right) < 1e-30, 1e-30, right)
                left = left / safe
        return left

    def _unary(self) -> np.ndarray:
        if self._peek().type == _TokenType.MINUS:
            self._advance()
            return -self._unary()
        return self._primary()

    def _primary(self) -> np.ndarray:
        tok = self._peek()

        if tok.type == _TokenType.NUMBER:
            self._advance()
            return np.float64(tok.value)

        if tok.type == _TokenType.SIGNAL:
            self._advance()
            return self._resolve_signal(tok.value)

        if tok.type == _TokenType.FUNC:
            self._advance()
            self._expect(_TokenType.LPAREN)
            arg = self._expr()
            self._expect(_TokenType.RPAREN)
            return self._apply_func(tok.value, arg)

        if tok.type == _TokenType.LPAREN:
            self._advance()
            result = self._expr()
            self._expect(_TokenType.RPAREN)
            return result

        raise ValueError(f"Unexpected token: {tok.type.name} ({tok.value!r})")

    def _resolve_signal(self, name: str) -> np.ndarray:
        """Look up a signal by name, trying exact match then case-insensitive."""
        if name in self.variables:
            return np.real(self.variables[name])

        name_lower = name.lower()
        for key, val in self.variables.items():
            if key.lower() == name_lower:
                return np.real(val)

        available = ", ".join(sorted(self.variables.keys()))
        raise ValueError(f"Unknown signal {name!r}. Available: {available}")

    @staticmethod
    def _apply_func(name: str, arg: np.ndarray) -> np.ndarray:
        """Apply a built-in function to a numpy array."""
        a = np.real(arg)

        if name == "abs":
            return np.abs(a)
        if name == "sqrt":
            return np.sqrt(np.maximum(a, 0.0))
        if name == "log10":
            return np.log10(np.maximum(np.abs(a), 1e-30))
        if name == "db":
            return 20.0 * np.log10(np.maximum(np.abs(a), 1e-30))

        raise ValueError(f"Unknown function: {name!r}")


# -- Public API ---------------------------------------------------------------


def evaluate_expression(expression: str, variables: dict[str, np.ndarray]) -> np.ndarray:
    """Evaluate a mathematical expression over waveform data.

    Uses a safe recursive-descent parser -- no eval() or exec().

    Args:
        expression: Math expression, e.g. "V(out) * I(R1)" or "dB(V(out))"
        variables: Dict mapping signal names to numpy arrays

    Returns:
        Resulting numpy array

    Raises:
        ValueError: On parse errors or unknown signals
    """
    tokens = _tokenize(expression)
    parser = _Parser(tokens, variables)
    return parser.parse()


class WaveformCalculator:
    """Evaluate expressions against all signals in a RawFile."""

    def __init__(self, raw_file: RawFile):
        self._raw = raw_file
        self._variables: dict[str, np.ndarray] = {}

        for var in raw_file.variables:
            self._variables[var.name] = raw_file.data[var.index]

    def calc(self, expression: str) -> np.ndarray:
        """Evaluate an expression against the loaded signals.

        Args:
            expression: Math expression referencing signal names from the raw file

        Returns:
            Resulting numpy array
        """
        return evaluate_expression(expression, self._variables)

    def available_signals(self) -> list[str]:
        """List all signal names available for expressions."""
        return sorted(self._variables.keys())
