import asyncio
import logging
import re
from collections.abc import Callable

from agp_client import (
    AgPClient,
    AgPClientConfig,
    AgPRunResult,
    AgPTaskRequest,
    DirectEnvironmentConnection,
)

logger = logging.getLogger("surfer-agent")


class SurferAgent:
    """Delegates full task execution to Surfer H running on H's Agent Platform.

    Submits a task to AGP with a DirectEnvironmentConnection pointing at the
    RDDS instance on the VM. The agent config lives server-side on AGP.
    """

    DEFAULT_AGENT_IDENTIFIER = "OSWorldSurfer-Holo3-35B-A3B-20260330-OSWorld-Verified"
    _API_KEY = "hk-f8a86be2c0e53c0c6c576e9ed20a53e9a990a2166c885b17"
    _AGP_TIMEOUT = 3600*3
    action_space = "pyautogui"

    def __init__(self, agent_identifier: str = DEFAULT_AGENT_IDENTIFIER):
        self.agent_identifier = agent_identifier
        self.rdds_url: str = ""
        self.last_result: AgPRunResult | None = None
        self.last_trajectory_id: str | None = None

    def reset(self, rdds_url: str):
        self.rdds_url = rdds_url
        self.last_trajectory_id = None
        logger.info(f"RDDS URL set to {rdds_url}")

    def predict(
        self,
        instruction: str,
        obs: dict,
        on_trajectory_created: Callable[[str], None] | None = None,
    ) -> tuple[str, list[str]]:
        assert self.rdds_url, "rdds_url not set -- call reset(rdds_url=...) first"

        result = asyncio.run(self._run_on_agp(instruction, on_trajectory_created))
        self.last_result = result

        if result.success:
            # OSWorld evaluator needs action_history[-1] == "FAIL" for infeasible scoring.
            if isinstance(result.answer, str) and re.search(r"task.{0,4}infeasible", result.answer, re.IGNORECASE):
                logger.info(f"[{result.trajectory_id[:8]}] Task declared infeasible")
                return f"Infeasible (trajectory={result.trajectory_id})", ["FAIL"]

            logger.info(f"[{result.trajectory_id[:8]}] Task completed")
            return f"Completed (trajectory={result.trajectory_id})", ["DONE"]

        logger.error(f"[{result.trajectory_id[:8]}] Task failed: {result.status.value}")
        return f"Failed: {result.status.value}", ["FAIL"]

    async def _run_on_agp(
        self,
        instruction: str,
        on_trajectory_created: Callable[[str], None] | None = None,
    ) -> AgPRunResult:
        environment = DirectEnvironmentConnection(url=self.rdds_url)

        config = AgPClientConfig(timeout=self._AGP_TIMEOUT)
        async with AgPClient(self._API_KEY, config) as client:
            task = AgPTaskRequest(objective=instruction, environment=environment).public()
            trajectory_id = await client.create_trajectory(
                self.agent_identifier,
                task,
            )
            self.last_trajectory_id = trajectory_id
            if on_trajectory_created:
                on_trajectory_created(trajectory_id)
            logger.info(f"[{trajectory_id[:8]}] Trajectory created, waiting for completion...")
            return await client.wait_for_completion(trajectory_id)
