#
#   Imandra Inc.
#
#   codelogician/sketch/container.py
#

from pydantic import BaseModel
from rich.panel import Panel
from rich.pretty import Pretty

from .sketch import Sketch


class SketchContainer(BaseModel):
    """
    Container for helping us keep track of the sketches
    """

    sketches: dict[str, Sketch] = {}

    def add(self, new_sketch: Sketch) -> None:
        """
        Add spetch to the container
        """
        if new_sketch.sketch_id in self.sketches:
            raise ValueError(
                f'Sketch with specified ID is already present: {new_sketch.sketch_id}'
            )

        self.sketches[new_sketch.sketch_id] = new_sketch

    def from_id(self, sketch_id: str) -> Sketch | None:
        """
        Return sketch from sketch ID, None if it's missing
        """
        return self.sketches.get(sketch_id)

    def ids(self) -> list[str]:
        """
        List of available sketch IDs
        """
        return list(self.sketches.keys())

    def list_sketches(self) -> list[dict[str, str]]:
        """
        List available sketches
        """
        return [
            {
                'sketch_id': sketch_id,
                'anchor_model': self.sketches[sketch_id].anchor_model_path,
            }
            for sketch_id in self.sketches
        ]

    def toJSON(self):
        """
        Return JSON dict
        """
        return self.model_dump_json()

    @staticmethod
    def fromJSON(j: str | dict):
        """
        Create a SketchesContainer object from json dict
        """
        if isinstance(j, str):
            return SketchContainer.model_validate_json(j)
        else:
            return SketchContainer.model_validate(j)

    def __rich__(self):
        """Return the list of rich representations of the sketches we hold"""

        pretty = Pretty(self.toJSON(), indent_guides=True)
        return Panel(pretty, title='SketchContainer', border_style='green')
