from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Any

from ._exceptions import APIError
from ._streaming import ChatCompletionStream, ChatCompletionStreamManager
from .types.budget import BudgetStatus
from .types.rungate_metadata import RungateMetadata

if TYPE_CHECKING:
    from ._client import Rungate


class Run:
    """An active Rungate run with budget tracking and tool call logging.

    All LLM calls made through this run automatically include the run ID.
    Budget is updated in real-time from response headers (for non-streaming)
    or after stream consumption (for streaming).

    Usage as context manager (recommended):
        with rg.run(name="research") as run:
            response = run.chat.completions.create(model="gpt-4o", messages=[...])
            print(f"Budget remaining: ${run.budget.remaining_usd}")
            run.log_tool_call(name="web_search", request={...}, response={...})
        # Run auto-completes on exit, auto-fails on exception

    Usage without context manager (for long-running agents):
        run = rg.runs.create(name="multi-session-task")
        response = run.chat.completions.create(model="gpt-4o", messages=[...])
        run.complete()

    Thread safety: Run objects are NOT thread-safe. Each thread or coroutine
    should use its own Run instance. Budget tracking uses non-atomic
    operations that can race under concurrent access.
    """

    def __init__(
        self,
        client: Rungate,
        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 = RunChat(self)
        self._completed = False

    def _update_budget(self, meta: RungateMetadata) -> None:
        """Update budget tracking from Rungate response metadata."""
        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
        )

    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.

        Args:
            name: Tool name (e.g., "web_search", "database_query").
            request: Tool call request payload.
            response: Tool call response payload.
            provider: Tool provider (default: "sdk").
            latency_ms: Execution latency in milliseconds.
            cost_usd: Cost of the tool call (default: 0).
            status: Step status ("success", "error", "timeout").
            error_message: Error message if status is "error".

        Returns:
            The created step data from the server.
        """
        resp = 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", {})

    def complete(self, *, reason: str = "agent_signaled", metadata: dict[str, Any] | None = None) -> dict[str, Any]:
        """Explicitly complete this run.

        Safe to call multiple times — subsequent calls are no-ops.
        Handles 409 (already completed) from the server gracefully.
        """
        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 = 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

    def fail(self, *, reason: str = "error", metadata: dict[str, Any] | None = None) -> dict[str, Any]:
        """Mark this run as failed."""
        return self.complete(reason=reason, metadata=metadata)

    def __enter__(self) -> Run:
        return self

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


class RunChat:
    """Chat resource scoped to a run. Auto-injects run_id on all calls."""

    def __init__(self, run: Run) -> None:
        self.completions = RunCompletions(run)


class RunCompletions:
    """Completions scoped to a run. Auto-injects run_id, updates budget."""

    def __init__(self, run: Run) -> None:
        self._run = run

    def create(self, **kwargs: Any) -> Any:
        """Create a chat completion within this run.

        Automatically injects run_id and updates budget from response.
        Accepts the same parameters as client.chat.completions.create().

        For streaming, use run.chat.completions.stream() instead — it
        properly updates budget after stream consumption.
        """
        kwargs["run_id"] = self._run.id
        result = self._run._client.chat.completions.create(**kwargs)

        # Only update budget for non-streaming responses (ChatCompletion)
        # Streaming results (RungateStream) don't have metadata until consumed
        from .types.chat.completion import ChatCompletion

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

        return result

    def stream(self, **kwargs: Any) -> RunChatCompletionStreamManager:
        """Stream a chat completion within this run.

        Returns a context manager that auto-updates budget after consumption.
        """
        kwargs["run_id"] = self._run.id
        inner_manager = self._run._client.chat.completions.stream(**kwargs)
        return RunChatCompletionStreamManager(self._run, inner_manager)


class RunChatCompletionStreamManager:
    """Wraps ChatCompletionStreamManager to update run budget after streaming."""

    def __init__(self, run: Run, inner: ChatCompletionStreamManager) -> None:
        self._run = run
        self._inner = inner
        self._stream: ChatCompletionStream | None = None

    def __enter__(self) -> ChatCompletionStream:
        self._stream = self._inner.__enter__()
        return self._stream

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


class Runs:
    """Run management resource on the client."""

    def __init__(self, client: Rungate) -> None:
        self._client = client

    def create(
        self,
        *,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Run:
        """Create a new explicit run.

        The run ID is generated client-side (UUID) and sent via header
        on the first LLM call, which triggers server-side run creation.
        """
        run_id = str(uuid.uuid4())
        return Run(self._client, run_id, name=name, metadata=metadata)

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

        Use this when restarting a process that was working on a
        long-running run. The run must already exist on the server.

        Note: Budget tracking starts fresh from this point. The server
        maintains the authoritative run totals; client-side budget
        tracking only reflects calls made in this session.
        """
        return Run(self._client, run_id, name=name, metadata=metadata)
