client.py

"""Hawk SDK client implementations (sync and async)."""

from __future__ import annotations

from typing import Any
from urllib.parse import urlparse

import httpx

from ._version import __version__
from .errors import parse_error
from .retry import DEFAULT_RETRY_CONFIG, RetryConfig, with_retry, with_retry_sync
from .streaming import AsyncStreamReader, StreamReader
from .types import (
    ChatRequest,
    ChatResponse,
    HealthResponse,
    Message,
    PaginatedResponse,
    SessionDetail,
    SessionSummary,
    StatsResponse,
)

DEFAULT_BASE_URL = "http://127.0.0.1:4590"
DEFAULT_TIMEOUT = 30.0
DEFAULT_POOL_CONNECTIONS = 10
DEFAULT_POOL_MAXSIZE = 100


def _build_headers(api_key: str | None) -> dict[str, str]:
    """Build standard HTTP headers for Hawk API requests."""
    headers: dict[str, str] = {
        "Accept": "application/json",
        "User-Agent": f"hawk-sdk-python/{__version__}",
    }
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    return headers


def _validate_base_url(url: str) -> str:
    """Validate that base_url uses http/https and has a hostname."""
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise ValueError(f"base_url must use http or https scheme, got: {parsed.scheme}")
    if not parsed.hostname:
        raise ValueError("base_url must have a hostname")
    return url


class HawkClient:
    """Synchronous client for the Hawk daemon API.

    Usage:
        with HawkClient() as client:
            health = client.health()
            response = client.chat("Hello!")
    """

    def __init__(
        self,
        base_url: str = DEFAULT_BASE_URL,
        api_key: str | None = None,
        retry_config: RetryConfig | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        pool_connections: int = DEFAULT_POOL_CONNECTIONS,
        pool_maxsize: int = DEFAULT_POOL_MAXSIZE,
    ) -> None:
        _validate_base_url(base_url)
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._retry_config = retry_config or DEFAULT_RETRY_CONFIG
        self._timeout = timeout
        self._pool_connections = pool_connections
        self._pool_maxsize = pool_maxsize
        self._client = httpx.Client(
            base_url=self._base_url,
            timeout=self._timeout,
            headers=_build_headers(api_key),
            limits=httpx.Limits(
                max_keepalive_connections=pool_connections,
                max_connections=pool_maxsize,
            ),
        )

    def __repr__(self) -> str:
        masked = "***" + self._api_key[-4:] if self._api_key else "None"
        return f"HawkClient(base_url='{self._base_url}', api_key='{masked}')"

    def __enter__(self) -> HawkClient:
        return self

    def __exit__(self, *args: object) -> None:
        self.close()

    def close(self) -> None:
        """Close the underlying HTTP client."""
        self._client.close()

    def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
        """Make an HTTP request, raising typed errors on failure."""
        response = self._client.request(method, path, **kwargs)
        if response.status_code >= 400:
            raise parse_error(response)
        return response

    def health(self) -> HealthResponse:
        """Check daemon connectivity and health."""

        def _do() -> HealthResponse:
            resp = self._request("GET", "/v1/health")
            return HealthResponse.model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)

    def chat(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        model: str | None = None,
        max_turns: int | None = None,
        autonomy: str | None = None,
        cwd: str | None = None,
        agent: str | None = None,
    ) -> ChatResponse:
        """Send a prompt and return the complete response."""
        request = ChatRequest(
            prompt=prompt,
            session_id=session_id,
            model=model,
            max_turns=max_turns,
            autonomy=autonomy,
            cwd=cwd,
            agent=agent,
        )

        def _do() -> ChatResponse:
            resp = self._request(
                "POST",
                "/v1/chat",
                json=request.model_dump(exclude_none=True, by_alias=True),
            )
            return ChatResponse.model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)

    def chat_stream(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        model: str | None = None,
        max_turns: int | None = None,
        autonomy: str | None = None,
        cwd: str | None = None,
        agent: str | None = None,
    ) -> StreamReader:
        """Send a prompt and stream the response via SSE."""
        request = ChatRequest(
            prompt=prompt,
            session_id=session_id,
            model=model,
            max_turns=max_turns,
            autonomy=autonomy,
            cwd=cwd,
            agent=agent,
        )

        def _do() -> StreamReader:
            response = self._client.send(
                self._client.build_request(
                    "POST",
                    "/v1/chat",
                    json=request.model_dump(exclude_none=True, by_alias=True),
                    headers={"Accept": "text/event-stream"},
                ),
                stream=True,
            )

            if response.status_code >= 400:
                response.read()
                raise parse_error(response)

            return StreamReader(response)

        return with_retry_sync(_do, self._retry_config)

    def create_session(
        self,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> SessionDetail:
        """Create a new session."""
        body: dict[str, Any] = {}
        if name is not None:
            body["name"] = name
        if metadata is not None:
            body["metadata"] = metadata

        def _do() -> SessionDetail:
            resp = self._request("POST", "/v1/sessions", json=body)
            return SessionDetail.model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)

    def get_session(self, session_id: str) -> SessionDetail:
        """Get a session by ID."""

        def _do() -> SessionDetail:
            resp = self._request("GET", f"/v1/sessions/{session_id}")
            return SessionDetail.model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)

    def list_sessions(self, limit: int = 20, offset: int = 0) -> PaginatedResponse[SessionSummary]:
        """List sessions with pagination."""

        def _do() -> PaginatedResponse[SessionSummary]:
            params: dict[str, Any] = {"limit": limit}
            if offset > 0:
                params["offset"] = offset
            resp = self._request("GET", "/v1/sessions", params=params)
            return PaginatedResponse[SessionSummary].model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)

    def delete_session(self, session_id: str) -> None:
        """Delete a session by ID."""

        def _do() -> None:
            self._request("DELETE", f"/v1/sessions/{session_id}")

        return with_retry_sync(_do, self._retry_config)

    def list_messages(
        self, session_id: str, limit: int = 50, offset: int = 0
    ) -> PaginatedResponse[Message]:
        """List messages for a session with pagination."""

        def _do() -> PaginatedResponse[Message]:
            params: dict[str, Any] = {"limit": limit}
            if offset > 0:
                params["offset"] = offset
            resp = self._request("GET", f"/v1/sessions/{session_id}/messages", params=params)
            return PaginatedResponse[Message].model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)

    def stats(self) -> StatsResponse:
        """Get aggregated usage statistics."""

        def _do() -> StatsResponse:
            resp = self._request("GET", "/v1/stats")
            return StatsResponse.model_validate(resp.json())

        return with_retry_sync(_do, self._retry_config)


class AsyncHawkClient:
    """Asynchronous client for the Hawk daemon API.

    Usage:
        async with AsyncHawkClient() as client:
            health = await client.health()
            response = await client.chat("Hello!")
    """

    def __init__(
        self,
        base_url: str = DEFAULT_BASE_URL,
        api_key: str | None = None,
        retry_config: RetryConfig | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        pool_connections: int = DEFAULT_POOL_CONNECTIONS,
        pool_maxsize: int = DEFAULT_POOL_MAXSIZE,
    ) -> None:
        _validate_base_url(base_url)
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._retry_config = retry_config or DEFAULT_RETRY_CONFIG
        self._timeout = timeout
        self._pool_connections = pool_connections
        self._pool_maxsize = pool_maxsize
        self._client = httpx.AsyncClient(
            base_url=self._base_url,
            timeout=self._timeout,
            headers=_build_headers(api_key),
            limits=httpx.Limits(
                max_keepalive_connections=pool_connections,
                max_connections=pool_maxsize,
            ),
        )

    def __repr__(self) -> str:
        masked = "***" + self._api_key[-4:] if self._api_key else "None"
        return f"AsyncHawkClient(base_url='{self._base_url}', api_key='{masked}')"

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

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

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

    async def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
        """Make an async HTTP request, raising typed errors on failure."""
        response = await self._client.request(method, path, **kwargs)
        if response.status_code >= 400:
            raise parse_error(response)
        return response

    async def health(self) -> HealthResponse:
        """Check daemon connectivity and health."""

        async def _do() -> HealthResponse:
            resp = await self._request("GET", "/v1/health")
            return HealthResponse.model_validate(resp.json())

        return await with_retry(_do, self._retry_config)

    async def chat(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        model: str | None = None,
        max_turns: int | None = None,
        autonomy: str | None = None,
        cwd: str | None = None,
        agent: str | None = None,
    ) -> ChatResponse:
        """Send a prompt and return the complete response."""
        request = ChatRequest(
            prompt=prompt,
            session_id=session_id,
            model=model,
            max_turns=max_turns,
            autonomy=autonomy,
            cwd=cwd,
            agent=agent,
        )

        async def _do() -> ChatResponse:
            resp = await self._request(
                "POST",
                "/v1/chat",
                json=request.model_dump(exclude_none=True, by_alias=True),
            )
            return ChatResponse.model_validate(resp.json())

        return await with_retry(_do, self._retry_config)

    async def chat_stream(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        model: str | None = None,
        max_turns: int | None = None,
        autonomy: str | None = None,
        cwd: str | None = None,
        agent: str | None = None,
    ) -> AsyncStreamReader:
        """Send a prompt and stream the response via SSE."""
        request = ChatRequest(
            prompt=prompt,
            session_id=session_id,
            model=model,
            max_turns=max_turns,
            autonomy=autonomy,
            cwd=cwd,
            agent=agent,
        )

        async def _do() -> AsyncStreamReader:
            response = await self._client.send(
                self._client.build_request(
                    "POST",
                    "/v1/chat",
                    json=request.model_dump(exclude_none=True, by_alias=True),
                    headers={"Accept": "text/event-stream"},
                ),
                stream=True,
            )

            if response.status_code >= 400:
                await response.aread()
                raise parse_error(response)

            return AsyncStreamReader(response)

        return await with_retry(_do, self._retry_config)

    async def create_session(
        self,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> SessionDetail:
        """Create a new session."""
        body: dict[str, Any] = {}
        if name is not None:
            body["name"] = name
        if metadata is not None:
            body["metadata"] = metadata

        async def _do() -> SessionDetail:
            resp = await self._request("POST", "/v1/sessions", json=body)
            return SessionDetail.model_validate(resp.json())

        return await with_retry(_do, self._retry_config)

    async def get_session(self, session_id: str) -> SessionDetail:
        """Get a session by ID."""

        async def _do() -> SessionDetail:
            resp = await self._request("GET", f"/v1/sessions/{session_id}")
            return SessionDetail.model_validate(resp.json())

        return await with_retry(_do, self._retry_config)

    async def list_sessions(
        self, limit: int = 20, offset: int = 0
    ) -> PaginatedResponse[SessionSummary]:
        """List sessions with pagination."""

        async def _do() -> PaginatedResponse[SessionSummary]:
            params: dict[str, Any] = {"limit": limit}
            if offset > 0:
                params["offset"] = offset
            resp = await self._request("GET", "/v1/sessions", params=params)
            return PaginatedResponse[SessionSummary].model_validate(resp.json())

        return await with_retry(_do, self._retry_config)

    async def delete_session(self, session_id: str) -> None:
        """Delete a session by ID."""

        async def _do() -> None:
            await self._request("DELETE", f"/v1/sessions/{session_id}")

        return await with_retry(_do, self._retry_config)

    async def list_messages(
        self, session_id: str, limit: int = 50, offset: int = 0
    ) -> PaginatedResponse[Message]:
        """List messages for a session with pagination."""

        async def _do() -> PaginatedResponse[Message]:
            params: dict[str, Any] = {"limit": limit}
            if offset > 0:
                params["offset"] = offset
            resp = await self._request("GET", f"/v1/sessions/{session_id}/messages", params=params)
            return PaginatedResponse[Message].model_validate(resp.json())

        return await with_retry(_do, self._retry_config)

    async def stats(self) -> StatsResponse:
        """Get aggregated usage statistics."""

        async def _do() -> StatsResponse:
            resp = await self._request("GET", "/v1/stats")
            return StatsResponse.model_validate(resp.json())

        return await with_retry(_do, self._retry_config)