#!/usr/bin/env python
"""
A/B Testing Framework
Manages feature flags and variant assignment for experiments.
"""

import hashlib
import logging
from typing import Dict, Any, Optional, List
from enum import Enum

logger = logging.getLogger(__name__)


class Variant(Enum):
    """Standard variant types for A/B tests."""
    CONTROL = "control"
    TREATMENT = "treatment"
    TREATMENT_A = "treatment_a"
    TREATMENT_B = "treatment_b"


# Feature Flags Configuration
FEATURE_FLAGS = {
    # Tutorial experiments
    "tutorial_flow": {
        "enabled": True,
        "variants": {
            "control": 0.5,      # 50% get original tutorial
            "treatment": 0.5     # 50% get new tutorial flow
        },
        "description": "Test new tutorial flow effectiveness"
    },

    # Monetization experiments
    "premium_pricing": {
        "enabled": True,
        "variants": {
            "control": 0.33,     # $4.99
            "treatment_a": 0.33,  # $2.99
            "treatment_b": 0.34   # $6.99
        },
        "description": "Test different premium subscription prices"
    },

    # Dating system experiments
    "dating_match_algorithm": {
        "enabled": False,
        "variants": {
            "control": 0.5,      # Original compatibility algorithm
            "treatment": 0.5     # Enhanced AI-based matching
        },
        "description": "Test enhanced dating match algorithm"
    },

    # Retention experiments
    "daily_reward_amount": {
        "enabled": True,
        "variants": {
            "control": 0.5,      # Standard rewards
            "treatment": 0.5     # Increased rewards
        },
        "description": "Test impact of increased daily rewards on retention"
    },

    # Engagement experiments
    "notification_frequency": {
        "enabled": True,
        "variants": {
            "control": 0.33,     # Standard notifications
            "treatment_a": 0.33,  # More frequent
            "treatment_b": 0.34   # Less frequent
        },
        "description": "Test optimal notification frequency"
    },

    # UI/UX experiments
    "conversation_ui": {
        "enabled": False,
        "variants": {
            "control": 0.5,      # Original UI
            "treatment": 0.5     # New conversation interface
        },
        "description": "Test new conversation UI design"
    },

    # Game mechanics experiments
    "energy_regeneration": {
        "enabled": True,
        "variants": {
            "control": 0.5,      # Standard regeneration rate
            "treatment": 0.5     # Faster regeneration
        },
        "description": "Test faster energy regeneration impact"
    }
}


class ABTestingManager:
    """Manages A/B test variant assignment and tracking."""

    def __init__(self, feature_flags: Optional[Dict] = None):
        """
        Initialize the A/B testing manager.

        Args:
            feature_flags: Optional custom feature flags dictionary
        """
        self.feature_flags = feature_flags or FEATURE_FLAGS
        self._cache = {}  # Cache player assignments

    def get_variant(self, player_id: str, flag_name: str) -> str:
        """
        Get the assigned variant for a player and feature flag.
        Uses consistent hashing to ensure same player always gets same variant.

        Args:
            player_id: Player's unique identifier
            flag_name: Name of the feature flag

        Returns:
            Variant name (e.g., 'control', 'treatment')
        """
        # Check cache first
        cache_key = f"{player_id}:{flag_name}"
        if cache_key in self._cache:
            return self._cache[cache_key]

        # Check if flag exists and is enabled
        if flag_name not in self.feature_flags:
            logger.warning(f"Feature flag '{flag_name}' not found, defaulting to control")
            return "control"

        flag = self.feature_flags[flag_name]
        if not flag.get("enabled", False):
            logger.debug(f"Feature flag '{flag_name}' is disabled, returning control")
            return "control"

        # Use consistent hashing to assign variant
        variant = self._assign_variant(player_id, flag_name, flag["variants"])

        # Cache the assignment
        self._cache[cache_key] = variant

        logger.debug(f"Player {player_id} assigned to '{variant}' for flag '{flag_name}'")
        return variant

    def _assign_variant(self, player_id: str, flag_name: str, variants: Dict[str, float]) -> str:
        """
        Assign a variant using consistent hashing.

        Args:
            player_id: Player's unique identifier
            flag_name: Feature flag name
            variants: Dictionary of variant names to weights

        Returns:
            Assigned variant name
        """
        # Create hash from player_id and flag_name
        hash_input = f"{player_id}:{flag_name}".encode('utf-8')
        hash_value = int(hashlib.md5(hash_input).hexdigest(), 16)

        # Normalize hash to 0-1 range
        normalized = (hash_value % 10000) / 10000.0

        # Assign variant based on cumulative weights
        cumulative = 0.0
        for variant_name, weight in variants.items():
            cumulative += weight
            if normalized <= cumulative:
                return variant_name

        # Fallback to last variant (should not happen if weights sum to 1.0)
        return list(variants.keys())[-1]

    def is_feature_enabled(self, flag_name: str) -> bool:
        """
        Check if a feature flag is enabled.

        Args:
            flag_name: Name of the feature flag

        Returns:
            True if enabled, False otherwise
        """
        if flag_name not in self.feature_flags:
            return False
        return self.feature_flags[flag_name].get("enabled", False)

    def enable_feature(self, flag_name: str):
        """
        Enable a feature flag.

        Args:
            flag_name: Name of the feature flag
        """
        if flag_name in self.feature_flags:
            self.feature_flags[flag_name]["enabled"] = True
            logger.info(f"Enabled feature flag '{flag_name}'")

    def disable_feature(self, flag_name: str):
        """
        Disable a feature flag.

        Args:
            flag_name: Name of the feature flag
        """
        if flag_name in self.feature_flags:
            self.feature_flags[flag_name]["enabled"] = False
            logger.info(f"Disabled feature flag '{flag_name}'")

    def get_all_flags(self) -> Dict[str, Any]:
        """
        Get all feature flags configuration.

        Returns:
            Dictionary of all feature flags
        """
        return self.feature_flags

    def get_player_variants(self, player_id: str) -> Dict[str, str]:
        """
        Get all variant assignments for a player.

        Args:
            player_id: Player's unique identifier

        Returns:
            Dictionary mapping flag names to variant assignments
        """
        assignments = {}
        for flag_name in self.feature_flags:
            assignments[flag_name] = self.get_variant(player_id, flag_name)
        return assignments

    def clear_cache(self, player_id: Optional[str] = None):
        """
        Clear cached variant assignments.

        Args:
            player_id: Optional player ID. If None, clears all cache.
        """
        if player_id:
            # Clear only for specific player
            keys_to_remove = [k for k in self._cache if k.startswith(f"{player_id}:")]
            for key in keys_to_remove:
                del self._cache[key]
        else:
            # Clear all cache
            self._cache.clear()


# Global A/B testing manager
_ab_manager = None


def get_ab_manager() -> ABTestingManager:
    """Get or create the global A/B testing manager instance."""
    global _ab_manager
    if _ab_manager is None:
        _ab_manager = ABTestingManager()
    return _ab_manager


def get_variant(player_id: str, flag_name: str) -> str:
    """
    Get variant assignment for a player (convenience function).

    Args:
        player_id: Player's unique identifier
        flag_name: Feature flag name

    Returns:
        Variant name
    """
    return get_ab_manager().get_variant(player_id, flag_name)


def is_in_treatment(player_id: str, flag_name: str) -> bool:
    """
    Check if a player is in treatment variant.

    Args:
        player_id: Player's unique identifier
        flag_name: Feature flag name

    Returns:
        True if in treatment, False otherwise
    """
    variant = get_variant(player_id, flag_name)
    return variant.startswith('treatment')


# Example usage in game code:
"""
from experiments.ab_testing import get_variant, is_in_treatment

# Get specific variant
variant = get_variant(player_id, 'tutorial_flow')
if variant == 'treatment':
    # Show new tutorial
    show_new_tutorial()
else:
    # Show original tutorial
    show_original_tutorial()

# Simple boolean check
if is_in_treatment(player_id, 'energy_regeneration'):
    energy_regen_rate = 2.0  # Faster
else:
    energy_regen_rate = 1.0  # Standard
"""