# 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.

"""
JSON-RPC 2.0 types for the Agentic Data Protocol (ADP).

This module defines the base JSON-RPC types used throughout the ADP protocol.
"""

from typing import Annotated, Any, Literal

from pydantic import BaseModel, ConfigDict, Field

# Protocol constants
JSONRPC_VERSION: Literal["2.0"] = "2.0"
LATEST_PROTOCOL_VERSION = "2026-01-20"
SUPPORTED_PROTOCOL_VERSIONS = [LATEST_PROTOCOL_VERSION]

# Type aliases
RequestId = str | int
ProgressToken = str | int
Cursor = str


class RequestMeta(BaseModel):
    """Metadata for a request."""

    model_config = ConfigDict(extra="allow", populate_by_name=True)

    progress_token: ProgressToken | None = Field(default=None, alias="progressToken")


class RequestParams(BaseModel):
    """Common params for any request."""

    model_config = ConfigDict(extra="allow", populate_by_name=True)

    meta: RequestMeta | None = Field(default=None, alias="_meta")


class Result(BaseModel):
    """Common result type for all successful responses."""

    model_config = ConfigDict(extra="allow", populate_by_name=True)

    meta: dict[str, Any] | None = Field(default=None, alias="_meta")


class EmptyResult(Result):
    """A response that indicates success but carries no data."""

    pass


class JSONRPCError(BaseModel):
    """JSON-RPC error object."""

    code: int = Field(description="The error type that occurred.")
    message: str = Field(description="A short description of the error.")
    data: Any | None = Field(default=None, description="Additional information about the error.")


class JSONRPCRequest(BaseModel):
    """A JSON-RPC request that expects a response."""

    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    jsonrpc: Literal["2.0"] = JSONRPC_VERSION
    id: RequestId
    method: str
    params: dict[str, Any] | None = None


class JSONRPCResultResponse(BaseModel):
    """A successful (non-error) response to a request."""

    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    jsonrpc: Literal["2.0"] = JSONRPC_VERSION
    id: RequestId
    result: dict[str, Any]


class JSONRPCErrorResponse(BaseModel):
    """A response to a request that indicates an error occurred."""

    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    jsonrpc: Literal["2.0"] = JSONRPC_VERSION
    id: RequestId | None = None
    error: JSONRPCError


JSONRPCResponse = Annotated[
    JSONRPCResultResponse | JSONRPCErrorResponse,
    Field(discriminator=None),
]

JSONRPCMessage = JSONRPCRequest | JSONRPCResultResponse | JSONRPCErrorResponse


# Standard JSON-RPC error codes
PARSE_ERROR = -32700
INVALID_REQUEST = -32600
METHOD_NOT_FOUND = -32601
INVALID_PARAMS = -32602
INTERNAL_ERROR = -32603

# ADP-specific error codes (range -32000 to -32099)
RESOURCE_NOT_FOUND = -32001
VALIDATION_FAILED = -32002
UNAUTHORIZED = -32003
EXECUTION_FAILED = -32004
