"""Asynchronous Baponi client."""

from __future__ import annotations

import os
import warnings
from types import TracebackType
from typing import Self

import httpx

from baponi._base import (
    async_execute_with_retry,
    build_cancel_url,
    build_execute_body,
    build_execute_url,
    build_execution_url,
    build_output_url,
    build_request_timeout,
    default_headers,
    raise_for_status,
    streaming_headers,
    validate_execute_params,
    validate_trace_id,
    validate_webhook_url,
    webhook_headers,
)
from baponi._constants import (
    DEFAULT_BASE_URL,
    DEFAULT_MAX_RETRIES,
    DEFAULT_TIMEOUT,
    FILES_DOWNLOAD_URL_PATH,
    FILES_LIST_PATH,
    FILES_UPLOAD_URL_PATH,
)
from baponi._streaming import AsyncExecutionStream
from baponi._webhook import AsyncExecutionHandle
from baponi.exceptions import BaponiError
from baponi.types import CancelResponse, ExecutionStatus, FileEntry, SandboxResult, SignedUrl, WebhookAccepted


class AsyncBaponi:
    """Asynchronous client for the Baponi sandbox execution API.

    Usage::

        async with AsyncBaponi() as client:
            result = await client.execute("print('hello')")
            print(result.stdout)

    Or without context manager::

        client = AsyncBaponi()
        result = await client.execute("print('hello')")
        await client.close()
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
        http_client: httpx.AsyncClient | None = None,
        _client_id: str = "baponi-python/direct",
    ) -> None:
        resolved_key = api_key or os.environ.get("BAPONI_API_KEY")
        if not resolved_key:
            raise BaponiError(
                "API key is required. Pass api_key= or set BAPONI_API_KEY environment variable."
            )
        self._api_key = resolved_key
        self._base_url = base_url.rstrip("/")
        self._max_retries = max_retries
        self._execute_url = build_execute_url(self._base_url)
        self._headers = default_headers(self._api_key, _client_id)

        if not self._base_url.startswith("https://"):
            warnings.warn(
                f"Base URL {self._base_url!r} does not use HTTPS. "
                "API keys will be sent in plaintext.",
                stacklevel=2,
            )

        if http_client is not None:
            self._client = http_client
            self._owns_client = False
        else:
            self._client = httpx.AsyncClient(timeout=httpx.Timeout(timeout), http2=True)
            self._owns_client = True

    async def execute(
        self,
        code: str,
        *,
        language: str = "python",
        timeout: int = 30,
        thread_id: str | None = None,
        metadata: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
        env_vars: dict[str, str] | None = None,
    ) -> SandboxResult:
        """Execute code in a sandboxed environment.

        Args:
            code: The code to execute.
            language: Programming language. One of: python, bash, node.
            timeout: Execution timeout in seconds (1-3600).
            thread_id: Pass for persistent state (files, pip packages)
                across calls.
            metadata: Key-value metadata for audit trail and filtering.
            sub_paths: Optional list of full sandbox paths (starting with
                /data/{slug}/) to narrow storage mounts at runtime.
            env_vars: Environment variables to set in the sandbox.
                Keys must be uppercase (e.g. MY_VAR). Max 50 vars.

        Returns:
            SandboxResult with stdout, stderr, exit_code, and timing info.

        Raises:
            ValueError: If input validation fails (client-side).
            AuthenticationError: If the API key is invalid (401).
            RateLimitError: If rate limited (429).
            ThreadBusyError: If the thread is busy (409).
            APITimeoutError: If execution timed out on server (504).
            ServerError: If server error (500/503).
            BaponiError: For any other API error.
        """
        validate_execute_params(
            code, language, timeout, thread_id, metadata, sub_paths, env_vars
        )
        body = build_execute_body(
            code, language, timeout, thread_id, metadata, sub_paths, env_vars
        )
        request_timeout = build_request_timeout(timeout)

        return await async_execute_with_retry(
            self._client,
            url=self._execute_url,
            headers=self._headers,
            body=body,
            timeout=request_timeout,
            max_retries=self._max_retries,
        )

    async def execute_stream(
        self,
        code: str,
        *,
        language: str = "python",
        timeout: int = 30,
        thread_id: str | None = None,
        metadata: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
        env_vars: dict[str, str] | None = None,
    ) -> AsyncExecutionStream:
        """Execute code with real-time NDJSON streaming.

        Returns an AsyncExecutionStream that yields typed events.

        Usage::

            async with await client.execute_stream("print('hi')") as stream:
                async for event in stream:
                    print(event)
                result = stream.get_final_result()
        """
        validate_execute_params(
            code, language, timeout, thread_id, metadata, sub_paths, env_vars
        )
        body = build_execute_body(
            code, language, timeout, thread_id, metadata, sub_paths, env_vars
        )
        request_timeout = build_request_timeout(timeout)
        headers = streaming_headers(self._api_key, self._headers["X-Client-Id"])

        request = self._client.build_request(
            "POST", self._execute_url, json=body, headers=headers, timeout=request_timeout
        )
        response = await self._client.send(request, stream=True)

        if not response.is_success:
            try:
                await response.aread()
                raise_for_status(response)
            finally:
                await response.aclose()

        return AsyncExecutionStream(
            response,
            base_url=self._base_url,
            headers=self._headers,
            client=self._client,
        )

    async def execute_webhook(
        self,
        code: str,
        *,
        language: str = "python",
        timeout: int = 30,
        thread_id: str | None = None,
        metadata: dict[str, str] | None = None,
        sub_paths: list[str] | None = None,
        env_vars: dict[str, str] | None = None,
        webhook_url: str | None = None,
    ) -> AsyncExecutionHandle:
        """Execute code in webhook (async) mode.

        The server accepts the request immediately and returns a 202 with
        a trace_id. Use the returned AsyncExecutionHandle to poll for
        status, wait for completion, or cancel.
        """
        validate_execute_params(
            code, language, timeout, thread_id, metadata, sub_paths, env_vars
        )
        body = build_execute_body(
            code, language, timeout, thread_id, metadata, sub_paths, env_vars
        )
        if webhook_url is not None:
            validate_webhook_url(webhook_url)
            body["webhook_url"] = webhook_url
        request_timeout = build_request_timeout(timeout)
        headers = webhook_headers(self._api_key, self._headers["X-Client-Id"])

        response = await self._client.post(
            self._execute_url, json=body, headers=headers, timeout=request_timeout
        )
        raise_for_status(response)
        accepted = WebhookAccepted.model_validate(response.json())

        return AsyncExecutionHandle(
            accepted.trace_id,
            base_url=self._base_url,
            headers=self._headers,
            client=self._client,
        )

    async def get_execution(self, trace_id: str) -> ExecutionStatus:
        """Get the status of an execution by trace_id."""
        validate_trace_id(trace_id)
        url = build_execution_url(self._base_url, trace_id)
        response = await self._client.get(url, headers=self._headers)
        raise_for_status(response)
        return ExecutionStatus.model_validate(response.json())

    async def cancel_execution(self, trace_id: str) -> CancelResponse:
        """Cancel a running execution."""
        validate_trace_id(trace_id)
        url = build_cancel_url(self._base_url, trace_id)
        response = await self._client.post(url, headers=self._headers)
        raise_for_status(response)
        return CancelResponse.model_validate(response.json())

    async def reconnect_stream(
        self, trace_id: str, *, after_seq: int = 0
    ) -> AsyncExecutionStream:
        """Reconnect to an existing execution's NDJSON stream."""
        validate_trace_id(trace_id)
        url = build_output_url(self._base_url, trace_id)
        headers = streaming_headers(self._api_key, self._headers["X-Client-Id"])
        params = {"after_seq": str(after_seq)} if after_seq > 0 else {}

        request = self._client.build_request(
            "GET", url, headers=headers, params=params
        )
        response = await self._client.send(request, stream=True)

        if not response.is_success:
            try:
                await response.aread()
                raise_for_status(response)
            finally:
                await response.aclose()

        return AsyncExecutionStream(
            response,
            base_url=self._base_url,
            headers=self._headers,
            client=self._client,
        )

    async def upload_url(
        self,
        *,
        path: str,
        content_length: int,
        source: str = "thread",
        id: str,
        content_type: str | None = None,
    ) -> SignedUrl:
        """Generate a signed upload URL for direct file upload."""
        body: dict[str, object] = {
            "source": source,
            "id": id,
            "path": path,
            "content_length": content_length,
        }
        if content_type is not None:
            body["content_type"] = content_type
        response = await self._client.post(
            self._base_url + FILES_UPLOAD_URL_PATH,
            json=body,
            headers=self._headers,
        )
        raise_for_status(response)
        return SignedUrl.model_validate(response.json())

    async def download_url(
        self,
        *,
        path: str,
        source: str = "thread",
        id: str,
    ) -> SignedUrl:
        """Generate a signed download URL for direct file download."""
        response = await self._client.post(
            self._base_url + FILES_DOWNLOAD_URL_PATH,
            json={"source": source, "id": id, "path": path},
            headers=self._headers,
        )
        raise_for_status(response)
        return SignedUrl.model_validate(response.json())

    async def list_files(
        self,
        *,
        source: str = "thread",
        id: str,
        path: str | None = None,
    ) -> list[FileEntry]:
        """List files at a given path."""
        body: dict[str, object] = {"source": source, "id": id}
        if path is not None:
            body["path"] = path
        response = await self._client.post(
            self._base_url + FILES_LIST_PATH,
            json=body,
            headers=self._headers,
        )
        raise_for_status(response)
        data = response.json()
        return [FileEntry.model_validate(f) for f in data.get("files", [])]

    async def close(self) -> None:
        """Close the underlying HTTP client if internally created."""
        if self._owns_client:
            await self._client.aclose()

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

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        await self.close()
