from dataclasses import dataclass
from typing import TypeVar

from kirin import interp

from bloqade.shuttle.dialects.path import Parallel, Path, Play, dialect
from bloqade.shuttle.visualizer.interp import PathVisualizer
from bloqade.shuttle.visualizer.renderers import RendererInterface


@dataclass
class ParallelRuntime:
    paths: tuple[Path, ...]

    def __post_init__(self):
        if any(not isinstance(p, Path) for p in self.paths):
            raise interp.InterpreterError("path is not canonicalized")


@dialect.register(key="path.visualizer")
class PathVisualizerMethods(interp.MethodTable):

    Renderer = TypeVar("Renderer", bound=RendererInterface)

    @interp.impl(Play)
    def play(self, _interp: PathVisualizer[Renderer], frame: interp.Frame, stmt: Play):

        path = frame.get(stmt.path)
        # path is generated by "main" interpreter so
        # we need to render the path before showing it
        if isinstance(path, Path):

            _interp.renderer.render_path(path)
        elif isinstance(path, ParallelRuntime):
            for p in path.paths:
                _interp.renderer.render_path(p)
        else:
            raise interp.InterpreterError(
                f"Expected a Path or tuple of Paths, got {path}"
            )

        return ()

    @interp.impl(Parallel)
    def parallel(
        self, _interp: "PathVisualizer[Renderer]", frame: interp.Frame, stmt: Parallel
    ):
        return (ParallelRuntime(frame.get_values(stmt.paths)),)
