import yaml
from pathlib import Path

from promptlock.exceptions import PromptNotFound
from promptlock.template import PromptTemplate


class PromptRegistry:
    """Stores and retrieves versioned prompts from a local YAML file."""

    def __init__(self, path: str = "prompts.yaml"):
        self.path = Path(path)
        self._data = self._load_file()

    def _load_file(self) -> dict:
        if self.path.exists():
            with open(self.path, "r") as f:
                return yaml.safe_load(f) or {}
        return {}

    def _save_file(self):
        with open(self.path, "w") as f:
            yaml.dump(self._data, f, default_flow_style=False, allow_unicode=True)

    def save(self, name: str, version: str, text: str):
        """Save a prompt version to the registry.

        Args:
            name: prompt name (e.g. 'summarizer')
            version: version tag (e.g. 'v1.0')
            text: the prompt string with optional {variables}
        """
        if name not in self._data:
            self._data[name] = {}
        self._data[name][version] = text
        self._save_file()

    def load(self, name: str, version: str = "latest") -> PromptTemplate:
        """Load a prompt from the registry as a PromptTemplate.

        Args:
            name: prompt name
            version: version tag, or 'latest' to get the most recent

        Raises:
            PromptNotFound: if the name or version does not exist
        """
        if name not in self._data:
            raise PromptNotFound(name=name)

        versions = self._data[name]

        if version == "latest":
            version = list(versions.keys())[-1]

        if version not in versions:
            raise PromptNotFound(name=name, version=version)

        return PromptTemplate(text=versions[version])

    def list_prompts(self) -> dict:
        """Return all prompt names and their available versions."""
        return {name: list(versions.keys()) for name, versions in self._data.items()}

    def delete(self, name: str, version: str = None):
        """Delete a prompt or a specific version of a prompt.

        Args:
            name: prompt name
            version: if provided, delete only that version. Otherwise delete all versions.

        Raises:
            PromptNotFound: if the name or version does not exist
        """
        if name not in self._data:
            raise PromptNotFound(name=name)

        if version:
            if version not in self._data[name]:
                raise PromptNotFound(name=name, version=version)
            del self._data[name][version]
            if not self._data[name]:
                del self._data[name]
        else:
            del self._data[name]

        self._save_file()

    def __repr__(self) -> str:
        return f"PromptRegistry(path='{self.path}', prompts={list(self._data.keys())})"