# Copyright (c) Metis. All rights reserved.

"""Data types for A/B experiments.

This module defines the core data structures for the experiments system:
- Variant: A single variant in an experiment with its config
- ExperimentConfig: Definition of an experiment with its variants
- Exposure: Record of a user being exposed to a variant
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List


@dataclass
class Variant:
    """A single variant in an A/B experiment.

    Attributes:
        key: Unique identifier for this variant (e.g., "control", "treatment")
        weight: Relative weight for traffic allocation (e.g., 50 = 50% of traffic)
        config: Arbitrary configuration dict passed to the user's code

    Example::

        variant = Variant(
            key="treatment",
            weight=50,
            config={"prompt": "You are a concise assistant..."}
        )
    """

    key: str
    weight: int
    config: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ExperimentConfig:
    """Configuration for an A/B experiment.

    Attributes:
        key: Unique experiment identifier (e.g., "prompt-test-v1")
        variants: List of variants in this experiment
        traffic_percent: Percentage of users to include (0-100)

    Example::

        config = ExperimentConfig(
            key="prompt-test-v1",
            variants=[
                Variant(key="control", weight=50, config={"prompt": "..."}),
                Variant(key="treatment", weight=50, config={"prompt": "..."}),
            ],
            traffic_percent=100,
        )
    """

    key: str
    variants: List[Variant]
    traffic_percent: int = 100


@dataclass
class Exposure:
    """Record of a user being exposed to an experiment variant.

    This is logged when get_variant() is called and the user is assigned
    to a variant. Exposures are batched and sent to the backend asynchronously.

    Attributes:
        user_id: The user/subject identifier
        experiment_id: The experiment key
        variant_key: The variant the user was assigned to
        timestamp: ISO 8601 timestamp of the exposure
    """

    user_id: str
    experiment_id: str
    variant_key: str
    timestamp: str
