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

"""
Common types for the Agentic Data Protocol (ADP).

This module defines the common types used across ADP requests and responses.
"""

from typing import Any, Literal, Union

from pydantic import BaseModel, ConfigDict
from pydantic import Field as PydanticField

# Type aliases
ResourceId = str
TraceId = str

IntentClass = Literal["LOOKUP", "QUERY", "INGEST", "REVISE", "*"]

PredicateOperator = Literal[
    "EQ",  # Equal
    "NEQ",  # Not equal
    "GT",  # Greater than
    "LT",  # Less than
    "GTE",  # Greater than or equal
    "LTE",  # Less than or equal
    "CONTAINS",  # Contains substring or element
    "IN",  # Value in set
    "LIKE",  # Simple pattern match (case-sensitive)
    "ILIKE",  # Simple pattern match (case-insensitive)
    "SIMILAR",  # Vector similarity search
]

LogicOperator = Literal["AND", "OR", "NOT"]

IssueSeverity = Literal["BLOCKING", "WARNING"]

ConsistencyLevel = Literal["STRONG", "EVENTUAL"]

FieldType = Literal[
    "STRING",
    "INTEGER",
    "FLOAT",
    "BOOLEAN",
    "DATE",
    "TIMESTAMP",
    "VECTOR",
    "BLOB",
    "JSON",
]

PredicateUsage = Literal["REQUIRED", "OPTIONAL"]

ValidationIssueCode = Literal[
    "MISSING_REQUIRED_PREDICATE",
    "INVALID_FORMAT",
    "FIELD_NOT_PERMITTED",
    "FIELD_NOT_FOUND",
    "INVALID_OPERATOR",
    "INVALID_VALUE",
    "CARDINALITY_EXCEEDED",
]


class SimilarValue(BaseModel):
    """Structured value for SIMILAR operator (vector similarity search)."""

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

    text: str | None = PydanticField(default=None, description="Text content to search for similarity.")
    blob: str | None = PydanticField(default=None, description="Binary content as Base64 or URI reference.")
    top: int | None = PydanticField(default=None, description="Maximum number of results to return.")
    threshold: float | None = PydanticField(default=None, description="Similarity threshold (0.0 to 1.0).")
    distance_function: str | None = PydanticField(
        default=None, alias="distanceFunction", description="Distance function to use."
    )
    vector: list[float] | None = PydanticField(
        default=None, description="Pre-computed embedding vector for direct vector similarity search."
    )


PredicateValue = str | int | float | bool | list[str | int | float | bool] | SimilarValue


class Predicate(BaseModel):
    """A single predicate in an Intent IR."""

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

    field_id: str = PydanticField(alias="fieldId", description="The field identifier to filter on.")
    op: PredicateOperator = PydanticField(description="The comparison operator.")
    value: PredicateValue = PydanticField(description="The value to compare against.")


class IdentityPredicate(BaseModel):
    """An identity predicate used for unique key lookups."""

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

    field_id: str = PydanticField(alias="fieldId", description="The field identifier to filter on.")
    op: Literal["EQ"] = PydanticField(description="The comparison operator. Must be EQ for identity lookups.")
    value: str | int | float | bool = PydanticField(description="The value to compare against.")


class PredicateGroup(BaseModel):
    """A group of predicates with a logic operator."""

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

    op: LogicOperator = PydanticField(description="The logic operator to use.")
    predicates: list[Union["Predicate", "PredicateGroup"]] = PydanticField(description="The predicates to use.")


# Resolve forward reference
PredicateGroup.model_rebuild()

PredicateExpression = Predicate | PredicateGroup
"""A predicate expression that can be either a single predicate or a group of predicates."""


class Implementation(BaseModel):
    """Describes the ADP implementation."""

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

    name: str = PydanticField(description="Name of the implementation.")
    version: str = PydanticField(description="Version of the implementation.")


class ClientCapabilities(BaseModel):
    """Capabilities a client may support."""

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

    experimental: dict[str, dict[str, Any]] | None = PydanticField(
        default=None, description="Experimental, non-standard capabilities that the client supports."
    )


class ServerCapabilities(BaseModel):
    """Capabilities that a server may support."""

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

    experimental: dict[str, dict[str, Any]] | None = PydanticField(
        default=None, description="Experimental, non-standard capabilities that the server supports."
    )
    supported_intent_classes: list[IntentClass] | None = PydanticField(
        default=None, alias="supportedIntentClasses", description="Supported intent classes by this server."
    )


class Resource(BaseModel):
    """Base resource information."""

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

    resource_id: ResourceId = PydanticField(
        alias="resourceId", description="The unique domain-qualified identifier for this resource."
    )
    version: int = PydanticField(description="The version number of this resource's schema.")
    intent_classes: list[IntentClass] = PydanticField(
        alias="intentClasses", description="The intent classes this resource supports."
    )
    description: str | None = PydanticField(default=None, description="A brief description of the resource.")
    semantic_description: str | None = PydanticField(
        default=None, alias="semanticDescription", description="A detailed semantic description of the resource."
    )
    tags: list[str] | None = PydanticField(default=None, description="Tags for categorization and filtering.")


class VectorMetadata(BaseModel):
    """Metadata for VECTOR fields."""

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

    dimensions: int = PydanticField(description="Number of dimensions in the vector.")
    distance_function: str | None = PydanticField(
        default=None, alias="distanceFunction", description="Distance function to use for similarity search."
    )


class FieldMetadata(BaseModel):
    """Metadata about a field's characteristics."""

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

    cardinality: int | str | None = PydanticField(default=None, description="Number of unique values in this field.")
    format: str | None = PydanticField(default=None, description="Expected format pattern.")
    whitelist_only: bool | None = PydanticField(
        default=None, alias="whitelistOnly", description="If true, only values from provided samples are valid."
    )
    hint: str | None = PydanticField(default=None, description="A hint for the Agent about how to use this field.")
    vector: VectorMetadata | None = PydanticField(
        default=None, description="For VECTOR fields, specifies vector dimensions and distance function."
    )
    samples: list[Any] | None = PydanticField(default=None, description="Sample values for this field.")


class FieldDefinition(BaseModel):
    """Definition of a single field in a resource."""

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

    field_id: str = PydanticField(alias="fieldId", description="The field identifier.")
    type: FieldType | None = PydanticField(default=None, description="The data type of this field.")
    description: str | None = PydanticField(default=None, description="Human-readable description of the field.")
    samples: list[Any] | None = PydanticField(default=None, description="Sample values for this field.")
    is_masked: bool | None = PydanticField(
        default=None, alias="isMasked", description="If true, this field is masked/redacted by policy."
    )
    is_searchable: bool | None = PydanticField(
        default=None, alias="isSearchable", description="If true, this field supports similarity search."
    )
    metadata: FieldMetadata | None = PydanticField(default=None, description="Additional metadata about the field.")


# Alias for backward compatibility
Field = FieldDefinition


class PredicateCapability(BaseModel):
    """Definition of a predicate capability for a field."""

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

    field_id: str = PydanticField(alias="fieldId", description="The field this predicate applies to.")
    usage: PredicateUsage = PydanticField(description="Whether this predicate is required or optional.")
    operators: list[PredicateOperator] = PydanticField(description="Allowed operators for this predicate.")


class ProjectionCapability(BaseModel):
    """Definition of a projection capability."""

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

    field_id: str = PydanticField(alias="fieldId", description="The field that can be projected.")


class MutableCapability(BaseModel):
    """Definition of a mutable field capability (for WRITE operations)."""

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

    field_id: str = PydanticField(alias="fieldId", description="The field that can be mutated.")
    constraints: dict[str, Any] | None = PydanticField(default=None, description="Constraints on the mutation.")


class Capabilities(BaseModel):
    """Capabilities available for a resource."""

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

    predicates: list[PredicateCapability] | None = PydanticField(
        default=None, description="Available predicate capabilities."
    )
    projections: list[ProjectionCapability] | None = PydanticField(
        default=None, description="Available projection capabilities."
    )
    mutables: list[MutableCapability] | None = PydanticField(
        default=None, description="Available mutable field capabilities (for WRITE operations)."
    )


class UsageContract(BaseModel):
    """The usage contract returned by DESCRIBE."""

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

    fields: list[FieldDefinition] = PydanticField(description="Field definitions for this resource.")
    capabilities: Capabilities = PydanticField(description="Capabilities available for this resource and intent.")


class SortOrder(BaseModel):
    """Sort order for a field."""

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

    direction: Literal["ASC", "DESC"] = PydanticField(description="The direction of the sort order.")
    field_id: str = PydanticField(alias="fieldId", description="The field to sort by.")


class ValidationIssue(BaseModel):
    """A single validation issue."""

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

    code: ValidationIssueCode = PydanticField(description="The type of validation issue.")
    field: str | None = PydanticField(default=None, description="The field associated with this issue.")
    severity: IssueSeverity = PydanticField(description="Severity of this issue.")
    message: str = PydanticField(description="Human-readable message describing the issue.")
    correction_hint: str | None = PydanticField(
        default=None, alias="correctionHint", description="Correction hint for the Agent."
    )


class ExecutionMetadata(BaseModel):
    """Metadata about execution performance and characteristics."""

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

    duration_ms: int | None = PydanticField(
        default=None, alias="durationMs", description="Time taken to execute the query in milliseconds."
    )
    source_system: str | None = PydanticField(
        default=None, alias="sourceSystem", description="The backend system that served the request."
    )
    consistency: ConsistencyLevel | None = PydanticField(
        default=None, description="Data consistency level of the result."
    )
