# Copyright (c) Metis. All rights reserved.

"""Deterministic variant assignment for A/B experiments.

This module provides the core hashing logic that ensures:
1. The same user always gets the same variant (deterministic)
2. Assignment is uniformly distributed across variants
3. Traffic allocation is respected

The algorithm:
1. Hash (experiment_key + user_id) → deterministic bucket [0, 1)
2. If bucket >= traffic_percent/100 → user not in experiment
3. Otherwise, normalize bucket within traffic and assign to variant by weight
"""

import hashlib
from typing import Dict, List, Optional

from .types import ExperimentConfig, Variant


def compute_hash(experiment_key: str, user_id: str) -> float:
    """Compute a deterministic hash bucket for a user in an experiment.

    Uses SHA256 for high-quality uniform distribution. Takes the first
    8 hex characters (32 bits) and normalizes to [0, 1).

    Args:
        experiment_key: The experiment identifier
        user_id: The user/subject identifier

    Returns:
        A float in [0, 1) that is deterministic for this (experiment, user) pair

    Example::

        bucket = compute_hash("prompt-test-v1", "user-alice-789")
        # Always returns the same value for this pair
    """
    raw = hashlib.sha256(f"{experiment_key}:{user_id}".encode()).hexdigest()
    # Use first 8 hex chars (32 bits) for good distribution.
    # Divide by 2^32 to produce a value in [0, 1).
    return int(raw[:8], 16) / 0x100000000


def assign_variant(config: ExperimentConfig, user_id: str) -> Optional[Variant]:
    """Assign a user to a variant in an experiment.

    This is the core assignment algorithm:
    1. Compute deterministic bucket for (experiment, user)
    2. Check if user is in traffic allocation
    3. If yes, assign to variant based on cumulative weights

    Args:
        config: The experiment configuration with variants and traffic %
        user_id: The user/subject identifier

    Returns:
        The assigned Variant, or None if user is not in the experiment
        (either due to traffic allocation or empty variants)

    Example::

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

        variant = assign_variant(config, "user-alice-789")
        # Returns Variant(key="control", ...) or Variant(key="treatment", ...)
    """
    if not config.variants:
        return None

    bucket = compute_hash(config.key, user_id)

    # Check traffic allocation
    # If bucket >= traffic_percent/100, user is not in experiment
    traffic_threshold = config.traffic_percent / 100.0
    if bucket >= traffic_threshold:
        return None

    # Normalize bucket within the traffic allocation
    # This ensures users in the experiment are evenly distributed across variants
    normalized_bucket = bucket / traffic_threshold

    # Assign to variant based on cumulative weights
    total_weight = sum(v.weight for v in config.variants)
    if total_weight <= 0:
        return config.variants[0]  # Fallback to first variant

    cumulative = 0.0
    for variant in config.variants:
        cumulative += variant.weight / total_weight
        if normalized_bucket < cumulative:
            return variant

    # Fallback (should not happen due to floating point, but just in case)
    return config.variants[-1]


def compute_balanced_assignments(
    config: ExperimentConfig,
    user_ids: List[str],
) -> Dict[str, Optional[Variant]]:
    """Compute balanced variant assignments for a known set of users.

    Unlike ``assign_variant`` which maps each user's raw hash bucket to a
    variant (causing unequal splits at small n), this function sorts all
    users by their hash value and divides the sorted list proportionally
    across variants according to their weights.

    This guarantees the split is as close to the declared weights as
    possible for any n, while remaining fully deterministic — the same
    (experiment_key, user_ids) pair always produces the same assignments.

    Requires knowing all users upfront (offline/batch use only).  In
    production where users arrive one at a time, use ``assign_variant``
    instead.

    Args:
        config: Experiment configuration with variants and traffic percent.
        user_ids: The complete list of users to assign.

    Returns:
        Dict mapping each user_id to their assigned Variant, or None if the
        user falls outside the traffic allocation.

    Example::

        assignments = compute_balanced_assignments(config, ["alice", "bob", "charlie", "diana"])
        # Guaranteed: 2 control, 2 treatment (for equal weights)
    """
    if not config.variants or not user_ids:
        return {uid: None for uid in user_ids}

    traffic_threshold = config.traffic_percent / 100.0

    # Compute hashes once per user.
    user_hashes = [(uid, compute_hash(config.key, uid)) for uid in user_ids]

    # Separate users inside/outside traffic allocation.
    in_traffic = sorted(
        [(uid, h) for uid, h in user_hashes if h < traffic_threshold],
        key=lambda x: x[1],
    )
    out_of_traffic = {uid for uid, h in user_hashes if h >= traffic_threshold}

    assignments: Dict[str, Optional[Variant]] = {uid: None for uid in out_of_traffic}

    if not in_traffic:
        return assignments

    total_weight = sum(v.weight for v in config.variants)
    if total_weight <= 0:
        for uid, _ in in_traffic:
            assignments[uid] = config.variants[0]
        return assignments

    n = len(in_traffic)
    idx = 0

    for i, variant in enumerate(config.variants):
        if i < len(config.variants) - 1:
            count = round(n * variant.weight / total_weight)
        else:
            count = n - idx  # absorb any floating-point remainder

        for uid, _ in in_traffic[idx : idx + count]:
            assignments[uid] = variant
        idx += count

    return assignments
