"""
Utilities for managing agent submissions.  Currently, these are aimed at
internal agent usage; tasks generally expect submissions to be made by
assigning `state.output.completion`.

Use `get_submission_manager()` to retrieve the current submission manager,
which can be written to with `manager.write_submission()` and read with
`manager.get_submission()`.

The `submit_tool` in this module internally writes to the submission manager,
so e.g. in a ReAct loop you could check
`get_submission_manager().has_submission()` to determine when the agent is
done.  See the README for more details.
"""

import logging

from inspect_ai.tool import Tool, ToolResult, tool
from inspect_ai.util import StoreModel, store_as

logger = logging.getLogger(__name__)


class MissingSubmissionError(Exception):
    """Raised when attempting to retrieve a submission that does not exist."""


class Submission(StoreModel):
    submission: str | None = None


class SubmissionManager(StoreModel):
    """Manager for handling submissions."""

    submission: str | None = None

    def write_submission(self, submission: str) -> None:
        """Store or overwrite the submission."""
        self.submission = submission

    def get_submission(self) -> str:
        """Retrieve the stored submission. Raises MissingSubmissionError if no submission exists."""
        if self.submission is None:
            raise MissingSubmissionError("No submission has been made.")
        return self.submission

    def has_submission(self) -> bool:
        """Check if a submission exists."""
        return self.submission is not None

    def clear_submission(self) -> None:
        """Clear the stored submission; i.e., `has_submission()` will return False after this call."""
        self.submission = None


def get_submission_manager() -> SubmissionManager:
    """Get the SubmissionManager instance from the current sample Store."""
    return store_as(SubmissionManager)


@tool
def submit_tool() -> Tool:
    async def submit(submission: str) -> ToolResult:
        """Submit an answer for evaluation.

        Args:
          submission (str): Submitted answer
        """
        get_submission_manager().write_submission(submission)
        return "Submission successful."

    return submit
