"""Waveform annotation and marker export.

Export analysis results as annotations for GTKWave, Verdi, or generic JSON.
Also supports exporting annotated VCD files with marker signals.

Usage:
    markers = ewal.Markers()
    markers.add(index, "OVERFLOW", color="red", signal="tb.dut.counter")
    markers.to_gtkwave_savefile("annotations.gtkw")
    markers.to_json("annotations.json")
"""

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

if TYPE_CHECKING:
    from ewal.trace import Trace


@dataclass
class Marker:
    index: int
    label: str
    color: str = "yellow"
    signal: str = ""
    timestamp: int = 0


class Markers:
    """Collection of waveform annotations/markers."""

    def __init__(self):
        self._markers: list[Marker] = []

    def add(self, index: int, label: str, color: str = "yellow",
            signal: str = "", timestamp: int = 0):
        """Add a marker at a given index."""
        self._markers.append(Marker(
            index=index, label=label, color=color,
            signal=signal, timestamp=timestamp,
        ))

    def __len__(self):
        return len(self._markers)

    def __iter__(self):
        return iter(self._markers)

    def as_signal(self, label: str, max_index: int = 0) -> dict[int, int]:
        """Convert markers with a given label into a signal (index -> 1/0)."""
        if max_index == 0 and self._markers:
            max_index = max(m.index for m in self._markers)
        result = {}
        for m in self._markers:
            if m.label == label:
                result[m.index] = 1
        return result

    # ── GTKWave export ──

    def to_gtkwave_savefile(self, path: str, signals: list[str] = None):
        """Export as GTKWave .gtkw save file with named markers.

        GTKWave uses named markers A-Z at specific timestamps.
        """
        lines = [
            "[*]",
            "[*] eWAL Generated Annotations",
            "[*]",
        ]

        # Add signals to display
        if signals:
            for sig in signals:
                lines.append(sig)

        # GTKWave supports markers A-Z (26 max)
        marker_names = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        for i, m in enumerate(self._markers[:26]):
            ts = m.timestamp if m.timestamp else m.index
            lines.append(f"[marker] {marker_names[i]} {ts}")

        # Named markers via time-names
        lines.append("[*]")
        for m in self._markers:
            ts = m.timestamp if m.timestamp else m.index
            lines.append(f"[time_name] {ts} {m.label}")

        with open(path, "w") as f:
            f.write("\n".join(lines) + "\n")

    # ── Verdi export ──

    def to_verdi_fsdb_annotation(self, path: str):
        """Export as Verdi-compatible annotation script."""
        lines = [
            "# eWAL Generated Annotations for Verdi",
            "# Source this in Verdi's Tcl console",
        ]
        for i, m in enumerate(self._markers):
            ts = m.timestamp if m.timestamp else m.index
            color_map = {
                "red": "red", "green": "green", "blue": "blue",
                "yellow": "yellow", "orange": "orange",
            }
            vc = color_map.get(m.color, "yellow")
            lines.append(f'debAddMarker {ts} -name "{m.label}" -color {vc}')
            if m.signal:
                lines.append(f'# Signal: {m.signal}')

        with open(path, "w") as f:
            f.write("\n".join(lines) + "\n")

    # ── JSON export ──

    def to_json(self, path: str):
        """Export markers as JSON."""
        data = []
        for m in self._markers:
            data.append({
                "index": m.index,
                "label": m.label,
                "color": m.color,
                "signal": m.signal,
                "timestamp": m.timestamp,
            })
        with open(path, "w") as f:
            json.dump(data, f, indent=2)


def export_vcd(trace: "Trace", path: str, extra_signals: dict[str, dict[int, int]] = None):
    """Export a VCD file with optional extra marker signals.

    Args:
        trace: Source trace
        path: Output VCD path
        extra_signals: Dict of signal_name -> {index: value} for marker signals
    """
    if extra_signals is None:
        extra_signals = {}

    max_idx = trace.max_index
    signal_names = list(trace._core.signal_names())

    # Assign VCD IDs
    vcd_ids = {}
    next_id = ord("!")
    for sig in signal_names:
        vcd_ids[sig] = chr(next_id)
        next_id += 1
    for sig in extra_signals:
        vcd_ids[sig] = chr(next_id)
        next_id += 1

    with open(path, "w") as f:
        # Header
        f.write("$date\n  eWAL generated\n$end\n")
        f.write("$version\n  eWAL 0.1.0\n$end\n")
        f.write("$timescale\n  1ns\n$end\n")

        # Signals
        f.write("$scope module top $end\n")
        for sig in signal_names:
            w = trace._core.signal_width(sig)
            f.write(f"$var wire {w} {vcd_ids[sig]} {sig} $end\n")
        for sig in extra_signals:
            f.write(f"$var wire 1 {vcd_ids[sig]} {sig} $end\n")
        f.write("$upscope $end\n")
        f.write("$enddefinitions $end\n")

        # Value changes
        for idx in range(max_idx + 1):
            ts = trace._core.timestamp(idx)
            f.write(f"#{ts}\n")

            for sig in signal_names:
                val = trace._get_signal_value(sig, idx)
                w = trace._core.signal_width(sig)
                vid = vcd_ids[sig]
                if w == 1:
                    f.write(f"{val if val is not None else 'x'}{vid}\n")
                else:
                    if val is not None:
                        f.write(f"b{val:0{w}b} {vid}\n")
                    else:
                        f.write(f"bx {vid}\n")

            for sig, vals in extra_signals.items():
                vid = vcd_ids[sig]
                v = vals.get(idx, 0)
                f.write(f"{v}{vid}\n")


def _patch_trace_markers():
    from ewal.trace import Trace

    def _export_vcd(self, path, extra_signals=None):
        return export_vcd(self, path, extra_signals)

    Trace.export_vcd = _export_vcd


_patch_trace_markers()
