from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Any

from ._exceptions import APIError
from ._streaming import AsyncChatCompletionStream, AsyncChatCompletionStreamManager
from .types.budget import BudgetStatus
from .types.rungate_metadata import RungateMetadata

if TYPE_CHECKING:
    from ._client import AsyncRungate


class AsyncRun:
    """Async variant of Run. See Run for full documentation."""

    def __init__(
        self,
        client: AsyncRungate,
        run_id: str,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self._client = client
        self.id = run_id
        self.name = name
        self.metadata = metadata or {}
        self.budget = BudgetStatus()
        self.chat = AsyncRunChat(self)
        self._completed = False

    def _update_budget(self, meta: RungateMetadata) -> None:
        if meta.cost_usd is not None:
            self.budget.total_cost += meta.cost_usd
        if meta.tokens_used is not None:
            self.budget.total_tokens += meta.tokens_used
        if meta.budget_remaining_usd is not None:
            self.budget.remaining_usd = meta.budget_remaining_usd
        if meta.budget_remaining_tokens is not None:
            self.budget.remaining_tokens = meta.budget_remaining_tokens
        self.budget.at_limit = (self.budget.remaining_usd is not None and self.budget.remaining_usd <= 0) or (
            self.budget.remaining_tokens is not None and self.budget.remaining_tokens <= 0
        )

    async def log_tool_call(
        self,
        *,
        name: str,
        request: dict[str, Any] | None = None,
        response: dict[str, Any] | None = None,
        provider: str = "sdk",
        latency_ms: float = 0,
        cost_usd: float = 0,
        tokens: dict[str, int] | None = None,
        status: str = "success",
        error_message: str | None = None,
    ) -> dict[str, Any]:
        """Log a non-LLM tool call as a step in this run."""
        resp = await self._client._request(
            "POST",
            f"/v1/runs/{self.id}/steps",
            json={
                "type": "tool_call",
                "name": name,
                "provider": provider,
                "provider_key_id": "sdk",
                "request": request or {},
                "response": response or {},
                "tokens": tokens or {"input": 0, "output": 0},
                "latency_ms": latency_ms,
                "cost_usd": cost_usd,
                "status": status,
                "error_message": error_message,
            },
        )
        return resp.body.get("data", {})

    async def complete(
        self, *, reason: str = "agent_signaled", metadata: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        """Explicitly complete this run."""
        if self._completed:
            return {}
        merged_metadata = {**self.metadata, **(metadata or {})}
        body: dict[str, Any] = {"ended_reason": reason}
        if merged_metadata:
            body["metadata"] = merged_metadata
        try:
            resp = await self._client._request("POST", f"/v1/runs/{self.id}/complete", json=body)
            self._completed = True
            return resp.body.get("data", {})
        except APIError as e:
            if e.status_code in (404, 409):
                self._completed = True
                return {}
            raise

    async def fail(self, *, reason: str = "error", metadata: dict[str, Any] | None = None) -> dict[str, Any]:
        return await self.complete(reason=reason, metadata=metadata)

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

    async def __aexit__(self, exc_type: type | None, exc_val: BaseException | None, exc_tb: Any) -> None:
        if not self._completed:
            try:
                if exc_type is not None:
                    await self.fail(reason=f"exception: {exc_type.__name__}")
                else:
                    await self.complete()
            except Exception:
                pass


class AsyncRunChat:
    def __init__(self, run: AsyncRun) -> None:
        self.completions = AsyncRunCompletions(run)


class AsyncRunCompletions:
    def __init__(self, run: AsyncRun) -> None:
        self._run = run

    async def create(self, **kwargs: Any) -> Any:
        """Create a chat completion within this run (async)."""
        kwargs["run_id"] = self._run.id
        result = await self._run._client.chat.completions.create(**kwargs)

        from .types.chat.completion import ChatCompletion

        if isinstance(result, ChatCompletion):
            self._run._update_budget(result.rungate)
        return result

    def stream(self, **kwargs: Any) -> AsyncRunChatCompletionStreamManager:
        kwargs["run_id"] = self._run.id
        inner_manager = self._run._client.chat.completions.stream(**kwargs)
        return AsyncRunChatCompletionStreamManager(self._run, inner_manager)


class AsyncRunChatCompletionStreamManager:
    def __init__(self, run: AsyncRun, inner: AsyncChatCompletionStreamManager) -> None:
        self._run = run
        self._inner = inner
        self._stream: AsyncChatCompletionStream | None = None

    async def __aenter__(self) -> AsyncChatCompletionStream:
        self._stream = await self._inner.__aenter__()
        return self._stream

    async def __aexit__(self, *args: Any) -> None:
        try:
            if self._stream:
                self._run._update_budget(self._stream.rungate)
        finally:
            await self._inner.__aexit__(*args)


class AsyncRuns:
    def __init__(self, client: AsyncRungate) -> None:
        self._client = client

    def create(
        self,
        *,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> AsyncRun:
        run_id = str(uuid.uuid4())
        return AsyncRun(self._client, run_id, name=name, metadata=metadata)

    def resume(
        self,
        run_id: str,
        *,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> AsyncRun:
        """Resume an existing run by ID (async)."""
        return AsyncRun(self._client, run_id, name=name, metadata=metadata)
