from __future__ import annotations

from typing import Any

from rlm_core import DataSourceRegistry
from rlm_core.protocols.environment import EnvCapabilities, ExecutionResult, RLMEnvironment

from .restricted_executor import RestrictedExecutor
from .subcalls import RclSubcallClient


class RecallEnvironment(RLMEnvironment):
    """Recall REPL environment backed by RestrictedPython."""

    def __init__(
        self,
        *,
        executor: RestrictedExecutor,
        registry: DataSourceRegistry,
        subcalls: RclSubcallClient,
        max_output_chars: int,
    ) -> None:
        self._executor = executor
        self._registry = registry
        self._subcalls = subcalls
        self._max_output_chars = max_output_chars
        self._globals: dict[str, Any] = {
            "answer": {"content": "", "ready": False},
            "llm_query": subcalls.llm_query,
            "llm_batch": subcalls.llm_batch,
            "batch_llm_query": subcalls.llm_batch,
            "chunk_data": subcalls.chunk_data,
        }
        self._globals.update(self._registry.globals())

    def capabilities(self) -> EnvCapabilities:
        return {
            "tools_in_root": False,
            "max_output_chars": self._max_output_chars,
        }

    def globals(self) -> dict[str, Any]:
        return self._globals

    def prompt_fragment(self) -> str:
        return self._registry.prompt_fragment()

    def execute(self, code: str) -> ExecutionResult:
        output = self._executor.execute(code, extra_globals=self._globals)
        return ExecutionResult(
            stdout=output,
            stderr="",
            timed_out=False,
            exit_code=0,
            files_written=[],
        )

    def close(self) -> None:
        self._executor.close()
