"""The public OpenEnv training session contract over a remote native OpenCode server."""

from pathlib import Path

from openenv.core.harness import (
    ResourceSession,
    ToolResult,
    TrainingTrace,
    VerifyResult,
)
from openenv.core.mcp_client import MCPToolClient


class RemoteSession(ResourceSession):
    def __init__(self, server, row, model, sampling):
        self.client = MCPToolClient(server, message_timeout_s=1800).sync()
        self.row, self.model, self.sampling = row, model, sampling
        self.result = None

    def initial_messages(self):
        return [{"role": "user", "content": self.row["instruction"]}]

    def list_tools(self):
        return []

    def call_tool(self, name, arguments):
        return ToolResult(error="OpenCode owns its tool loop")

    def wait_for_completion(self, timeout_s=None):
        self.client.reset()
        folder = Path(self.row["folder"])
        self.result = self.client.call_tool(
            "run_rollout",
            split=folder.parent.parent.name,
            task_name=folder.name,
            model=self.model,
            sampling=self.sampling,
        )
        return 0

    def fetch_training_trace(self):
        if self.result is None:
            raise RuntimeError("Rollout has not completed")
        if "capture_error" in self.result:
            raise ValueError(self.result["capture_error"])
        return TrainingTrace.model_validate(self.result["training_trace"])

    def verify(self, transcript, final_state=None):
        return VerifyResult(
            env_reward=self.result["correctness"] if self.result else None, done=True
        )

    def close(self):
        self.client.close()


class RemoteTaskFactory:
    def __init__(self, args, tasks, *, sampling):
        self.args, self.sampling = args, sampling
        self.tasks = {row["instruction"]: row for row in tasks}

    def create(self, task, seed=None, episode_id=None):
        return RemoteSession(
            self.args.server,
            self.tasks[task[-1]["content"]],
            self.args.model,
            self.sampling,
        )