"""Signal references and lazy expression trees for eWAL's embedded DSL."""

from __future__ import annotations
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ewal.trace import Trace


def _coerce(val) -> "SignalExpr":
    """Coerce Python primitives into SignalExpr constants."""
    if isinstance(val, SignalExpr):
        return val
    if isinstance(val, SignalRef):
        return val._to_expr()
    if isinstance(val, bool):
        return SignalExpr("const", int(val))
    if isinstance(val, (int, float, str)):
        return SignalExpr("const", val)
    raise TypeError(f"Cannot coerce {type(val)} into SignalExpr")


class SignalRef:
    """Reference to a signal in a trace.

    Attribute access chains build hierarchical signal names.
    Operators build lazy SignalExpr trees.
    """

    def __init__(self, trace: "Trace", path: str):
        object.__setattr__(self, "_trace", trace)
        object.__setattr__(self, "_path", path)

    def __getattr__(self, name: str) -> "SignalRef":
        if name.startswith("_"):
            raise AttributeError(name)
        path = self._path
        new_path = f"{path}.{name}" if path else name
        return SignalRef(self._trace, new_path)

    def __getitem__(self, offset):
        if isinstance(offset, int):
            # Relative evaluation: sig[-2] means sig@-2
            return SignalExpr("rel_eval", self._to_expr(), offset)
        if isinstance(offset, slice):
            # Bit slicing: sig[7:4]
            high = offset.start if offset.start is not None else 0
            low = offset.stop if offset.stop is not None else 0
            return SignalExpr("slice", self._to_expr(), high, low)
        raise TypeError(f"Invalid index type: {type(offset)}")

    @property
    def value(self):
        """Eagerly evaluate: get the signal's current value."""
        return self._trace._get_signal_value(self._path)

    @property
    def width(self) -> int:
        return self._trace._core.signal_width(self._path)

    # Operator overloads — all return SignalExpr (lazy)
    def __and__(self, other):
        return SignalExpr("and", self._to_expr(), _coerce(other))

    def __rand__(self, other):
        return SignalExpr("and", _coerce(other), self._to_expr())

    def __or__(self, other):
        return SignalExpr("or", self._to_expr(), _coerce(other))

    def __ror__(self, other):
        return SignalExpr("or", _coerce(other), self._to_expr())

    def __invert__(self):
        return SignalExpr("not", self._to_expr())

    def __gt__(self, other):
        return SignalExpr("gt", self._to_expr(), _coerce(other))

    def __lt__(self, other):
        return SignalExpr("lt", self._to_expr(), _coerce(other))

    def __ge__(self, other):
        return SignalExpr("gte", self._to_expr(), _coerce(other))

    def __le__(self, other):
        return SignalExpr("lte", self._to_expr(), _coerce(other))

    def __eq__(self, other):
        return SignalExpr("eq", self._to_expr(), _coerce(other))

    def __ne__(self, other):
        return SignalExpr("neq", self._to_expr(), _coerce(other))

    def __add__(self, other):
        return SignalExpr("add", self._to_expr(), _coerce(other))

    def __sub__(self, other):
        return SignalExpr("sub", self._to_expr(), _coerce(other))

    def __hash__(self):
        return hash(("SignalRef", self._path))

    def __repr__(self):
        return f"SignalRef({self._path!r})"

    def __bool__(self):
        return bool(self.value)

    # Convenience methods
    def rising(self):
        return ~self[-1] & self._to_expr()

    def falling(self):
        return self[-1] & ~self._to_expr()

    def stable(self):
        return SignalExpr("eq", self._to_expr(), self[-1])

    def unstable(self):
        return SignalExpr("neq", self._to_expr(), self[-1])

    def signed(self):
        return SignalExpr("signed", self._to_expr())

    def is_x(self):
        """Check if signal contains X values."""
        # This requires runtime check, not expressible as pure expr
        return SignalExpr("is_x", self._to_expr())

    def _to_expr(self) -> "SignalExpr":
        return SignalExpr("signal", self)

    def _compile(self, trace: "Trace"):
        """Compile this signal ref to a PyExpr."""
        from ewal._core import PyExpr
        sid = trace._core.signal_id(self._path)
        if sid is None:
            raise ValueError(f"Unknown signal: {self._path}")
        return PyExpr.signal(sid)


class SignalExpr:
    """Lazy expression tree node.

    Built by operator overloads on SignalRef. Not evaluated until
    passed to find(), whenever(), etc.
    """

    def __init__(self, op: str, *args):
        self.op = op
        self.args = args

    # Propagate operator overloads
    def __and__(self, other):
        return SignalExpr("and", self, _coerce(other))

    def __rand__(self, other):
        return SignalExpr("and", _coerce(other), self)

    def __or__(self, other):
        return SignalExpr("or", self, _coerce(other))

    def __ror__(self, other):
        return SignalExpr("or", _coerce(other), self)

    def __invert__(self):
        return SignalExpr("not", self)

    def __gt__(self, other):
        return SignalExpr("gt", self, _coerce(other))

    def __lt__(self, other):
        return SignalExpr("lt", self, _coerce(other))

    def __ge__(self, other):
        return SignalExpr("gte", self, _coerce(other))

    def __le__(self, other):
        return SignalExpr("lte", self, _coerce(other))

    def __eq__(self, other):
        return SignalExpr("eq", self, _coerce(other))

    def __ne__(self, other):
        return SignalExpr("neq", self, _coerce(other))

    def __add__(self, other):
        return SignalExpr("add", self, _coerce(other))

    def __sub__(self, other):
        return SignalExpr("sub", self, _coerce(other))

    def __hash__(self):
        return id(self)

    def __repr__(self):
        return f"SignalExpr({self.op}, {self.args!r})"

    def __bool__(self):
        raise TypeError(
            "Cannot use SignalExpr as boolean directly. "
            "Use .value for eager evaluation or pass to find()/whenever()."
        )

    def compile(self, trace: "Trace"):
        """Compile this expression tree into a Rust-side PyExpr."""
        return _compile_expr(self, trace)


def _compile_expr(expr, trace: "Trace"):
    """Recursively compile a SignalExpr/SignalRef tree to a PyExpr."""
    from ewal._core import PyExpr

    if isinstance(expr, SignalRef):
        return expr._compile(trace)

    if not isinstance(expr, SignalExpr):
        raise TypeError(f"Cannot compile {type(expr)}")

    op = expr.op
    args = expr.args

    if op == "signal":
        return args[0]._compile(trace)
    elif op == "const":
        val = args[0]
        if isinstance(val, bool):
            return PyExpr.const_int(int(val))
        elif isinstance(val, int):
            return PyExpr.const_int(val)
        elif isinstance(val, float):
            return PyExpr.const_float(val)
        elif isinstance(val, str):
            return PyExpr.const_str(val)
        else:
            return PyExpr.const_int(int(val))
    elif op == "rel_eval":
        inner = _compile_expr(args[0], trace)
        return inner.rel_eval(args[1])
    elif op == "not":
        return _compile_expr(args[0], trace).not_()
    elif op == "and":
        return _compile_expr(args[0], trace).and_(_compile_expr(args[1], trace))
    elif op == "or":
        return _compile_expr(args[0], trace).or_(_compile_expr(args[1], trace))
    elif op == "eq":
        return _compile_expr(args[0], trace).eq_(_compile_expr(args[1], trace))
    elif op == "neq":
        return _compile_expr(args[0], trace).neq(_compile_expr(args[1], trace))
    elif op == "gt":
        return _compile_expr(args[0], trace).gt_(_compile_expr(args[1], trace))
    elif op == "lt":
        return _compile_expr(args[0], trace).lt_(_compile_expr(args[1], trace))
    elif op == "gte":
        return _compile_expr(args[0], trace).gte(_compile_expr(args[1], trace))
    elif op == "lte":
        return _compile_expr(args[0], trace).lte(_compile_expr(args[1], trace))
    elif op == "add":
        return _compile_expr(args[0], trace).add(_compile_expr(args[1], trace))
    elif op == "sub":
        return _compile_expr(args[0], trace).sub(_compile_expr(args[1], trace))
    elif op == "slice":
        return _compile_expr(args[0], trace).slice(args[1], args[2])
    elif op == "signed":
        return _compile_expr(args[0], trace).signed()
    else:
        raise ValueError(f"Unknown expression op: {op}")
