"""Read capacity models for serverless indexes (2026-01.alpha API)."""

from __future__ import annotations

from typing import Any

from pinecone.models._base import BaseModel


class ReadCapacityOnDemandSpec(BaseModel, tag="OnDemand", tag_field="mode"):
    """On-demand read capacity specification (request model).

    On-demand capacity automatically scales based on usage. You pay only for
    what you use with no reserved capacity. Ideal for variable workloads.

    Example:
        >>> capacity = ReadCapacityOnDemandSpec()
    """

    # mode is implicitly "OnDemand" via tag


class ReadCapacityManualScaling(BaseModel):
    """Manual scaling configuration for dedicated read capacity.

    Fixed number of replicas and shards for predictable performance.

    Attributes:
        shards: Number of shards for horizontal scaling
        replicas: Number of replicas for high availability

    Example:
        >>> manual = ReadCapacityManualScaling(shards=1, replicas=1)
    """

    shards: int
    replicas: int


class ReadCapacityDedicatedInner(BaseModel):
    """Inner dedicated capacity configuration.

    Attributes:
        node_type: Node type ("b1", "s1", etc.)
        scaling: Scaling strategy ("Manual" or "Auto")
        manual: Manual scaling configuration (when scaling="Manual")
        auto: Auto scaling configuration (when scaling="Auto", optional)

    Example:
        >>> dedicated = ReadCapacityDedicatedInner(
        ...     node_type="b1",
        ...     scaling="Manual",
        ...     manual=ReadCapacityManualScaling(shards=1, replicas=1),
        ... )
    """

    node_type: str
    scaling: str  # Manual, Auto
    manual: ReadCapacityManualScaling | None = None
    auto: dict[str, Any] | None = None  # For future auto-scaling support


class ReadCapacityDedicatedSpec(BaseModel, tag="Dedicated", tag_field="mode"):
    """Dedicated read capacity specification (request model).

    Dedicated capacity provides reserved compute resources with predictable
    performance and pricing. You pay for reserved capacity regardless of usage.

    Attributes:
        dedicated: Dedicated capacity configuration

    Example:
        >>> capacity = ReadCapacityDedicatedSpec(
        ...     dedicated=ReadCapacityDedicatedInner(
        ...         node_type="b1",
        ...         scaling="Manual",
        ...         manual=ReadCapacityManualScaling(shards=1, replicas=1),
        ...     )
        ... )
    """

    dedicated: ReadCapacityDedicatedInner


class ReadCapacityStatus(BaseModel):
    """Read capacity status.

    Indicates the current state of read capacity provisioning.

    Attributes:
        state: Current state ("Ready", "Initializing", "Migrating", etc.)
        current_shards: Current number of shards (may be None during provisioning)
        current_replicas: Current number of replicas (may be None during provisioning)

    Example:
        >>> status = ReadCapacityStatus(state="Ready", current_shards=1, current_replicas=1)
    """

    state: str  # Ready, Initializing, Migrating, etc.
    current_shards: int | None = None
    current_replicas: int | None = None


class ReadCapacityDedicatedResponseInner(BaseModel):
    """Dedicated capacity details in API responses.

    Describes the current state of dedicated read capacity.

    Attributes:
        node_type: Node type ("b1", "s1", etc.)
        scaling: Scaling strategy ("Manual" or "Auto")
        manual: Manual scaling configuration (when scaling="Manual")
        auto: Auto scaling configuration (when scaling="Auto", optional)

    Example:
        >>> dedicated = ReadCapacityDedicatedResponseInner(
        ...     node_type="b1",
        ...     scaling="Manual",
        ...     manual=ReadCapacityManualScaling(shards=1, replicas=1),
        ... )
    """

    node_type: str
    scaling: str  # Manual, Auto
    manual: ReadCapacityManualScaling | None = None
    auto: dict[str, Any] | None = None  # For future auto-scaling support


class ReadCapacityDedicatedResponse(BaseModel, tag="Dedicated", tag_field="mode"):
    """Dedicated capacity in API responses.

    Describes the current state of dedicated read capacity for a serverless index.

    Attributes:
        dedicated: Dedicated capacity details
        status: Provisioning status

    Note:
        The "mode" field is automatically set to "Dedicated" by msgspec's
        tagged union system and should not be included in the struct definition.

    Example:
        >>> response = ReadCapacityDedicatedResponse(
        ...     dedicated=ReadCapacityDedicatedResponseInner(
        ...         node_type="b1",
        ...         scaling="Manual",
        ...         manual=ReadCapacityManualScaling(shards=1, replicas=1),
        ...     ),
        ...     status=ReadCapacityStatus(state="Ready"),
        ... )
    """

    dedicated: ReadCapacityDedicatedResponseInner
    status: ReadCapacityStatus


class ReadCapacityOnDemandResponse(BaseModel, tag="OnDemand", tag_field="mode"):
    """On-demand capacity in API responses.

    Describes the current state of on-demand read capacity for a serverless index.

    Attributes:
        status: Provisioning status

    Note:
        The "mode" field is automatically set to "OnDemand" by msgspec's
        tagged union system and should not be included in the struct definition.

    Example:
        >>> response = ReadCapacityOnDemandResponse(status=ReadCapacityStatus(state="Ready"))
    """

    status: ReadCapacityStatus


# Union type for API responses (can be OnDemand or Dedicated)
ReadCapacityResponse = ReadCapacityOnDemandResponse | ReadCapacityDedicatedResponse


# Union type for request models
ReadCapacitySpec = ReadCapacityOnDemandSpec | ReadCapacityDedicatedSpec


__all__ = [
    "ReadCapacityDedicatedInner",
    "ReadCapacityDedicatedResponse",
    "ReadCapacityDedicatedResponseInner",
    "ReadCapacityDedicatedSpec",
    "ReadCapacityManualScaling",
    "ReadCapacityOnDemandResponse",
    "ReadCapacityOnDemandSpec",
    "ReadCapacityResponse",
    "ReadCapacitySpec",
    "ReadCapacityStatus",
]
