# client.py
# A simple Python client for calling the environment server
# You could do this with httpx directly but wrapping it is cleaner
#
# Use this in inference.py or anywhere else you need to talk to the server

import httpx
from typing import Optional
from .models import Observation, StepResult, Action


class StudyEnvClient:
    """
    HTTP client for the Study Optimizer environment server.

    Usage:
        client = StudyEnvClient("http://localhost:8000")
        obs = client.reset()
        result = client.step(0)  # 0=study
    """

    def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        # using httpx because requests doesn't support async
        # TODO: add async versions of these methods later
        self._client = httpx.Client(timeout=self.timeout)

    def health(self) -> dict:
        """Quick check to see if the server is up."""
        resp = self._client.get(f"{self.base_url}/health")
        resp.raise_for_status()
        return resp.json()

    def reset(self, task_id: Optional[str] = None) -> Observation:
        """Reset environment and get first observation."""
        payload = {}
        if task_id:
            payload["task_id"] = task_id

        resp = self._client.post(f"{self.base_url}/reset", json=payload)
        resp.raise_for_status()
        data = resp.json()

        # took me a while to figure this out - the server wraps in 'observation' key
        if "observation" in data:
            return Observation(**data["observation"])
        return Observation(**data)

    def step(self, action: int) -> StepResult:
        """Take an action and get the result."""
        payload = {"action": action}
        resp = self._client.post(f"{self.base_url}/step", json=payload)
        resp.raise_for_status()
        data = resp.json()
        return StepResult(**data)

    def get_state(self) -> dict:
        """Get the full current state of the environment."""
        resp = self._client.get(f"{self.base_url}/state")
        resp.raise_for_status()
        return resp.json()

    def close(self):
        """Clean up the HTTP client."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()