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

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

This module defines the Intent types used in VALIDATE and EXECUTE operations.
"""

from typing import Annotated, Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from .common import IdentityPredicate, PredicateExpression, ResourceId, SortOrder


class LookupIntent(BaseModel):
    """Intent for retrieving a single entity by its unique identifier."""

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

    intent_class: Literal["LOOKUP"] = Field(alias="intentClass")
    resource_id: ResourceId = Field(alias="resourceId", description="The resource to operate on.")
    key: IdentityPredicate = Field(description="The identity predicate specifying the unique key to lookup.")
    projections: list[str] | None = Field(default=None, description="Fields to project.")


class QueryIntent(BaseModel):
    """Intent for retrieving a set of entities based on criteria."""

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

    intent_class: Literal["QUERY"] = Field(alias="intentClass")
    resource_id: ResourceId = Field(alias="resourceId", description="The resource to operate on.")
    predicates: PredicateExpression = Field(description="Predicates for filtering data.")
    projections: list[str] | None = Field(default=None, description="Fields to project.")
    order_by: list[SortOrder] | None = Field(default=None, alias="orderBy", description="Fields to order results by.")
    limit: int | None = Field(default=None, description="The maximum number of results to return.")


class IngestIntent(BaseModel):
    """Intent for creating or appending new data."""

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

    intent_class: Literal["INGEST"] = Field(alias="intentClass")
    resource_id: ResourceId = Field(alias="resourceId", description="The resource to operate on.")
    payload: list[dict[str, Any]] = Field(
        description="The data payload to ingest. An array of records (key-value maps)."
    )


class ReviseIntent(BaseModel):
    """Intent for updating existing data."""

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

    intent_class: Literal["REVISE"] = Field(alias="intentClass")
    resource_id: ResourceId = Field(alias="resourceId", description="The resource to operate on.")
    predicates: PredicateExpression = Field(description="Predicates to identify the records to update.")
    payload: dict[str, Any] = Field(description="The data payload containing the fields to update.")


Intent = Annotated[
    LookupIntent | QueryIntent | IngestIntent | ReviseIntent,
    Field(discriminator="intent_class"),
]
