from __future__ import annotations

import json
from typing import Any
from urllib import request

from llmserve_env.models import EpisodeLog, ServeAction, ServeObservation, ServeState


class LLMServeEnv:
    def __init__(self, base_url: str) -> None:
        self.base_url = base_url.rstrip("/")
        self.session_id: str | None = None

    @classmethod
    def from_url(cls, base_url: str) -> "LLMServeEnv":
        return cls(base_url=base_url)

    @classmethod
    def from_hub(cls, repo_id: str) -> "LLMServeEnv":
        return cls(base_url=f"https://huggingface.co/spaces/{repo_id}")

    def reset(self, task_id: str, seed: int | None = None) -> ServeObservation:
        payload = self._post("/reset", {"task_id": task_id, "seed": seed})
        self.session_id = payload.get("session_id")
        return self._parse_observation_payload(payload)

    def step(self, action: dict[str, Any] | ServeAction) -> tuple[ServeObservation, float, bool, dict[str, Any]]:
        action_payload = action.model_dump(mode="json") if isinstance(action, ServeAction) else action
        body: dict[str, Any] = {"action": action_payload}
        if self.session_id is not None:
            body["session_id"] = self.session_id
        payload = self._post("/step", body)
        observation = self._parse_observation_payload(payload)
        if payload.get("session_id") and self.session_id is None:
            self.session_id = str(payload["session_id"])
        return observation, float(payload["reward"]), bool(payload["done"]), observation.metadata

    def state(self) -> ServeState:
        path = f"/state?session_id={self.session_id}" if self.session_id is not None else "/state"
        payload = self._get(path)
        return ServeState.model_validate(payload)

    def tasks(self) -> dict[str, Any]:
        return self._get("/tasks")

    def grade(self, log: EpisodeLog | None = None) -> dict[str, Any]:
        body = {} if log is None else {"episode_log": log.model_dump(mode="json")}
        return self._post("/grader", body)

    def baseline(self, task_id: str | None = None, use_openai: bool = False, model: str | None = None) -> dict[str, Any]:
        params = []
        if task_id:
            params.append(f"task_id={task_id}")
        if use_openai:
            params.append("use_openai=true")
        if model:
            params.append(f"model={model}")
        suffix = f"?{'&'.join(params)}" if params else ""
        return self._get(f"/baseline{suffix}")

    def _parse_observation_payload(self, payload: dict[str, Any]) -> ServeObservation:
        observation_payload = dict(payload["observation"])
        observation_payload["reward"] = payload.get("reward")
        observation_payload["done"] = payload.get("done", False)
        return ServeObservation.model_validate(observation_payload)

    def _get(self, path: str) -> dict[str, Any]:
        with request.urlopen(f"{self.base_url}{path}") as response:
            return json.loads(response.read().decode("utf-8"))

    def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        body = json.dumps(payload).encode("utf-8")
        headers = {"Content-Type": "application/json"}
        req = request.Request(f"{self.base_url}{path}", data=body, headers=headers, method="POST")
        with request.urlopen(req) as response:
            return json.loads(response.read().decode("utf-8"))