"""Asynchronous Areteos client."""

from __future__ import annotations

import asyncio
from typing import Any
from uuid import UUID

import httpx

from . import _jsonapi as ja
from .client import _JSONAPI_HEADERS, _raise_for_status
from .exceptions import TimeoutError
from .models import (
    ApprovalRequest,
    ApprovalStatus,
    Edge,
    Run,
    RunStatus,
    Step,
    StepRun,
    Workflow,
)


class AreteosAsyncClient:
    """Asynchronous client for the Areteos Workflow Runtime API.

    Usage::

        async with AreteosAsyncClient(
            base_url="https://areteos.example.com",
            api_key="areteos_...",
            engagement_id="your-engagement-uuid",
        ) as client:
            workflow = await client.create_workflow(
                name="My Pipeline",
                steps=[Step.python("process", "Process", code="...")],
            )
            run = await client.trigger(workflow.current_version_id, input={"key": "value"})
            result = await client.wait_for_run(run.id)
    """

    def __init__(
        self,
        *,
        base_url: str,
        api_key: str,
        engagement_id: str | UUID,
        timeout: float = 30.0,
    ) -> None:
        self.engagement_id = str(engagement_id)
        self._http = httpx.AsyncClient(
            base_url=f"{base_url.rstrip('/')}/api/json",
            headers={
                "Authorization": f"Bearer {api_key}",
                **_JSONAPI_HEADERS,
            },
            timeout=timeout,
        )

    async def close(self) -> None:
        await self._http.aclose()

    async def __aenter__(self) -> AreteosAsyncClient:
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()

    # --- Internal ---

    async def _get(self, path: str, **params: Any) -> dict[str, Any]:
        resp = await self._http.get(path, params=params)
        _raise_for_status(resp)
        result: dict[str, Any] = resp.json()
        return result

    async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
        resp = await self._http.post(path, json=body)
        _raise_for_status(resp)
        result: dict[str, Any] = resp.json()
        return result

    async def _patch(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
        resp = await self._http.patch(path, json=body)
        _raise_for_status(resp)
        result: dict[str, Any] = resp.json()
        return result

    # --- Workflows ---

    async def create_workflow(
        self,
        name: str,
        steps: list[Step],
        *,
        description: str | None = None,
        edges: list[Edge] | None = None,
        change_note: str | None = None,
    ) -> Workflow:
        """Register a workflow with steps and edges in a single call."""
        attrs: dict[str, Any] = {
            "name": name,
            "engagement_id": self.engagement_id,
            "steps": [s.model_dump(exclude_none=True) for s in steps],
            "edges": [e.model_dump(exclude_none=True) for e in (edges or [])],
        }
        if description is not None:
            attrs["description"] = description
        if change_note is not None:
            attrs["change_note"] = change_note

        body = await self._post(
            "/workflows/create-with-version",
            ja.wrap_create("workflows", attrs),
        )
        return Workflow.model_validate(ja.unwrap_one(body))

    async def get_workflow(self, workflow_id: str | UUID) -> Workflow:
        body = await self._get(f"/workflows/{workflow_id}")
        return Workflow.model_validate(ja.unwrap_one(body))

    async def list_workflows(self) -> list[Workflow]:
        body = await self._get("/workflows")
        return [Workflow.model_validate(item) for item in ja.unwrap_many(body)]

    # --- Runs ---

    async def trigger(
        self,
        workflow_version_id: str | UUID,
        *,
        input: dict[str, Any] | None = None,
    ) -> Run:
        """Trigger a new run of a workflow version."""
        attrs: dict[str, Any] = {
            "workflow_version_id": str(workflow_version_id),
            "engagement_id": self.engagement_id,
        }
        if input is not None:
            attrs["input"] = input

        body = await self._post("/runs", ja.wrap_create("runs", attrs))
        return Run.model_validate(ja.unwrap_one(body))

    async def get_run(self, run_id: str | UUID) -> Run:
        body = await self._get(f"/runs/{run_id}")
        return Run.model_validate(ja.unwrap_one(body))

    async def list_runs(self) -> list[Run]:
        body = await self._get("/runs")
        return [Run.model_validate(item) for item in ja.unwrap_many(body)]

    async def cancel_run(self, run_id: str | UUID) -> Run:
        body = await self._patch(
            f"/runs/{run_id}/cancel",
            ja.wrap_update("runs", str(run_id), {}),
        )
        return Run.model_validate(ja.unwrap_one(body))

    async def pause_run(self, run_id: str | UUID) -> Run:
        body = await self._patch(
            f"/runs/{run_id}/pause",
            ja.wrap_update("runs", str(run_id), {}),
        )
        return Run.model_validate(ja.unwrap_one(body))

    async def wait_for_run(
        self,
        run_id: str | UUID,
        *,
        poll_interval: float = 5.0,
        timeout: float = 600.0,
        auto_approve: bool = False,
    ) -> Run:
        """Poll until a run reaches a terminal state.

        Args:
            run_id: The run ID to poll.
            poll_interval: Seconds between polls.
            timeout: Maximum seconds to wait before raising TimeoutError.
            auto_approve: If True, automatically approve any HITL requests.

        Returns:
            The completed Run.

        Raises:
            TimeoutError: If the run does not complete within the timeout.
        """
        deadline = asyncio.get_event_loop().time() + timeout

        while True:
            run = await self.get_run(run_id)

            if run.is_terminal:
                return run

            if run.status == RunStatus.PAUSED and auto_approve:
                await self._auto_approve_pending()

            if asyncio.get_event_loop().time() >= deadline:
                raise TimeoutError(
                    f"Run {run_id} did not complete within {timeout}s (status: {run.status})"
                )

            await asyncio.sleep(poll_interval)

    # --- Step Runs ---

    async def get_step_runs(self, run_id: str | UUID) -> list[StepRun]:
        body = await self._get("/step-runs", **{"filter[run_id]": str(run_id)})
        return [StepRun.model_validate(item) for item in ja.unwrap_many(body)]

    async def get_step_run(self, step_run_id: str | UUID) -> StepRun:
        body = await self._get(f"/step-runs/{step_run_id}")
        return StepRun.model_validate(ja.unwrap_one(body))

    # --- Approval Requests ---

    async def list_approval_requests(
        self, *, status: ApprovalStatus | None = None
    ) -> list[ApprovalRequest]:
        params = {}
        if status is not None:
            params["filter[status]"] = str(status)
        body = await self._get("/approval-requests", **params)
        return [ApprovalRequest.model_validate(item) for item in ja.unwrap_many(body)]

    async def respond_to_approval(
        self,
        approval_id: str | UUID,
        *,
        status: ApprovalStatus,
        response_text: str | None = None,
        edited_output: dict[str, Any] | None = None,
    ) -> ApprovalRequest:
        attrs: dict[str, Any] = {"status": str(status)}
        if response_text is not None:
            attrs["response_text"] = response_text
        if edited_output is not None:
            attrs["edited_output"] = edited_output

        body = await self._patch(
            f"/approval-requests/{approval_id}",
            ja.wrap_update("approval-requests", str(approval_id), attrs),
        )
        return ApprovalRequest.model_validate(ja.unwrap_one(body))

    async def approve(
        self, approval_id: str | UUID, *, response_text: str | None = None
    ) -> ApprovalRequest:
        return await self.respond_to_approval(
            approval_id, status=ApprovalStatus.APPROVED, response_text=response_text
        )

    async def reject(
        self, approval_id: str | UUID, *, response_text: str | None = None
    ) -> ApprovalRequest:
        return await self.respond_to_approval(
            approval_id, status=ApprovalStatus.REJECTED, response_text=response_text
        )

    # --- Helpers ---

    async def _auto_approve_pending(self) -> None:
        pending = await self.list_approval_requests(status=ApprovalStatus.PENDING)
        for req in pending:
            await self.approve(req.id, response_text="Auto-approved by SDK")
