"""Waveform diffing — structural comparison between traces.

Usage:
    diff = ewal.diff(golden, failing)
    print(diff.first_divergence)
    print(diff.divergent_signals)
"""

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

if TYPE_CHECKING:
    from ewal.trace import Trace


@dataclass
class Divergence:
    signal: str
    index: int


@dataclass
class RootCause:
    signal: str
    diverged_at: int
    leads_by: int


@dataclass
class WaveformDiff:
    """Result of comparing two traces."""
    first_divergence: Optional[Divergence]
    divergent_signals: list[str]
    identical_signals: list[str]
    divergence_points: dict[str, int]  # signal -> first diverging index
    _trace_a: Any = None
    _trace_b: Any = None

    def summary(self) -> str:
        lines = []
        if self.first_divergence:
            lines.append(
                f"First divergence: signal '{self.first_divergence.signal}' "
                f"at index {self.first_divergence.index}"
            )
        else:
            lines.append("Traces are identical.")
        lines.append(f"Divergent signals: {len(self.divergent_signals)}")
        lines.append(f"Identical signals: {len(self.identical_signals)}")
        return "\n".join(lines)

    def to_dataframe(self, signals: list[str] = None, window: int = 10):
        """Export divergence details as a DataFrame.

        Produces per-index rows around each divergence point:
            index | signal | golden_value | failing_value | first_diff

        Args:
            signals: Signals to include (default: divergent signals)
            window: Number of indices to show around each divergence
        """
        import pandas as pd

        if signals is None:
            signals = self.divergent_signals

        if not self._trace_a or not self._trace_b:
            # Fallback: summary-only format
            rows = []
            for sig in signals:
                div_idx = self.divergence_points.get(sig)
                rows.append({
                    "signal": sig,
                    "first_divergence_index": div_idx,
                    "divergent": div_idx is not None,
                })
            return pd.DataFrame(rows)

        rows = []
        for sig in signals:
            div_idx = self.divergence_points.get(sig)
            if div_idx is None:
                continue
            start = max(0, div_idx)
            end = min(div_idx + window,
                      min(self._trace_a.max_index, self._trace_b.max_index) + 1)
            for idx in range(start, end):
                val_a = self._trace_a._get_signal_value(sig, idx)
                val_b = self._trace_b._get_signal_value(sig, idx)
                rows.append({
                    "index": idx,
                    "signal": sig,
                    "golden_value": val_a,
                    "failing_value": val_b,
                    "first_diff": idx == div_idx,
                })
        return pd.DataFrame(rows)

    def root_causes(self, outputs: list[str], inputs: list[str],
                    window: int = 50) -> list[RootCause]:
        """Trace back from output divergence to find input root causes.

        Looks for input signals that diverged within `window` cycles
        before each output divergence.
        """
        results = []
        for out_sig in outputs:
            out_div = self.divergence_points.get(out_sig)
            if out_div is None:
                continue
            for in_sig in inputs:
                in_div = self.divergence_points.get(in_sig)
                if in_div is not None and 0 < out_div - in_div <= window:
                    results.append(RootCause(
                        signal=in_sig,
                        diverged_at=in_div,
                        leads_by=out_div - in_div,
                    ))
        results.sort(key=lambda r: r.leads_by)
        return results


def diff(trace_a: "Trace", trace_b: "Trace",
         signals: list[str] = None,
         tolerance: dict[str, float] = None) -> WaveformDiff:
    """Compare two traces and find divergence points.

    Args:
        trace_a: First trace (e.g., golden)
        trace_b: Second trace (e.g., failing)
        signals: Optional subset of signals to compare
        tolerance: Dict of signal -> tolerance for floating-point comparison

    Returns:
        WaveformDiff with structured results.
    """
    if tolerance is None:
        tolerance = {}

    # Determine signals to compare
    if signals is None:
        sigs_a = set(trace_a._core.signal_names())
        sigs_b = set(trace_b._core.signal_names())
        signals = sorted(sigs_a & sigs_b)

    max_idx = min(trace_a.max_index, trace_b.max_index)

    divergent = []
    identical = []
    divergence_points = {}
    first_div = None

    for sig in signals:
        sig_first_div = None
        tol = tolerance.get(sig, 0)

        for idx in range(max_idx + 1):
            val_a = trace_a._get_signal_value(sig, idx)
            val_b = trace_b._get_signal_value(sig, idx)

            if tol > 0:
                # Floating-point comparison
                try:
                    if abs(float(val_a or 0) - float(val_b or 0)) > tol:
                        sig_first_div = idx
                        break
                except (TypeError, ValueError):
                    if val_a != val_b:
                        sig_first_div = idx
                        break
            else:
                if val_a != val_b:
                    sig_first_div = idx
                    break

        if sig_first_div is not None:
            divergent.append(sig)
            divergence_points[sig] = sig_first_div
            if first_div is None or sig_first_div < first_div.index:
                first_div = Divergence(signal=sig, index=sig_first_div)
        else:
            identical.append(sig)

    return WaveformDiff(
        first_divergence=first_div,
        divergent_signals=divergent,
        identical_signals=identical,
        divergence_points=divergence_points,
        _trace_a=trace_a,
        _trace_b=trace_b,
    )
