"""eWAL — Embedded Waveform Analysis Library.

A Pythonic embedded DSL for waveform analysis with a Rust core.
"""

from __future__ import annotations
from typing import Optional

from ewal.trace import Trace
from ewal.signal import SignalRef, SignalExpr

__version__ = "0.1.0"

# Global trace registry for multi-trace support
_traces: dict[str, Trace] = {}


def load(path: str, tid: str = "DEFAULT", **kwargs) -> Trace:
    """Load a waveform trace file. Format auto-detected from extension.

    Args:
        path: Path to trace file (.vcd, .fst, .csv)
        tid: Trace ID for multi-trace support

    Returns:
        Trace object for analysis.
    """
    from ewal._core import PyTrace
    from ewal.plugins import get_reader_for_extension

    ext = path.rsplit(".", 1)[-1].lower() if "." in path else ""
    dot_ext = f".{ext}"

    # Check plugin registry first
    reader_cls = get_reader_for_extension(dot_ext)
    if reader_cls is not None:
        core = reader_cls(path)
        t = Trace(core, tid=tid)
        _traces[tid] = t
        return t

    if ext == "vcd":
        core = PyTrace.from_vcd(path)
    elif ext == "csv":
        raise NotImplementedError("CSV format support coming soon")
    else:
        # Try VCD as fallback
        core = PyTrace.from_vcd(path)

    t = Trace(core, tid=tid)
    _traces[tid] = t
    return t


def virtual_trace(max_index: int = 30, time_step: int = 1) -> Trace:
    """Create a virtual trace with no backing file.

    Virtual signals are defined with defsig() and computed on demand.
    """
    core = _VirtualTraceCore(max_index, time_step)
    return Trace(core, tid="VIRTUAL")


class _VirtualTraceCore:
    """Pure-Python core for virtual traces (no Rust backing)."""

    def __init__(self, max_idx: int, time_step: int):
        self._max_index = max_idx
        self._time_step = time_step
        self._signal_names: list[str] = []
        self._scopes: list[str] = []

    def signal_names(self):
        return self._signal_names

    def scopes(self):
        return self._scopes

    def max_index(self):
        return self._max_index

    def timestamp(self, index):
        return index * self._time_step

    def signal_value(self, name, index):
        return None

    def signal_value_raw(self, name, index):
        return 0

    def signal_width(self, name):
        return 1

    def signal_id(self, name):
        try:
            return self._signal_names.index(name)
        except ValueError:
            return None

    def signal_values_range(self, name, start, end):
        return [None] * (end - start)

    def resolve(self, scope, short_name):
        return short_name

    def local_signals(self, scope):
        return []

    def local_scopes(self, scope):
        return []

    def find(self, expr):
        return []

    def eval_expr(self, expr, index):
        return None

    def eval_expr_bool(self, expr, index):
        return False


# ── Multi-trace support ──

def step_all(n: int = 1):
    """Step all loaded traces by n."""
    for t in _traces.values():
        t.step(n)


def whenever_all(condition):
    """Register a handler that runs on all traces in lockstep."""
    def decorator(fn):
        # Store for run_all
        if not hasattr(whenever_all, "_handlers"):
            whenever_all._handlers = []
        whenever_all._handlers.append((condition, fn))
        return fn
    return decorator


def run_all():
    """Run all registered multi-trace handlers."""
    handlers = getattr(whenever_all, "_handlers", [])
    if not handlers or not _traces:
        return

    traces = list(_traces.values())
    max_idx = min(t.max_index for t in traces)

    for idx in range(max_idx + 1):
        for t in traces:
            t._index = idx
            t._virtual_cache.clear()
        for cond, fn in handlers:
            if cond is True:
                fn()
            elif callable(cond):
                if cond():
                    fn()


# ── Import feature modules to activate monkey-patches ──
import ewal.export       # adds to_numpy, to_dataframe, to_csv, to_parquet, to_arrow
import ewal.assertions   # adds assert_property, check_assertions
import ewal.fsm          # adds extract_fsm
import ewal.markers      # adds export_vcd
import ewal.notebook     # adds show, plot, heatmap
import ewal.query        # adds query
import ewal.debug        # adds debug

# ── Waveform diffing ──
from ewal.diffing import diff, WaveformDiff

# ── Streaming ──
from ewal.streaming import stream, live, parallel_find

# ── Markers ──
from ewal.markers import Markers

# ── Correlation ──
from ewal.correlation import correlate, find_predictors

# ── Plugin system ──
from ewal.plugins import register_format, TraceReader

# ── Stdlib convenience functions ──

def rising(sig: SignalRef) -> SignalExpr:
    """Detect rising edge of signal."""
    return sig.rising()


def falling(sig: SignalRef) -> SignalExpr:
    """Detect falling edge of signal."""
    return sig.falling()


def stable(sig: SignalRef) -> SignalExpr:
    """Detect when signal is unchanged from previous step."""
    return sig.stable()


def unstable(sig: SignalRef) -> SignalExpr:
    """Detect when signal changed from previous step."""
    return sig.unstable()
