"""Pipeline visualization and stage tracking.

Tracks instructions through pipeline stages, detects stalls/flushes/bubbles,
computes CPI, and exports visualizations.

Usage:
    from ewal.pipeline import Pipeline, Stage

    pipe = Pipeline(clock=t.tb.clk, stages=[
        Stage("fetch", valid=t.tb.dut.fetch_valid, pc=t.tb.dut.fetch_pc),
        ...
    ])
    analysis = pipe.analyze(t)
"""

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

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


@dataclass
class Stage:
    """A pipeline stage definition."""
    name: str
    valid: Any = None    # SignalRef: stage has valid instruction
    pc: Any = None       # SignalRef: program counter / instruction ID
    stall: Any = None    # SignalRef: stage is stalled
    flush: Any = None    # SignalRef: stage is flushed
    retire: Any = None   # SignalRef: instruction retires (writeback only)


@dataclass
class InstructionRecord:
    """Tracking record for a single instruction through the pipeline."""
    pc: int
    stage_cycles: dict[str, int] = field(default_factory=dict)  # stage -> entry cycle
    latency: int = 0
    stalled: bool = False
    flushed: bool = False


@dataclass
class PipelineAnalysis:
    """Results of pipeline analysis."""
    total_cycles: int = 0
    retired: int = 0
    stall_cycles: int = 0
    flush_count: int = 0
    bubble_cycles: int = 0
    cpi: float = 0.0
    stalls_by_stage: dict[str, int] = field(default_factory=dict)
    flushes_by_stage: dict[str, int] = field(default_factory=dict)
    instructions: list[InstructionRecord] = field(default_factory=list)

    def to_html(self, path: str):
        """Export pipeline diagram as HTML."""
        lines = [
            "<!DOCTYPE html><html><head>",
            "<style>",
            "table { border-collapse: collapse; font-family: monospace; font-size: 11px; }",
            "td { border: 1px solid #ccc; padding: 2px 4px; text-align: center; min-width: 30px; }",
            ".valid { background: #4CAF50; color: white; }",
            ".stall { background: #FF9800; color: white; }",
            ".flush { background: #F44336; color: white; }",
            ".bubble { background: #E0E0E0; }",
            "th { background: #333; color: white; padding: 4px; }",
            "</style></head><body>",
            f"<h2>Pipeline Diagram — {self.total_cycles} cycles, "
            f"{self.retired} retired, CPI={self.cpi:.2f}</h2>",
            "<table><tr><th>Cycle</th>",
        ]
        stage_names = sorted(set(
            s for inst in self.instructions for s in inst.stage_cycles
        ))
        for s in stage_names:
            lines.append(f"<th>{s}</th>")
        lines.append("</tr>")

        # Build cycle -> stage -> info map
        cycle_map: dict[int, dict[str, str]] = {}
        for inst in self.instructions:
            for stage, cycle in inst.stage_cycles.items():
                if cycle not in cycle_map:
                    cycle_map[cycle] = {}
                cls = "valid"
                if inst.stalled:
                    cls = "stall"
                if inst.flushed:
                    cls = "flush"
                cycle_map[cycle][stage] = (f"0x{inst.pc:x}", cls)

        for cycle in sorted(cycle_map.keys())[:200]:  # limit to 200 rows
            lines.append(f"<tr><td>{cycle}</td>")
            for s in stage_names:
                if s in cycle_map.get(cycle, {}):
                    text, cls = cycle_map[cycle][s]
                    lines.append(f'<td class="{cls}">{text}</td>')
                else:
                    lines.append('<td class="bubble"></td>')
            lines.append("</tr>")

        lines.append("</table></body></html>")
        with open(path, "w") as f:
            f.write("\n".join(lines))

    def to_svg(self, path: str):
        """Export pipeline diagram as SVG."""
        stage_names = sorted(set(
            s for inst in self.instructions for s in inst.stage_cycles
        ))
        n_stages = len(stage_names)
        n_cycles = min(self.total_cycles, 200)
        col_w = 60
        row_h = 20
        header_h = 30
        w = col_w * (n_stages + 1) + 20
        h = header_h + row_h * n_cycles + 20

        lines = [
            f'<svg xmlns="http://www.w3.org/2000/svg" width="{w}" height="{h}">',
            '<style>text { font-family: monospace; font-size: 10px; }</style>',
        ]

        # Headers
        lines.append(f'<text x="10" y="20" font-weight="bold">Cycle</text>')
        for i, s in enumerate(stage_names):
            lines.append(f'<text x="{col_w * (i + 1) + 10}" y="20" font-weight="bold">{s}</text>')

        # Build map
        cycle_map: dict[int, dict[str, tuple]] = {}
        for inst in self.instructions:
            for stage, cycle in inst.stage_cycles.items():
                if cycle not in cycle_map:
                    cycle_map[cycle] = {}
                color = "#4CAF50"
                if inst.stalled:
                    color = "#FF9800"
                if inst.flushed:
                    color = "#F44336"
                cycle_map[cycle][stage] = (f"0x{inst.pc:x}", color)

        sorted_cycles = sorted(cycle_map.keys())[:n_cycles]
        for row, cycle in enumerate(sorted_cycles):
            y = header_h + row * row_h
            lines.append(f'<text x="10" y="{y + 15}">{cycle}</text>')
            for i, s in enumerate(stage_names):
                x = col_w * (i + 1) + 5
                if s in cycle_map.get(cycle, {}):
                    text, color = cycle_map[cycle][s]
                    lines.append(f'<rect x="{x}" y="{y}" width="{col_w - 5}" height="{row_h - 2}" fill="{color}" rx="2"/>')
                    lines.append(f'<text x="{x + 3}" y="{y + 13}" fill="white">{text}</text>')

        lines.append("</svg>")
        with open(path, "w") as f:
            f.write("\n".join(lines))

    def stall_heatmap(self, path: str):
        """Export stall heatmap as a simple text-based visualization."""
        lines = ["=== Stall Heatmap ==="]
        for stage, count in sorted(self.stalls_by_stage.items()):
            bar = "#" * min(count, 80)
            lines.append(f"  {stage:>12}: {bar} ({count})")
        result = "\n".join(lines)
        with open(path, "w") as f:
            f.write(result)
        return result


class Pipeline:
    """Pipeline analysis framework."""

    def __init__(self, clock, stages: list[Stage]):
        self.clock = clock
        self.stages = stages

    def analyze(self, trace: "Trace") -> PipelineAnalysis:
        """Run pipeline analysis across the trace."""
        from ewal.signal import SignalRef, _compile_expr

        analysis = PipelineAnalysis()

        # Sample at rising clock edges
        if isinstance(self.clock, SignalRef):
            compiled_rising = _compile_expr(self.clock.rising(), trace)
            clk_indices = [i for i in range(trace.max_index + 1)
                           if trace._core.eval_expr_bool(compiled_rising, i)]
        else:
            clk_indices = list(range(trace.max_index + 1))

        analysis.total_cycles = len(clk_indices)

        # Initialize per-stage counters
        for stage in self.stages:
            analysis.stalls_by_stage[stage.name] = 0
            analysis.flushes_by_stage[stage.name] = 0

        # Track instructions by PC
        active_instructions: dict[int, InstructionRecord] = {}  # pc -> record

        for idx in clk_indices:
            old = trace._index
            trace._index = idx
            trace._virtual_cache.clear()

            for stage in self.stages:
                # Check valid
                valid = False
                if stage.valid is not None:
                    valid = bool(stage.valid.value)

                # Check stall
                stalled = False
                if stage.stall is not None:
                    stalled = bool(stage.stall.value)
                    if stalled:
                        analysis.stall_cycles += 1
                        analysis.stalls_by_stage[stage.name] += 1

                # Check flush
                flushed = False
                if stage.flush is not None:
                    flushed = bool(stage.flush.value)
                    if flushed:
                        analysis.flush_count += 1
                        analysis.flushes_by_stage[stage.name] += 1

                if not valid:
                    analysis.bubble_cycles += 1
                    continue

                # Get PC
                pc = 0
                if stage.pc is not None:
                    pc_val = stage.pc.value
                    pc = int(pc_val) if pc_val is not None else 0

                # Track instruction
                if pc not in active_instructions:
                    active_instructions[pc] = InstructionRecord(pc=pc)

                inst = active_instructions[pc]
                if stage.name not in inst.stage_cycles:
                    inst.stage_cycles[stage.name] = idx

                if stalled:
                    inst.stalled = True
                if flushed:
                    inst.flushed = True

                # Check retirement
                if stage.retire is not None and bool(stage.retire.value):
                    analysis.retired += 1
                    if inst.stage_cycles:
                        first = min(inst.stage_cycles.values())
                        inst.latency = idx - first
                    analysis.instructions.append(inst)
                    if pc in active_instructions:
                        del active_instructions[pc]

            trace._index = old

        # CPI
        if analysis.retired > 0:
            analysis.cpi = analysis.total_cycles / analysis.retired

        return analysis
