import os
from typing import Any, Dict, Generic, List, Optional, Type

from llama_cloud.client import AsyncLlamaCloud
from tenacity import (
    WrappedFn,
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
)
import httpx

from .schema import (
    AgentDataT,
    ComparisonOperator,
    TypedAgentData,
    TypedAgentDataItems,
    TypedAggregateGroup,
    TypedAggregateGroupItems,
)


def agent_data_retry(func: WrappedFn) -> WrappedFn:
    """
    Decorator that adds automatic retry logic to agent data API calls.

    Applies exponential backoff retry strategy for common network-related exceptions:
    - Up to 3 retry attempts
    - Exponential wait time between 0.5s and 10s
    - Retries on timeout, connection, and HTTP status errors

    This ensures resilient API communication in distributed environments where
    temporary network issues or service unavailability may occur.
    """
    return retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(min=0.5, max=10),
        retry=retry_if_exception_type(
            (httpx.TimeoutException, httpx.ConnectError, httpx.HTTPStatusError)
        ),
    )(func)


def get_default_agent_id() -> str:
    """
    Retrieve the default agent ID from environment variables.

    Returns:
        The value of LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable,
        or None if not set

    Note:
        This provides a convenient way to configure agent ID globally
        via environment variables instead of passing it explicitly
        to each client instance.
    """
    return os.getenv("LLAMA_DEPLOY_DEPLOYMENT_NAME") or "_public"


class AsyncAgentDataClient(Generic[AgentDataT]):
    """
    Async client for managing agent-generated structured data with type safety.

    This client provides a high-level interface for CRUD operations, searching, and
    aggregation of structured data created by agents. It enforces type safety by
    validating all data against a specified Pydantic model type.

    The client is generic over AgentDataT, which must be a Pydantic BaseModel that
    defines the structure of your agent's data output.

    Example:
        ```python
        from pydantic import BaseModel
        from llama_cloud.client import AsyncLlamaCloud
        from llama_cloud_services.beta.agent_data import AsyncAgentDataClient

        class ExtractedPerson(BaseModel):
            name: str
            age: int
            email: str

        # Initialize client
        llama_client = AsyncLlamaCloud(token="your-api-key")
        agent_client = AsyncAgentDataClient(
            client=llama_client,
            type=ExtractedPerson,
            collection="extracted_people",
            agent_url_id="person-extraction-agent"
        )

        # Create data
        person = ExtractedPerson(name="John Doe", age=30, email="john@example.com")
        result = await agent_client.create_agent_data(person)

        # Search data
        results = await agent_client.search(
            filter={"age": {"gt": 25}},
            order_by="data.name",
            page_size=20
        )
        ```

    Type Parameters:
        AgentDataT: Pydantic BaseModel type that defines the structure of agent data
    """

    def __init__(
        self,
        type: Type[AgentDataT],
        collection: str = "default",
        agent_url_id: Optional[str] = None,
        client: Optional[AsyncLlamaCloud] = None,
        token: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        """
        Initialize the AsyncAgentDataClient.

        Args:
            type: Pydantic BaseModel class that defines the data structure.
                All agent data will be validated against this type.
            collection: Named collection within the agent for organizing data.
                Defaults to "default". Collections allow logical separation of
                different data types or workflows within the same agent.
            agent_url_id: Unique identifier for the agent. This normally appears in the
                url of an agent within the llama cloud platform. If not provided,
                will attempt to use the LLAMA_DEPLOY_DEPLOYMENT_NAME environment
                variable. Data can only be added to an already existing agent in the
                platform.
            client: AsyncLlamaCloud client instance for API communication. If not provided, will
                construct one from the provided api token and base url
            token: Llama Cloud API token. Reads from LLAMA_CLOUD_API_KEY if not provided
            base_url: Llama Cloud API token. Reads from LLAMA_CLOUD_BASE_URL if not provided, and
                defaults to https://api.cloud.llamaindex.ai

        Raises:
            ValueError: If agent_url_id is not provided and the
                LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable is not set

        Note:
            The client automatically applies retry logic to all API calls with
            exponential backoff for timeout, connection, and HTTP status errors.
        """

        self.agent_url_id = agent_url_id or get_default_agent_id()

        self.collection = collection
        if not client:
            client = AsyncLlamaCloud(
                token=token or os.getenv("LLAMA_CLOUD_API_KEY"),
                base_url=base_url or os.getenv("LLAMA_CLOUD_BASE_URL"),
            )
        self.client = client
        self.type = type

    @agent_data_retry
    async def get_item(self, item_id: str) -> TypedAgentData[AgentDataT]:
        raw_data = await self.client.beta.get_agent_data(
            item_id=item_id,
        )
        return TypedAgentData.from_raw(raw_data, validator=self.type)

    @agent_data_retry
    async def create_item(self, data: AgentDataT) -> TypedAgentData[AgentDataT]:
        raw_data = await self.client.beta.create_agent_data(
            agent_slug=self.agent_url_id,
            collection=self.collection,
            data=data.model_dump(),
        )
        return TypedAgentData.from_raw(raw_data, validator=self.type)

    @agent_data_retry
    async def update_item(
        self, item_id: str, data: AgentDataT
    ) -> TypedAgentData[AgentDataT]:
        raw_data = await self.client.beta.update_agent_data(
            item_id=item_id,
            data=data.model_dump(),
        )
        return TypedAgentData.from_raw(raw_data, validator=self.type)

    @agent_data_retry
    async def delete_item(self, item_id: str) -> None:
        await self.client.beta.delete_agent_data(item_id=item_id)

    @agent_data_retry
    async def search(
        self,
        filter: Optional[Dict[str, Dict[ComparisonOperator, Any]]] = None,
        order_by: Optional[str] = None,
        offset: Optional[int] = None,
        page_size: Optional[int] = None,
        include_total: bool = False,
    ) -> TypedAgentDataItems[AgentDataT]:
        """
        Search agent data with filtering, sorting, and pagination.
        Args:
            filter: Filter conditions to apply to the search. Dict mapping field names to FilterOperation objects. Filters only by data fields
                Examples:
                - {"age": {"gt": 18}} - age greater than 18
                - {"status": {"eq": "active"}} - status equals "active"
                - {"tags": {"includes": ["python", "ml"]}} - tags include "python" or "ml"
                - {"created_at": {"gte": "2024-01-01"}} - created after date
                - {"score": {"lt": 100, "gte": 50}} - score between 50 and 100
            order_by: Comma delimited list of fields to sort results by. Can order by standard agent fields like created_at, or by data fields. Data fields must be prefixed with "data.". If ordering desceding, use a " desc" suffix.
                Examples:
                - "data.name desc, created_at" - sort by name in descending order, and then by creation date
            page_size: Maximum number of items to return per page. Defaults to 10.
            offset: Number of items to skip from the beginning. Defaults to 0.
            include_total: Whether to include the total count in the response. Defaults to False to improve performance. It's recommended to only request on the first page.
        """
        raw = await self.client.beta.search_agent_data_api_v_1_beta_agent_data_search_post(
            agent_slug=self.agent_url_id,
            collection=self.collection,
            filter=filter,
            order_by=order_by,
            offset=offset,
            page_size=page_size,
            include_total=include_total,
        )
        return TypedAgentDataItems(
            items=[
                TypedAgentData.from_raw(item, validator=self.type) for item in raw.items
            ],
            has_more=raw.next_page_token is not None,
            total=raw.total_size,
        )

    @agent_data_retry
    async def aggregate(
        self,
        filter: Optional[Dict[str, Dict[ComparisonOperator, Any]]] = None,
        group_by: Optional[List[str]] = None,
        count: Optional[bool] = None,
        first: Optional[bool] = None,
        order_by: Optional[str] = None,
        offset: Optional[int] = None,
        page_size: Optional[int] = None,
    ) -> TypedAggregateGroupItems[AgentDataT]:
        """
        Aggregate agent data into groups according to the group_by fields.
        Args:
            filter: Filter conditions to apply to the search. Dict mapping field names to FilterOperation objects. Filters only by data fields
                See search for more details on filtering.
            group_by: List of fields to group by. Groups strictly by equality. Can only group by data fields.
                Examples:
                - ["name"] - group by name
                - ["name", "age"] - group by name and age
            count: Whether to include the count of items in each group.
            first: Whether to include the first item in each group.
            order_by: Comma delimited list of fields to sort results by. See search for more details on ordering.
            offset: Number of groups to skip from the beginning. Defaults to 0.
            page_size: Maximum number of groups to return per page.
        """
        raw = await self.client.beta.aggregate_agent_data_api_v_1_beta_agent_data_aggregate_post(
            agent_slug=self.agent_url_id,
            collection=self.collection,
            page_size=page_size,
            filter=filter,
            order_by=order_by,
            group_by=group_by,
            count=count,
            first=first,
            offset=offset,
        )
        return TypedAggregateGroupItems(
            items=[
                TypedAggregateGroup.from_raw(item, validator=self.type)
                for item in raw.items
            ],
            has_more=raw.next_page_token is not None,
            total=raw.total_size,
        )
