# Copyright 2026 Datastrato, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""High-level ADP client session.

Manages the lifecycle of a connection to an ADP server, including
initialization, capability negotiation, and all protocol operations.
"""

from __future__ import annotations

import asyncio
import base64
import itertools
import logging
from importlib.metadata import version as _pkg_version
from types import TracebackType
from typing import Any

from adp_sdk.shared import (
    deserialize_json,
    get_exception_for_error_code,
    serialize_message,
)
from adp_sdk.transports import Transport
from adp_sdk.types import (
    LATEST_PROTOCOL_VERSION,
    SUPPORTED_PROTOCOL_VERSIONS,
    ClientCapabilities,
    DescribeRequestParams,
    DescribeResult,
    DiscoverFilter,
    DiscoverRequestParams,
    DiscoverResult,
    EmptyResult,
    ExecuteRequestParams,
    ExecuteResult,
    Implementation,
    InitializeRequestParams,
    InitializeResult,
    IntentClass,
    JSONRPCRequest,
    RequestParams,
    ServerCapabilities,
    ValidateRequestParams,
    ValidateResult,
)
from adp_sdk.types.intents import Intent

logger = logging.getLogger(__name__)

_SDK_NAME = "adp-sdk-python"
_SDK_VERSION = _pkg_version("adp-sdk")


def basic_auth(username: str, password: str = "") -> str:
    """Build a Basic Authorization header value.

    Args:
        username: The username. Must be non-empty and must not contain a colon.
        password: The password (defaults to empty string).

    Returns:
        A string in the format ``"Basic <base64(username:password)>"``.

    Raises:
        ValueError: If ``username`` is empty or contains a colon character.
    """
    if not username:
        raise ValueError("username must not be empty")
    if ":" in username:
        raise ValueError("username must not contain a colon character")
    credentials = base64.b64encode(f"{username}:{password}".encode()).decode()
    return f"Basic {credentials}"


class ClientSession:
    """High-level ADP client session.

    Wraps a :class:`Transport` and provides typed methods for every ADP
    protocol operation. Handles JSON-RPC request/response lifecycle,
    request ID management, and error mapping.

    Usage::

        async with ClientSession(transport) as session:
            resources = await session.discover()
            contract = await session.describe("com.example:users", "QUERY")
    """

    def __init__(
        self,
        transport: Transport,
        client_info: Implementation | None = None,
        capabilities: ClientCapabilities | None = None,
        authorization: str | None = None,
    ) -> None:
        self._transport = transport
        self._client_info = client_info or Implementation(
            name=_SDK_NAME,
            version=_SDK_VERSION,
        )
        self._capabilities = capabilities or ClientCapabilities()
        self._authorization = authorization
        self._id_counter = itertools.count(1)
        self._send_lock = asyncio.Lock()
        self._initialized = False
        self._server_info: Implementation | None = None
        self._server_capabilities: ServerCapabilities | None = None

    # -- Properties ----------------------------------------------------------

    @property
    def server_info(self) -> Implementation | None:
        """Server implementation info, available after initialization."""
        return self._server_info

    @property
    def server_capabilities(self) -> ServerCapabilities | None:
        """Server capabilities, available after initialization."""
        return self._server_capabilities

    @property
    def initialized(self) -> bool:
        """Whether the session has completed initialization."""
        return self._initialized

    # -- Lifecycle -----------------------------------------------------------

    async def initialize(self) -> InitializeResult:
        """Send ``adp.initialize`` and negotiate protocol version.

        Returns:
            The server's initialization result.

        Raises:
            ADPError: If the server returns an error response.
            RuntimeError: If the session is already initialized.
        """
        if self._initialized:
            raise RuntimeError("Session already initialized")

        params = InitializeRequestParams(
            protocolVersion=LATEST_PROTOCOL_VERSION,
            capabilities=self._capabilities,
            clientInfo=self._client_info,
        )
        result = await self._send_request("adp.initialize", params, InitializeResult)

        if result.protocol_version not in SUPPORTED_PROTOCOL_VERSIONS:
            raise RuntimeError(f"Unsupported protocol version from the server: {result.protocol_version}")

        self._server_info = result.server_info
        self._server_capabilities = result.capabilities
        self._initialized = True

        logger.info(
            "Initialized: server=%s/%s, protocol=%s",
            result.server_info.name,
            result.server_info.version,
            result.protocol_version,
        )
        return result

    async def close(self) -> None:
        """Close the transport connection."""
        self._initialized = False
        await self._transport.close()

    async def __aenter__(self) -> ClientSession:
        await self._transport.start()
        await self.initialize()
        return self

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

    # -- Protocol operations -------------------------------------------------

    async def ping(self) -> EmptyResult:
        """Send ``adp.ping`` to check server liveness.

        Returns:
            An empty result on success.
        """
        self._check_initialized()
        return await self._send_request("adp.ping", RequestParams(), EmptyResult)

    async def discover(
        self,
        filter: DiscoverFilter | None = None,
        cursor: str | None = None,
    ) -> DiscoverResult:
        """Send ``adp.discover`` to list available resources.

        Args:
            filter: Optional filter criteria.
            cursor: Pagination cursor from a previous response.

        Returns:
            The discovered resources.
        """
        self._check_initialized()
        params = DiscoverRequestParams(filter=filter, cursor=cursor)
        return await self._send_request("adp.discover", params, DiscoverResult)

    async def describe(
        self,
        resource_id: str,
        intent_class: IntentClass,
        version: int | None = None,
        cursor: str | None = None,
    ) -> DescribeResult:
        """Send ``adp.describe`` to get a resource's usage contract.

        Args:
            resource_id: The unique identifier of the resource.
            intent_class: The intent class to describe.
            version: Optional version of the resource.
            cursor: Pagination cursor from a previous response.

        Returns:
            The resource description and usage contract.
        """
        self._check_initialized()
        params = DescribeRequestParams(
            resourceId=resource_id,
            intentClass=intent_class,
            version=version,
            cursor=cursor,
        )
        return await self._send_request("adp.describe", params, DescribeResult)

    async def validate(
        self,
        intent: Intent,
    ) -> ValidateResult:
        """Send ``adp.validate`` to check if an intent is valid.

        Args:
            intent: The intent to validate (must include resourceId).

        Returns:
            The validation result.
        """
        self._check_initialized()
        params = ValidateRequestParams(
            intent=intent,
        )
        return await self._send_request("adp.validate", params, ValidateResult)

    async def execute(
        self,
        intent: Intent,
        cursor: str | None = None,
    ) -> ExecuteResult:
        """Send ``adp.execute`` to run an intent against a resource.

        Args:
            intent: The intent to execute (must include resourceId).
            cursor: Pagination cursor from a previous response.

        Returns:
            The execution result.
        """
        self._check_initialized()
        params = ExecuteRequestParams(
            intent=intent,
            cursor=cursor,
        )
        return await self._send_request("adp.execute", params, ExecuteResult)

    # -- Internal helpers ----------------------------------------------------

    def _check_initialized(self) -> None:
        """Raise if the session has not been initialized."""
        if not self._initialized:
            raise RuntimeError("Session not initialized; call initialize() or use 'async with'")

    async def _send_request(
        self,
        method: str,
        params: RequestParams,
        result_type: type[Any],
    ) -> Any:
        """Send a JSON-RPC request and return the parsed result.

        Serializes concurrent calls with an internal lock to prevent
        response-ID mismatches on the underlying transport.

        Args:
            method: The JSON-RPC method name (e.g. ``adp.discover``).
            params: The request parameters.
            result_type: The Pydantic model to parse the result into.

        Returns:
            The parsed result object.

        Raises:
            ADPError: If the server responds with an error.
            ConnectionError: If the transport connection is lost.
        """
        async with self._send_lock:
            request_id = next(self._id_counter)
            params_dict = params.model_dump(by_alias=True, exclude_none=True)

            if self._authorization is not None:
                meta = params_dict.setdefault("_meta", {})
                meta["authorization"] = self._authorization

            request = JSONRPCRequest(
                id=request_id,
                method=method,
                params=params_dict,
            )

            await self._transport.write_message(serialize_message(request))

            raw = await self._transport.read_message()
            if raw is None:
                raise ConnectionError("Transport connection lost (EOF)")

            data = deserialize_json(raw)
            self._handle_error(data)

            if "result" not in data:
                raise ConnectionError(f"Malformed response: missing 'result' key (id={data.get('id')})")

            return result_type.model_validate(data["result"])

    def _handle_error(self, data: dict[str, Any]) -> None:
        """Raise an ADPError if the response contains an error."""
        error = data.get("error")
        if error is None:
            return
        raise get_exception_for_error_code(
            code=error.get("code", -32603),
            message=error.get("message", "Unknown error"),
            data=error.get("data"),
        )
