"""FSM extraction and state coverage analysis.

Automatically extracts state machines from waveform traces,
computes coverage, and exports as DOT/Mermaid.

Usage:
    fsm = t.extract_fsm(state_signal=t.tb.dut.state, clock=t.tb.clk)
    print(fsm.states)
    print(fsm.to_dot())
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Optional, TYPE_CHECKING
from collections import Counter

if TYPE_CHECKING:
    from ewal.trace import Trace
    from ewal.signal import SignalRef


@dataclass
class TransitionEdge:
    src: int
    dst: int
    count: int = 1
    first_seen: int = 0
    conditions: dict[str, Any] = field(default_factory=dict)


@dataclass
class FSMCoverage:
    state_pct: float
    transition_pct: float
    missing_states: set
    missing_transitions: set
    covered_states: set
    covered_transitions: set


class ExtractedFSM:
    """An extracted FSM from a waveform trace."""

    def __init__(self):
        self.states: set[int] = set()
        self.state_counts: Counter = Counter()
        self.transition_counts: Counter = Counter()
        self.transition_first_seen: dict[tuple, int] = {}
        self.transitions: list[TransitionEdge] = []
        self._observe_data: dict[tuple, dict[str, Any]] = {}

    def coverage(self, expected: dict) -> FSMCoverage:
        """Compute coverage against an expected FSM.

        Args:
            expected: Dict with "states" (set) and "transitions" (set of tuples).
        """
        exp_states = expected.get("states", set())
        exp_trans = expected.get("transitions", set())

        covered_s = self.states & exp_states
        covered_t = set(self.transition_counts.keys()) & exp_trans
        missing_s = exp_states - self.states
        missing_t = exp_trans - set(self.transition_counts.keys())

        s_pct = len(covered_s) / len(exp_states) if exp_states else 1.0
        t_pct = len(covered_t) / len(exp_trans) if exp_trans else 1.0

        return FSMCoverage(
            state_pct=s_pct,
            transition_pct=t_pct,
            missing_states=missing_s,
            missing_transitions=missing_t,
            covered_states=covered_s,
            covered_transitions=covered_t,
        )

    def illegal_transitions(self, expected_transitions: set) -> list[TransitionEdge]:
        """Find transitions not in the expected set."""
        illegal = []
        for (src, dst), count in self.transition_counts.items():
            if (src, dst) not in expected_transitions:
                illegal.append(TransitionEdge(
                    src=src, dst=dst, count=count,
                    first_seen=self.transition_first_seen.get((src, dst), 0),
                ))
        return illegal

    def to_dot(self, path: str = None) -> str:
        """Export as DOT/Graphviz format."""
        lines = ["digraph fsm {", "  rankdir=LR;", '  node [shape=circle];']

        for state in sorted(self.states):
            count = self.state_counts[state]
            lines.append(f'  s{state} [label="S{state}\\n({count} cycles)"];')

        for (src, dst), count in sorted(self.transition_counts.items()):
            lines.append(f'  s{src} -> s{dst} [label="{count}x"];')

        lines.append("}")
        result = "\n".join(lines)

        if path:
            with open(path, "w") as f:
                f.write(result)

        return result

    def to_mermaid(self) -> str:
        """Export as Mermaid stateDiagram."""
        lines = ["stateDiagram-v2"]

        for state in sorted(self.states):
            lines.append(f"    s{state}: S{state} ({self.state_counts[state]} cycles)")

        for (src, dst), count in sorted(self.transition_counts.items()):
            lines.append(f"    s{src} --> s{dst}: {count}x")

        return "\n".join(lines)


def extract_fsm(trace: "Trace", state_signal, clock=None,
                reset=None, observe: list = None) -> ExtractedFSM:
    """Extract an FSM from a state signal in the trace.

    Args:
        trace: The Trace object
        state_signal: SignalRef pointing to the state register
        clock: Optional clock (only sample at rising edges)
        reset: Optional reset signal (ignore states during reset)
        observe: Optional list of signals to record at each transition
    """
    from ewal.signal import SignalRef, _compile_expr

    fsm = ExtractedFSM()

    # Determine sample indices
    if clock is not None and isinstance(clock, SignalRef):
        compiled_rising = _compile_expr(clock.rising(), trace)
        sample_indices = [i for i in range(trace.max_index + 1)
                          if trace._core.eval_expr_bool(compiled_rising, i)]
    else:
        sample_indices = list(range(trace.max_index + 1))

    # Build reset mask
    reset_set = set()
    if reset is not None and isinstance(reset, SignalRef):
        compiled_reset = _compile_expr(reset._to_expr(), trace)
        reset_set = {i for i in sample_indices
                     if trace._core.eval_expr_bool(compiled_reset, i)}

    prev_state = None

    for idx in sample_indices:
        if idx in reset_set:
            prev_state = None
            continue

        old = trace._index
        trace._index = idx
        trace._virtual_cache.clear()
        current_state = state_signal.value
        trace._index = old

        if current_state is None:
            continue

        # Coerce to int
        try:
            current_state = int(current_state)
        except (TypeError, ValueError):
            continue

        fsm.states.add(current_state)
        fsm.state_counts[current_state] += 1

        if prev_state is not None and current_state != prev_state:
            key = (prev_state, current_state)
            fsm.transition_counts[key] += 1
            if key not in fsm.transition_first_seen:
                fsm.transition_first_seen[key] = idx

            # Capture observed signals
            if observe:
                conditions = {}
                old2 = trace._index
                trace._index = idx
                trace._virtual_cache.clear()
                for obs_sig in observe:
                    if isinstance(obs_sig, SignalRef):
                        name = obs_sig._path.split(".")[-1]
                        conditions[name] = obs_sig.value
                trace._index = old2
                fsm._observe_data[key] = conditions

        prev_state = current_state

    # Build transitions list
    for (src, dst), count in fsm.transition_counts.items():
        edge = TransitionEdge(
            src=src, dst=dst, count=count,
            first_seen=fsm.transition_first_seen.get((src, dst), 0),
            conditions=fsm._observe_data.get((src, dst), {}),
        )
        fsm.transitions.append(edge)

    return fsm


def _patch_trace_fsm():
    from ewal.trace import Trace

    def _extract_fsm(self, state_signal, clock=None, reset=None, observe=None):
        return extract_fsm(self, state_signal, clock=clock, reset=reset, observe=observe)

    Trace.extract_fsm = _extract_fsm


_patch_trace_fsm()
