# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <mpippi@gmail.com>
#
# SPDX-License-Identifier: MIT
"""
File-based prompt registry implementation that stores all prompts in a single JSON file.

This module provides functionality to store and retrieve prompts using a single JSON file
as the storage backend. The file contains an index of all prompts with their associated
metadata and content.
"""

from pathlib import Path

from pydantic import BaseModel

from banks.errors import InvalidPromptError, PromptNotFoundError
from banks.prompt import Prompt, PromptModel


class PromptRegistryIndex(BaseModel):
    """
    Model representing the registry index containing all prompts.

    Stores a list of PromptModel objects that represent all prompts in the registry.
    """

    prompts: list[PromptModel] = []


class FilePromptRegistry:
    """A prompt registry storing all prompt data in a single JSON file."""

    def __init__(self, registry_index: Path) -> None:
        """
        Initialize the file prompt registry.

        Args:
            registry_index: Path to the JSON file that will store the prompts

        Note:
            Creates parent directories if they don't exist.
        """
        self._index_fpath: Path = registry_index
        self._index: PromptRegistryIndex = PromptRegistryIndex(prompts=[])
        try:
            self._index = PromptRegistryIndex.model_validate_json(self._index_fpath.read_text())
        except FileNotFoundError:
            # init the user data folder
            self._index_fpath.parent.mkdir(parents=True, exist_ok=True)

    def get(self, *, name: str, version: str | None = None) -> Prompt:
        """
        Retrieve a prompt by name and version.

        Args:
            name: Name of the prompt to retrieve
            version: Version of the prompt (optional)

        Returns:
            The requested Prompt object

        Raises:
            PromptNotFoundError: If the requested prompt doesn't exist
        """
        _, model = self._get_prompt_model(name, version)
        return Prompt(**model.model_dump())

    def set(self, *, prompt: Prompt, overwrite: bool = False) -> None:
        """
        Store a prompt in the registry.

        Args:
            prompt: The Prompt object to store
            overwrite: Whether to overwrite an existing prompt

        Raises:
            InvalidPromptError: If prompt exists and overwrite=False
        """
        try:
            idx, p_model = self._get_prompt_model(prompt.name, prompt.version)
            if overwrite:
                self._index.prompts[idx] = PromptModel.from_prompt(prompt)
                self._save()
            else:
                msg = f"Prompt with name '{prompt.name}' already exists. Use overwrite=True to overwrite"
                raise InvalidPromptError(msg)
        except PromptNotFoundError:
            p_model = PromptModel.from_prompt(prompt)
            self._index.prompts.append(p_model)
            self._save()

    def _save(self) -> None:
        """
        Save the prompt index to the JSON file.

        Writes the current state of the registry to disk.
        """
        with open(self._index_fpath, "w", encoding="locale") as f:
            f.write(self._index.model_dump_json())

    def _get_prompt_model(self, name: str | None, version: str | None) -> tuple[int, PromptModel]:
        """
        Find a prompt model in the index by name and version.

        Args:
            name: Name of the prompt
            version: Version of the prompt

        Returns:
            Tuple of (index position, PromptModel)

        Raises:
            PromptNotFoundError: If the prompt doesn't exist in the index
        """
        for i, model in enumerate(self._index.prompts):
            if model.name == name and model.version == version:
                return i, model

        msg = f"cannot find prompt with name '{name}' and version '{version}'"
        raise PromptNotFoundError(msg)
