"""
Async client for the nono API.

This module provides a type-safe async client for interacting with the nono server.
"""

from typing import Any, Optional
from functools import cache
import static_refl.json as sj
import os
import httpx

from . import models


@cache
def _get_schema(t: type):
    """Get the static_refl schema for a given type."""
    return sj.Schema[t]


def serialize[T](t: type[T], value: T) -> Any:
    """Serialize a dataclass to a JSON-serializable dict."""
    schema = _get_schema(t)
    return schema.to_untyped(value)


def deserialize[T](t: type[T], data: Any) -> T:
    """Deserialize a JSON-serializable dict to a dataclass."""
    schema = _get_schema(t)
    return schema.to_typed(data)


class NonoClient:
    """
    Async client for the nono API.

    The client reads the server host and port from environment variables:
    - NONO_HOST: The host address (default: 127.0.0.1)
    - NONO_PORT: The port number (default: 8080)

    Example:
        ```python
        async with NonoClient() as client:
            config = await client.config_read()
            print(config.config)
        ```
    """

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
        auth_token: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        """
        Initialize the nono client.

        Args:
            host: Server host (default: from NONO_HOST env var or 127.0.0.1)
            port: Server port (default: from NONO_PORT env var or 8080)
            auth_token: Optional auth token for Bearer authentication
            base_url: Optional base URL (overrides host/port if provided)
        """
        if base_url:
            self.base_url = base_url
        else:
            _host = host or os.environ.get("NONO_HOST", "127.0.0.1")
            _port = port or int(os.environ.get("NONO_PORT", "8080"))
            self.base_url = f"http://{_host}:{_port}"
            print(f"Using nono server at {self.base_url}")

        self.auth_token = auth_token
        self._client: Optional[httpx.AsyncClient] = None

    async def __aenter__(self):
        """Async context manager entry."""
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            headers=self._get_headers(),
            timeout=3600.0,
        )
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        if self._client:
            await self._client.aclose()

    def _get_headers(self) -> dict:
        """Get request headers including auth if token is set."""
        headers = {"Content-Type": "application/json"}
        if self.auth_token:
            headers["Authorization"] = f"Bearer {self.auth_token}"
        return headers

    def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client."""
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                headers=self._get_headers(),
                timeout=3600.0,
            )
        return self._client

    async def close(self):
        """Close the HTTP client."""
        if self._client:
            await self._client.aclose()
            self._client = None

    # ========================================================================
    # Config API
    # ========================================================================

    async def config_read(self) -> models.ConfigReadResponse:
        """
        Read the server configuration.

        Returns:
            ConfigReadResponse with the configuration as JSON
        """
        client = self._get_client()
        response = await client.post("/v1/config/read", json={})
        response.raise_for_status()
        return deserialize(models.ConfigReadResponse, response.json())

    # ========================================================================
    # Entity API
    # ========================================================================

    async def entity_read(self, entity_id: str) -> Optional[Any]:
        """
        Read an entity by ID.

        Args:
            entity_id: The entity ID to read

        Returns:
            The entity data, or None if not found
        """
        client = self._get_client()
        response = await client.post("/v1/entity/read", json={"id": entity_id})
        response.raise_for_status()
        result = deserialize(models.EntityReadResponse, response.json())
        return result.entity

    async def entity_write(self, entity_id: str, entity: Any) -> bool:
        """
        Write an entity.

        Args:
            entity_id: The entity ID
            entity: The entity data (any JSON-serializable value)

        Returns:
            True if successful
        """
        client = self._get_client()
        response = await client.post(
            "/v1/entity/write",
            json=serialize(
                models.EntityWriteRequest,
                models.EntityWriteRequest(id=entity_id, entity=entity),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.EntityWriteResponse, response.json())
        return result.success

    # ========================================================================
    # Session API
    # ========================================================================

    async def session_create(
        self,
        name: str,
        data: Any,
        frame: models.CallingFrame,
    ) -> str:
        """
        Create a new session.

        Args:
            name: Session name
            data: Session data (any JSON-serializable value)
            frame: Initial calling frame

        Returns:
            The session ID
        """
        client = self._get_client()
        response = await client.post(
            "/v1/session/create",
            json=serialize(
                models.SessionCreateRequest,
                models.SessionCreateRequest(name=name, data=data, frame=frame),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.SessionCreateResponse, response.json())
        return result.session_id

    async def session_delete(self, session_id: str) -> bool:
        """
        Delete a session.

        Args:
            session_id: The session ID to delete

        Returns:
            True if the session was deleted
        """
        client = self._get_client()
        response = await client.post(
            "/v1/session/delete",
            json=serialize(
                models.SessionDeleteRequest,
                models.SessionDeleteRequest(session_id=session_id),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.SessionDeleteResponse, response.json())
        return result.success

    async def session_next(self, session_id: str) -> bool:
        """
        Execute the next operation in the session.

        Args:
            session_id: The session ID

        Returns:
            True if the operation succeeded
        """
        client = self._get_client()
        response = await client.post(
            "/v1/session/next",
            json=serialize(
                models.SessionNextRequest,
                models.SessionNextRequest(session_id=session_id),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.SessionNextResponse, response.json())
        return result.success

    async def session_get_data(self, session_id: str) -> Optional[Any]:
        """
        Get session data.

        Args:
            session_id: The session ID

        Returns:
            The session data, or None if session not found
        """
        client = self._get_client()
        response = await client.post(
            "/v1/session/get-data",
            json=serialize(
                models.SessionGetDataRequest,
                models.SessionGetDataRequest(session_id=session_id),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.SessionGetDataResponse, response.json())
        return result.data

    # ========================================================================
    # Frame API
    # ========================================================================

    async def frame_push_value(self, session_id: str, value: Any) -> bool:
        """
        Push a value to the current frame's value stack.

        Args:
            session_id: The session ID
            value: The value to push (any JSON-serializable value)

        Returns:
            True if successful
        """
        client = self._get_client()
        response = await client.post(
            "/v1/frame/push-value",
            json=serialize(
                models.FramePushValueRequest,
                models.FramePushValueRequest(session_id=session_id, value=value),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.FramePushValueResponse, response.json())
        return result.success

    async def frame_push_op(
        self,
        session_id: str,
        value: Any,
        op: models.Operation,
    ) -> bool:
        """
        Push a value and operation to the current frame.

        Args:
            session_id: The session ID
            value: The value to push to the value stack
            op: The operation to push to the op stack

        Returns:
            True if successful
        """
        client = self._get_client()
        response = await client.post(
            "/v1/frame/push-op",
            json=serialize(
                models.FramePushOpRequest,
                models.FramePushOpRequest(session_id=session_id, value=value, op=op),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.FramePushOpResponse, response.json())
        return result.success

    async def frame_read_data(self, session_id: str) -> Optional[Any]:
        """
        Read the current frame's data.

        Args:
            session_id: The session ID

        Returns:
            The frame data, or None if not found
        """
        client = self._get_client()
        response = await client.post(
            "/v1/frame/read-data",
            json=serialize(
                models.FrameReadDataRequest,
                models.FrameReadDataRequest(session_id=session_id),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.FrameReadDataResponse, response.json())
        return result.data

    async def frame_write_data(self, session_id: str, data: Any) -> bool:
        """
        Write data to the current frame.

        Args:
            session_id: The session ID
            data: The data to write (any JSON-serializable value)

        Returns:
            True if successful
        """
        client = self._get_client()
        response = await client.post(
            "/v1/frame/write-data",
            json=serialize(
                models.FrameWriteDataRequest,
                models.FrameWriteDataRequest(session_id=session_id, data=data),
            ),
        )
        response.raise_for_status()
        result = deserialize(models.FrameWriteDataResponse, response.json())
        return result.success
