"""
SQL Debug Environment — OpenEnv-compliant async client.

Implements the standard EnvClient interface:
  async with SQLDebugEnv(base_url="...") as client:
      result = await client.reset()
      result = await client.step(SQLFixAction(query="SELECT ..."))

Sync wrapper is also available:
  with SQLDebugEnv(base_url="...").sync() as client:
      result = client.reset()
      result = client.step(SQLFixAction(query="SELECT ..."))
"""

import asyncio
from dataclasses import dataclass
from typing import Optional
import httpx

from models import Action, Observation, State


@dataclass
class StepResult:
    """Returned by step() and reset()."""
    observation: Observation
    reward: float
    done: bool
    success: bool
    state: Optional[State] = None


class SQLDebugEnv:
    """
    Async HTTP client for the SQL Debug Environment.

    Compatible with the OpenEnv EnvClient interface.

    Usage (async):
        async with SQLDebugEnv(base_url="http://localhost:8000") as env:
            result = await env.reset(task_group="easy_syntax_fix")
            result = await env.step(SQLFixAction(query="SELECT 1"))

    Usage (sync):
        with SQLDebugEnv(base_url="http://localhost:8000").sync() as env:
            result = env.reset()
            result = env.step(SQLFixAction(query="SELECT 1"))
    """

    def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._session_id: Optional[str] = None
        self._client: Optional[httpx.AsyncClient] = None

    #Context Manager
    async def __aenter__(self):
        self._client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout)
        return self

    async def __aexit__(self, *args):
        if self._client:
            await self._client.aclose()
        self._client = None

    def sync(self) -> "SyncSQLDebugEnv":
        """Return a synchronous wrapper around this async client."""
        return SyncSQLDebugEnv(self)

    #Core OpenEnv API
    async def reset(
        self,
        task_group: Optional[str] = None,
        task_id: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> StepResult:
        """Initialize or restart an episode."""
        payload = {
            "session_id": self._session_id,
            "task_group": task_group,
            "task_id": task_id,
            "seed": seed,
        }
        resp = await self._post("/reset", payload)
        self._session_id = resp["session_id"]
        obs = Observation(**resp["observation"])
        state = State(**resp["state"]) if "state" in resp else None
        return StepResult(observation=obs, reward=0.0, done=False, success=False, state=state)

    async def step(self, action: Action) -> StepResult:
        """Submit a SQL fix attempt."""
        if not self._session_id:
            raise RuntimeError("No active session. Call reset() first.")
        payload = {
            "session_id": self._session_id,
            "query": action.query,
            "explanation": action.explanation or "",
        }
        resp = await self._post("/step", payload)
        obs = Observation(**resp["observation"])
        state = State(**resp["state"]) if "state" in resp else None
        return StepResult(
            observation=obs,
            reward=resp["reward"],
            done=resp["done"],
            success=resp["success"],
            state=state,
        )

    async def state(self) -> State:
        """Return current episode state without acting."""
        if not self._session_id:
            raise RuntimeError("No active session. Call reset() first.")
        client = self._client or httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout)
        try:
            r = await client.get("/state", params={"session_id": self._session_id})
            r.raise_for_status()
            return State(**r.json())
        finally:
            if not self._client:
                await client.aclose()

    async def _post(self, endpoint: str, payload: dict) -> dict:
        client = self._client
        own_client = False
        if client is None:
            client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout)
            own_client = True
        try:
            resp = await client.post(endpoint, json=payload)
            resp.raise_for_status()
            return resp.json()
        finally:
            if own_client:
                await client.aclose()

    @classmethod
    def from_hub(cls, repo_id: str, port: int = 8000, **kwargs) -> "SQLDebugEnv":
        """
        Connect to an environment hosted on Hugging Face Spaces.

        Example:
            env = SQLDebugEnv.from_hub("openenv/sql-debug-env")
        """
        space_name = repo_id.replace("/", "-").lower()
        base_url = f"https://{space_name}.hf.space"
        return cls(base_url=base_url, **kwargs)


class SyncSQLDebugEnv:
    """Synchronous wrapper around SQLDebugEnv."""

    def __init__(self, async_env: SQLDebugEnv):
        self._env = async_env
        self._loop: Optional[asyncio.AbstractEventLoop] = None

    def __enter__(self):
        self._loop = asyncio.new_event_loop()
        self._loop.run_until_complete(self._env.__aenter__())
        return self

    def __exit__(self, *args):
        self._loop.run_until_complete(self._env.__aexit__(*args))
        self._loop.close()

    def reset(self, **kwargs) -> StepResult:
        return self._loop.run_until_complete(self._env.reset(**kwargs))

    def step(self, action: Action) -> StepResult:
        return self._loop.run_until_complete(self._env.step(action))

    def state(self) -> State:
        return self._loop.run_until_complete(self._env.state())


SQLFixAction = Action