"""WAL compatibility shim — drop-in replacement for wal.core.Wal.

Provides the Wal class that accepts eval_str() calls but uses
eWAL's Rust backend for trace operations.
"""

from __future__ import annotations
from typing import Optional

import ewal
from ewal.trace import Trace


class Wal:
    """Drop-in replacement for wal.core.Wal.

    Usage:
        from ewal.compat import Wal

        w = Wal()
        w.load("counter.fst")
        # Use eWAL trace directly:
        t = w.trace("DEFAULT")
        indices = t.find(t.tb.dut.clk & t.tb.dut.overflow)
    """

    def __init__(self):
        self._traces: dict[str, Trace] = {}

    def load(self, path: str, tid: str = "DEFAULT", **kwargs):
        """Load a trace file."""
        t = ewal.load(path, tid=tid, **kwargs)
        self._traces[tid] = t

    def trace(self, tid: str = "DEFAULT") -> Trace:
        """Get the eWAL Trace object for a loaded trace."""
        if tid not in self._traces:
            raise KeyError(f"No trace loaded with id '{tid}'")
        return self._traces[tid]

    def step(self, steps: int = 1, tid: Optional[str] = None):
        """Step one or all traces."""
        if tid:
            self._traces[tid].step(steps)
        else:
            for t in self._traces.values():
                t.step(steps)

    def eval_str(self, txt: str, **args):
        """Parse and evaluate a WAL expression string.

        NOTE: This is a limited compatibility shim. It supports basic
        WAL operations but not the full WAL language. For full eWAL
        functionality, use the Trace object directly.
        """
        raise NotImplementedError(
            "eval_str() is not supported in eWAL. "
            "Use the Trace object directly: t = w.trace(); t.find(...)"
        )
