"""Trace class — the central object in eWAL."""

from __future__ import annotations
from typing import Any, Callable, Optional

from ewal.signal import SignalRef, SignalExpr, _coerce, _compile_expr


class Trace:
    """Represents a loaded waveform trace.

    Supports signal access via attribute traversal (t.tb.dut.clk),
    find/whenever for analysis, virtual signals, scopes, and more.
    """

    def __init__(self, core, tid: str = "DEFAULT"):
        self._core = core
        self._tid = tid
        self._index = 0
        self._whenever_handlers: list[tuple] = []  # (condition, callback)
        self._begin_handlers: list[Callable] = []
        self._end_handlers: list[Callable] = []
        self._aliases: dict[str, str] = {}
        self._virtual_signals: dict[str, Any] = {}  # name -> (expr_or_callable)
        self._virtual_cache: dict[tuple[str, int], Any] = {}
        self._scope_stack: list[str] = []
        self._group_stack: list[str] = []
        self._variables: dict[str, Any] = {}

    def __getattr__(self, name: str) -> SignalRef:
        if name.startswith("_"):
            raise AttributeError(name)
        # Check aliases first
        if name in self._aliases:
            return SignalRef(self, self._aliases[name])
        # Check virtual signals
        if name in self._virtual_signals:
            return _VirtualSignalRef(self, name)
        return SignalRef(self, name)

    # ── Index management ──

    @property
    def index(self) -> int:
        return self._index

    @index.setter
    def index(self, val: int):
        self._index = val

    @property
    def max_index(self) -> int:
        return self._core.max_index()

    @property
    def ts(self) -> int:
        return self._core.timestamp(self._index)

    def step(self, n: int = 1):
        new_idx = self._index + n
        if 0 <= new_idx <= self.max_index:
            self._index = new_idx

    # ── Signal metadata ──

    @property
    def signals(self) -> list[str]:
        sigs = list(self._core.signal_names())
        sigs.extend(self._virtual_signals.keys())
        return sigs

    @property
    def scopes(self) -> list[str]:
        return self._core.scopes()

    @property
    def local_signals(self) -> list[str]:
        scope = self._current_scope()
        return self._core.local_signals(scope)

    @property
    def local_scopes(self) -> list[str]:
        scope = self._current_scope()
        return self._core.local_scopes(scope)

    def width(self, name: str) -> int:
        name = self._resolve_name(name)
        return self._core.signal_width(name)

    # ── Signal access ──

    def _get_signal_value(self, name: str, index: int | None = None):
        if index is None:
            index = self._index
        # Check virtual signals
        if name in self._virtual_signals:
            return self._eval_virtual(name, index)
        name = self._resolve_name(name)
        return self._core.signal_value(name, index)

    def _get_signal_value_at(self, name: str, index: int):
        if name in self._virtual_signals:
            return self._eval_virtual(name, index)
        name = self._resolve_name(name)
        return self._core.signal_value(name, index)

    def _resolve_name(self, name: str) -> str:
        if name in self._aliases:
            return self._aliases[name]
        return name

    # ── Aliases ──

    def alias(self, **kwargs):
        """Create signal aliases: t.alias(clk="tb.dut.clk")"""
        for short, full in kwargs.items():
            self._aliases[short] = full

    # ── Virtual signals ──

    def defsig(self, name: str, expr_or_callable):
        """Define a virtual signal.

        Can be a SignalExpr (lazy) or a callable (evaluated per timestep).
        """
        self._virtual_signals[name] = expr_or_callable

    def _eval_virtual(self, name: str, index: int):
        cache_key = (name, index)
        if cache_key in self._virtual_cache:
            return self._virtual_cache[cache_key]

        defn = self._virtual_signals[name]
        old_index = self._index
        self._index = index
        try:
            if callable(defn):
                result = defn()
            elif isinstance(defn, (SignalExpr, SignalRef)):
                # Compile and evaluate
                compiled = _compile_expr(defn, self)
                result = self._core.eval_expr(compiled, index)
            else:
                result = defn
        finally:
            self._index = old_index

        self._virtual_cache[cache_key] = result
        return result

    # ── Find ──

    def find(self, condition) -> list[int]:
        """Find all indices where condition is true.

        condition can be:
          - SignalExpr/SignalRef: compiled to Rust for fast evaluation
          - callable: called at every timestep (slower, Python loop)
          - True/bool: matches all/no timesteps
        """
        if condition is True:
            return list(range(self.max_index + 1))
        if condition is False:
            return []

        if isinstance(condition, (SignalExpr, SignalRef)):
            compiled = _compile_expr(condition, self)
            return self._core.find(compiled)

        if callable(condition):
            results = []
            old_index = self._index
            for idx in range(self.max_index + 1):
                self._index = idx
                self._virtual_cache.clear()
                if condition():
                    results.append(idx)
            self._index = old_index
            return results

        raise TypeError(f"Invalid condition type: {type(condition)}")

    def parallel_find(self, condition) -> list[int]:
        """Find all indices where condition is true, using rayon multi-core parallelism.

        Only works with SignalExpr/SignalRef conditions (not callables).
        Falls back to sequential find for callables.
        """
        if condition is True:
            return list(range(self.max_index + 1))
        if condition is False:
            return []

        if isinstance(condition, (SignalExpr, SignalRef)):
            compiled = _compile_expr(condition, self)
            return self._core.parallel_find(compiled)

        # Fallback for callables
        return self.find(condition)

    def parallel_count(self, condition) -> int:
        """Count matching indices using rayon multi-core parallelism."""
        if condition is True:
            return self.max_index + 1
        if condition is False:
            return 0

        if isinstance(condition, (SignalExpr, SignalRef)):
            compiled = _compile_expr(condition, self)
            return self._core.parallel_count(compiled)

        return len(self.find(condition))

    # ── Whenever ──

    def whenever(self, condition, callback=None):
        """Register a handler to run at each timestep matching condition.

        Can be used as decorator:
            @t.whenever(clk)
            def on_clk():
                ...

        Or with callback:
            t.whenever(clk, lambda: print(t.index))
        """
        if callback is not None:
            self._whenever_handlers.append((condition, callback))
            return callback

        def decorator(fn):
            self._whenever_handlers.append((condition, fn))
            return fn

        return decorator

    # ── Begin/End ──

    def begin(self, fn):
        """Decorator: register a BEGIN handler."""
        self._begin_handlers.append(fn)
        return fn

    def on_begin(self, fn):
        self._begin_handlers.append(fn)

    def end(self, fn):
        """Decorator: register an END handler."""
        self._end_handlers.append(fn)
        return fn

    def on_end(self, fn):
        self._end_handlers.append(fn)

    def on(self, condition):
        """Decorator: register condition-action handler (alias for whenever)."""
        return self.whenever(condition)

    # ── Run ──

    def run(self):
        """Execute all registered whenever handlers across the trace."""
        # Run BEGIN handlers
        for fn in self._begin_handlers:
            fn()

        # Pre-compile conditions where possible
        # Each entry is (compiled_or_None, callback, original_callable_or_None)
        compiled_handlers = []
        for cond, callback in self._whenever_handlers:
            if cond is True:
                compiled_handlers.append((None, callback, None))
            elif isinstance(cond, (SignalExpr, SignalRef)):
                compiled = _compile_expr(cond, self)
                compiled_handlers.append((compiled, callback, None))
            elif callable(cond):
                compiled_handlers.append((None, callback, cond))
            else:
                compiled_handlers.append((None, callback, None))

        # Main loop
        for idx in range(self.max_index + 1):
            self._index = idx
            self._virtual_cache.clear()
            for compiled_cond, callback, orig_callable in compiled_handlers:
                if orig_callable is not None:
                    # Callable condition — evaluate in Python
                    if orig_callable():
                        callback()
                elif compiled_cond is None:
                    # Always true
                    callback()
                else:
                    # Rust-compiled condition
                    if self._core.eval_expr_bool(compiled_cond, idx):
                        callback()

        # Run END handlers
        for fn in self._end_handlers:
            fn()

    # ── Temporal pattern matching ──

    def match(self, pattern, clock=None):
        """Match a temporal pattern against this trace.

        Args:
            pattern: A SequencePattern from ewal.temporal.seq()
            clock: Optional clock signal for cycle-based matching

        Returns:
            List of MatchResult with start/end indices and captures.
        """
        from ewal.temporal import match_pattern
        return match_pattern(self, pattern, clock=clock)

    # ── Scopes ──

    def scope(self, scope_name: str):
        """Context manager for scope."""
        return _ScopeContext(self, scope_name)

    def group(self, group_prefix: str):
        """Context manager for signal group."""
        return _GroupContext(self, group_prefix)

    def _current_scope(self) -> str:
        return ".".join(self._scope_stack) if self._scope_stack else ""

    def resolve(self, short_name: str) -> str:
        """Resolve a short signal name in the current scope."""
        scope = self._current_scope()
        result = self._core.resolve(scope, short_name)
        return result if result else short_name

    def gresolve(self, suffix: str):
        """Resolve a signal name relative to the current group."""
        if self._group_stack:
            return self._get_signal_value(f"{self._group_stack[-1]}{suffix}")
        return self._get_signal_value(suffix)

    def groups(self, *suffixes: str) -> list[str]:
        """Find all signal groups matching the given suffixes."""
        all_signals = set(self._core.signal_names())
        result = []
        # Find common prefixes where all suffixes exist
        prefixes = set()
        for sig in all_signals:
            for suffix in suffixes:
                if sig.endswith(suffix):
                    prefix = sig[: len(sig) - len(suffix)]
                    prefixes.add(prefix)

        for prefix in sorted(prefixes):
            if all(f"{prefix}{s}" in all_signals for s in suffixes):
                result.append(prefix)
        return result

    # ── Variables ──

    def var(self, name: str, default=0):
        """Get or create a trace-scoped variable."""
        if name not in self._variables:
            self._variables[name] = default
        return self._variables[name]

    # ── Aggregations ──

    def collect(self, signal_or_expr) -> list:
        """Collect all values of a signal across all timesteps."""
        if isinstance(signal_or_expr, SignalRef):
            path = signal_or_expr._path
            return [self._get_signal_value_at(path, i)
                    for i in range(self.max_index + 1)]
        return []

    def fold(self, signal_or_expr, init, fn):
        """Fold/reduce over signal values across all timesteps."""
        acc = init
        if isinstance(signal_or_expr, SignalRef):
            path = signal_or_expr._path
            for i in range(self.max_index + 1):
                val = self._get_signal_value_at(path, i)
                acc = fn(acc, val)
        return acc

    def count(self, condition) -> int:
        """Count how many timesteps match condition."""
        return len(self.find(condition))

    # ── Sample/Trim ──

    def sample_at(self, timestamps: list[int]):
        """Sample signal values at specific timestamps."""
        # Find indices for given timestamps
        indices = []
        for ts in timestamps:
            for i in range(self.max_index + 1):
                if self._core.timestamp(i) >= ts:
                    indices.append(i)
                    break
        return indices

    def trim(self, end: int):
        """Limit trace analysis to indices 0..end."""
        # This is handled by limiting the run loop
        # For now, store as attribute
        self._trim_end = end

    # ── REPL ──

    def repl(self):
        """Drop into an interactive REPL at the current trace state."""
        try:
            from IPython import embed
            embed()
        except ImportError:
            import code
            code.interact(local={**globals(), **locals(), "t": self})


class _VirtualSignalRef(SignalRef):
    """A SignalRef that points to a virtual signal."""

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

    @property
    def value(self):
        return self._trace._eval_virtual(self._path, self._trace._index)

    def __getitem__(self, offset):
        if isinstance(offset, int):
            # For virtual signals, relative access evaluates at offset index
            return _VirtualRelativeAccess(self._trace, self._path, offset)
        return super().__getitem__(offset)


class _VirtualRelativeAccess:
    """Lazy relative access to a virtual signal value."""

    def __init__(self, trace: Trace, name: str, offset: int):
        self._trace = trace
        self._name = name
        self._offset = offset

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

    def __eq__(self, other):
        if isinstance(other, _VirtualRelativeAccess):
            return self.value == other.value
        return self.value == other

    def __ne__(self, other):
        if isinstance(other, _VirtualRelativeAccess):
            return self.value != other.value
        return self.value != other

    def __int__(self):
        return int(self.value)

    @property
    def value(self):
        idx = self._trace._index + self._offset
        if idx < 0:
            return 0
        if idx > self._trace.max_index:
            return 0
        return self._trace._eval_virtual(self._name, idx)


class _ScopeContext:
    def __init__(self, trace: Trace, scope: str):
        self._trace = trace
        self._scope = scope

    def __enter__(self):
        self._trace._scope_stack.append(self._scope)
        return self._trace

    def __exit__(self, *args):
        self._trace._scope_stack.pop()


class _GroupContext:
    def __init__(self, trace: Trace, group: str):
        self._trace = trace
        self._group = group

    def __enter__(self):
        self._trace._group_stack.append(self._group)
        return self._trace

    def __exit__(self, *args):
        self._trace._group_stack.pop()
