"""Assertion monitoring / SVA-like property checking.

Provides post-hoc assertion checking on waveform traces with
structured pass/fail reporting and coverage tracking.

Usage:
    from ewal.assertions import always, never, eventually, implies, rose, fell

    @t.assert_property("req_gnt", clock=clk, disable=rst)
    def p1():
        return implies(t.tb.req, eventually(t.tb.gnt, within=10))

    results = t.check_assertions()
"""

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

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


# ── Assertion primitives ──

class AssertionExpr:
    """Base class for assertion expressions."""
    pass


class AlwaysExpr(AssertionExpr):
    def __init__(self, condition):
        self.condition = condition


class NeverExpr(AssertionExpr):
    def __init__(self, condition):
        self.condition = condition


class EventuallyExpr(AssertionExpr):
    def __init__(self, condition, within: int = 100):
        self.condition = condition
        self.within = within


class ImpliesExpr(AssertionExpr):
    def __init__(self, antecedent, consequent):
        self.antecedent = antecedent
        self.consequent = consequent


def always(condition) -> AlwaysExpr:
    """Condition must be true at every checked cycle."""
    return AlwaysExpr(condition)


def never(condition) -> NeverExpr:
    """Condition must never be true."""
    return NeverExpr(condition)


def eventually(condition, within: int = 100) -> EventuallyExpr:
    """Condition must become true within N cycles."""
    return EventuallyExpr(condition, within=within)


def implies(antecedent, consequent) -> ImpliesExpr:
    """When antecedent is true, consequent must hold."""
    return ImpliesExpr(antecedent, consequent)


def rose(sig) -> Any:
    """Rising edge of signal."""
    return sig.rising()


def fell(sig) -> Any:
    """Falling edge of signal."""
    return sig.falling()


def stable(sig) -> Any:
    """Signal unchanged."""
    return sig.stable()


def past(sig, n: int = 1) -> Any:
    """Value of signal N steps ago."""
    return sig[-n]


# ── Assertion results ──

@dataclass
class AssertionFailure:
    cycle: int
    message: str = ""
    signal_values: dict[str, Any] = field(default_factory=dict)


@dataclass
class AssertionResult:
    name: str
    passed: bool
    checked: int
    triggered: int  # how many times antecedent fired
    failures: list[AssertionFailure] = field(default_factory=list)


class AssertionReport(dict):
    """Collection of assertion results, keyed by name."""

    def coverage(self) -> dict:
        result = {}
        for name, res in self.items():
            hit_rate = res.triggered / res.checked if res.checked > 0 else 0
            result[name] = {
                "checked": res.checked,
                "triggered": res.triggered,
                "hit_rate": hit_rate,
            }
        return result

    def to_junit(self, path: str):
        """Export results as JUnit XML for CI integration."""
        lines = ['<?xml version="1.0" encoding="UTF-8"?>']
        total = len(self)
        failures = sum(1 for r in self.values() if not r.passed)
        lines.append(f'<testsuite name="ewal_assertions" tests="{total}" failures="{failures}">')
        for name, result in self.items():
            if result.passed:
                lines.append(f'  <testcase name="{name}" />')
            else:
                lines.append(f'  <testcase name="{name}">')
                for f in result.failures[:10]:
                    lines.append(f'    <failure message="cycle {f.cycle}: {f.message}" />')
                lines.append(f'  </testcase>')
        lines.append('</testsuite>')
        with open(path, "w") as fh:
            fh.write("\n".join(lines))

    def to_json(self, path: str):
        """Export results as JSON."""
        data = {}
        for name, result in self.items():
            data[name] = {
                "passed": result.passed,
                "checked": result.checked,
                "triggered": result.triggered,
                "failures": [
                    {"cycle": f.cycle, "message": f.message}
                    for f in result.failures
                ],
            }
        with open(path, "w") as fh:
            json.dump(data, fh, indent=2)


# ── Property registration and checking ──

@dataclass
class RegisteredProperty:
    name: str
    fn: Callable
    clock: Any = None
    disable: Any = None


def _eval_cond(trace, cond, index: int) -> bool:
    """Evaluate a condition/assertion expression at a given index."""
    from ewal.signal import SignalExpr, SignalRef, _compile_expr

    old = trace._index
    trace._index = index
    trace._virtual_cache.clear()

    try:
        if isinstance(cond, AlwaysExpr):
            return _eval_cond(trace, cond.condition, index)

        if isinstance(cond, NeverExpr):
            return not _eval_cond(trace, cond.condition, index)

        if isinstance(cond, EventuallyExpr):
            for offset in range(cond.within + 1):
                check_idx = index + offset
                if check_idx > trace.max_index:
                    break
                if _eval_cond(trace, cond.condition, check_idx):
                    return True
            return False

        if isinstance(cond, ImpliesExpr):
            if not _eval_cond(trace, cond.antecedent, index):
                return True  # vacuous truth
            return _eval_cond(trace, cond.consequent, index)

        if isinstance(cond, (SignalExpr, SignalRef)):
            compiled = _compile_expr(cond, trace)
            return trace._core.eval_expr_bool(compiled, index)

        if isinstance(cond, bool):
            return cond

        if callable(cond):
            return bool(cond())

        return bool(cond)
    finally:
        trace._index = old


def check_assertions(trace) -> AssertionReport:
    """Run all registered assertions on the trace."""
    from ewal.signal import SignalRef, _compile_expr

    report = AssertionReport()

    for prop in getattr(trace, "_properties", []):
        result = _check_property(trace, prop)
        report[prop.name] = result

    return report


def _check_property(trace, prop: RegisteredProperty) -> AssertionResult:
    """Check a single property across the trace."""
    from ewal.signal import SignalRef, _compile_expr

    # Determine check indices (rising edges of clock, or all)
    if prop.clock is not None and isinstance(prop.clock, SignalRef):
        compiled_rising = _compile_expr(prop.clock.rising(), trace)
        check_indices = [i for i in range(trace.max_index + 1)
                         if trace._core.eval_expr_bool(compiled_rising, i)]
    else:
        check_indices = list(range(trace.max_index + 1))

    # Build disable mask
    disable_indices = set()
    if prop.disable is not None:
        from ewal.signal import SignalRef as SR, SignalExpr as SE
        if isinstance(prop.disable, (SR, SE)):
            compiled_dis = _compile_expr(prop.disable, trace)
            disable_indices = {i for i in check_indices
                               if trace._core.eval_expr_bool(compiled_dis, i)}

    failures = []
    checked = 0
    triggered = 0

    for idx in check_indices:
        if idx in disable_indices:
            continue
        checked += 1

        # Evaluate the property function to get the assertion expression
        old = trace._index
        trace._index = idx
        trace._virtual_cache.clear()
        assertion_expr = prop.fn()
        trace._index = old

        if assertion_expr is None:
            continue

        # Track triggers (for implies: when antecedent is true)
        if isinstance(assertion_expr, ImpliesExpr):
            if _eval_cond(trace, assertion_expr.antecedent, idx):
                triggered += 1
        else:
            triggered += 1

        # Evaluate
        if not _eval_cond(trace, assertion_expr, idx):
            failures.append(AssertionFailure(
                cycle=idx,
                message=f"Property '{prop.name}' failed at index {idx}",
            ))

    return AssertionResult(
        name=prop.name,
        passed=len(failures) == 0,
        checked=checked,
        triggered=triggered,
        failures=failures,
    )


def _patch_trace_assertions():
    """Add assertion methods to Trace."""
    from ewal.trace import Trace

    def assert_property(self, name: str, clock=None, disable=None):
        """Decorator: register an assertion property."""
        if not hasattr(self, "_properties"):
            self._properties = []

        def decorator(fn):
            self._properties.append(RegisteredProperty(
                name=name, fn=fn, clock=clock, disable=disable,
            ))
            return fn
        return decorator

    Trace.assert_property = assert_property

    def _check_assertions(self):
        return check_assertions(self)
    Trace.check_assertions = _check_assertions


_patch_trace_assertions()
