from typing import Annotated, Optional

from pydantic import AfterValidator, BaseModel, Field

from nuclia_models.worker.triggers import NameOperationFilter, Payload


def validate_filters(
    filters: Optional[list[NameOperationFilter]],
) -> Optional[list[NameOperationFilter]]:
    if filters is None:
        return []
    # Check that the filters are unique
    seen = set()
    for filter in filters:
        if filter.operation_type in seen:
            msg = f"Duplicate operation type for type `{filter.operation_type.name}` in filters"
            raise ValueError(msg)
        seen.add(filter.operation_type)
    return filters


FiltersType = Annotated[Optional[list[NameOperationFilter]], AfterValidator(validate_filters)]


class RunAgentsRequestBase(BaseModel):
    user_id: str = Field(..., title="The user ID of the user making the request")
    filters: FiltersType = Field(
        default_factory=list,  # type: ignore
        title="Filters to select which Data Augmentation Agents are applied to the text. If empty, \
        all configured agents for the Knowledge Box are applied.",
    )
    enable_webhooks: bool = Field(
        default=False,
        title="Whether to enable the triggering of the configured webhooks for each applied \
        Data Augmentation Agent",
    )


class RunTextAgentsRequest(RunAgentsRequestBase):
    texts: list[str] = Field(
        ...,
        title="The texts to run the agents on",
    )


class AugmentedBase(BaseModel):
    input_nuclia_tokens: float = Field(..., title="The number of input Nuclia tokens consumed for the field")
    output_nuclia_tokens: float = Field(
        ..., title="The number of output Nuclia tokens consumed for the field"
    )
    time: float = Field(..., title="The time taken to execute the Data Augmentation agents to the field")


class AugmentedText(AugmentedBase):
    payloads: list[Payload] = Field(
        ...,
        title="The list of payloads generated by the Data Augmentation agents for a given text",
    )


class RunTextAgentsResponse(BaseModel):
    results: list[AugmentedText] = Field(
        ...,
        title="Payloads generated by the agents for each text",
    )
