#!/usr/bin/env python
import json
from dataclasses import dataclass, field
import random
import logging
from typing import Any, Dict, List, Optional, Union, Callable, Awaitable
from typing import Set
import asyncio
import threading
import traceback
from datetime import datetime
from growthbook import FeatureRepository, feature_repo
from contextlib import asynccontextmanager

from .core import eval_feature as core_eval_feature, run_experiment
from .common_types import (
    Feature,
    GlobalContext,
    Options,
    Result,
    UserContext,
    EvaluationContext,
    StackContext,
    FeatureResult,
    FeatureRefreshStrategy,
    Experiment
)

logger = logging.getLogger("growthbook.growthbook_client")

class SingletonMeta(type):
    """Thread-safe implementation of Singleton pattern"""
    _instances: Dict[type, Any] = {}
    _lock = threading.Lock()

    def __call__(cls, *args, **kwargs):
        with cls._lock:
            if cls not in cls._instances:
                instance = super().__call__(*args, **kwargs)
                cls._instances[cls] = instance
        return cls._instances[cls]

class BackoffStrategy:
    """Exponential backoff with jitter for failed requests"""
    def __init__(
        self, 
        initial_delay: float = 1.0, 
        max_delay: float = 60.0, 
        multiplier: float = 2.0,
        jitter: float = 0.1
    ):
        self.initial_delay = initial_delay
        self.max_delay = max_delay
        self.multiplier = multiplier
        self.jitter = jitter
        self.current_delay = initial_delay
        self.attempt = 0

    def next_delay(self) -> float:
        """Calculate next delay with jitter"""
        delay = min(
            self.current_delay * (self.multiplier ** self.attempt), 
            self.max_delay
        )
        # Add random jitter
        jitter_amount = delay * self.jitter
        delay = delay + (random.random() * 2 - 1) * jitter_amount
        self.attempt += 1
        return max(delay, self.initial_delay)

    def reset(self) -> None:
        """Reset backoff state"""
        self.current_delay = self.initial_delay
        self.attempt = 0

class WeakRefWrapper:
    """A wrapper class to allow weak references for otherwise non-weak-referenceable objects."""
    def __init__(self, obj):
        self.obj = obj

class FeatureCache:
    """Thread-safe feature cache"""
    def __init__(self):
        self._cache: Dict[str, Dict[str, Any]] = {
            'features': {},
            'savedGroups': {}
        }
        self._lock = threading.Lock()

    def update(self, features: Dict[str, Any], saved_groups: Dict[str, Any]) -> None:
        """Simple thread-safe update of cache with new API data"""
        with self._lock:
            self._cache['features'] = dict(features)
            self._cache['savedGroups'] = dict(saved_groups)

    def get_current_state(self) -> Dict[str, Any]:
        """Get current cache state"""
        with self._lock:
            return {
                "features": dict(self._cache['features']),
                "savedGroups": self._cache['savedGroups']
            }

class EnhancedFeatureRepository(FeatureRepository, metaclass=SingletonMeta):
    def __init__(self,
                 api_host: str,
                 client_key: str,
                 decryption_key: str = "",
                 cache_ttl: int = 60,
                 http_connect_timeout: Optional[int] = None,
                 http_read_timeout: Optional[int] = None):
        FeatureRepository.__init__(self)
        self._api_host = api_host
        self._client_key = client_key
        self._decryption_key = decryption_key
        self._cache_ttl = cache_ttl
        self._refresh_lock = threading.Lock()
        self._refresh_task: Optional[asyncio.Task] = None
        self._stop_event = asyncio.Event()
        self._backoff = BackoffStrategy()
        self._feature_cache = FeatureCache()
        self._callbacks: List[Callable[[Dict[str, Any]], Awaitable[None]]] = []
        self._last_successful_refresh: Optional[datetime] = None
        self._refresh_in_progress = asyncio.Lock()
        self.http_connect_timeout = http_connect_timeout
        self.http_read_timeout = http_read_timeout

    @asynccontextmanager
    async def refresh_operation(self):
        """Context manager for feature refresh with proper cleanup"""
        if self._refresh_in_progress.locked():
            yield False
            return

        # async with self._refresh_in_progress:
        try:
            await self._refresh_in_progress.acquire()
            yield True
            self._backoff.reset()
            self._last_successful_refresh = datetime.now()
        except Exception as e:
            delay = self._backoff.next_delay()
            logger.error(f"Refresh failed, next attempt in {delay:.2f}s: {str(e)}")
            traceback.print_exc()
            raise
        finally:
            if self._refresh_in_progress.locked():
                self._refresh_in_progress.release()

    async def _handle_feature_update(self, data: Dict[str, Any]) -> None:
        """Update features with memory optimization"""
        # Directly update with new features
        self._feature_cache.update(
            data.get("features", {}),
            data.get("savedGroups", {})
        )

        # Create a copy of callbacks to avoid modification during iteration
        with self._refresh_lock:
            callbacks = self._callbacks.copy()

        for callback in callbacks:
            try:
                await callback(dict(self._feature_cache.get_current_state()))
            except Exception:
                traceback.print_exc()

    def add_callback(self, callback: Callable[[Dict[str, Any]], Awaitable[None]]) -> None:
        """Add callback to the list"""
        with self._refresh_lock:
            if callback not in self._callbacks:
                self._callbacks.append(callback)

    def remove_callback(self, callback: Callable[[Dict[str, Any]], Awaitable[None]]) -> None:
        """Remove callback from the list"""
        with self._refresh_lock:
            if callback in self._callbacks:
                self._callbacks.remove(callback)

    """
    _start_sse_refresh flow mimics a bridge pattern to connect a blocking, synchronous background thread 
    (the SSEClient) with your non-blocking, async main loop.

    Bridge - _maintain_sse_connection - runs on the main async loop, calls `startAutoRefresh` (which in turn spawns a thread)
    and waits indefinitely. (Awaiting a Future suspends the coroutine, costing zero CPU)

    The SSEClient runs in a separate thread, makes a blocking HTTP request, and invokes `on_event` synchronously.

    The Hand off - when the event arrives (we're still on the background thread), sse_handler uses `asyncio.run_coroutine_threadsafe` 
    to schedule the async processing `_handle_sse_event` onto the main event loop.
    """

    async def _handle_sse_event(self, event_data: Dict[str, Any]) -> None:
        """Process an event received from the SSE connection"""
        try:
            event_type = event_data.get("type")
            if event_type == "features-updated":
                response = await self.load_features_async(
                    self._api_host, self._client_key, self._decryption_key, self._cache_ttl
                )
                if response is not None:
                    await self._handle_feature_update(response)
            elif event_type == "features":
                data = event_data.get("data", "{}")
                if isinstance(data, str):
                    data = json.loads(data)
                await self._handle_feature_update(data)
        except Exception:
            logger.exception("Error handling SSE event")

    async def _start_sse_refresh(self) -> None:
        """Start SSE-based feature refresh"""
        with self._refresh_lock:
            if self._refresh_task is not None:  # Already running
                return

            # SSEClient invokes `on_event` synchronously from a background thread.
            main_loop = asyncio.get_running_loop()

            # We must not pass an `async def` callback here (it would never be awaited).
            def sse_handler(event_data: Dict[str, Any]) -> None:
                # Schedule async processing onto the main event loop.
                try:
                    asyncio.run_coroutine_threadsafe(self._handle_sse_event(event_data), main_loop)
                except Exception:
                    logger.exception("Failed to schedule SSE event handler")

            async def _maintain_sse_connection() -> None:
                """
                Start SSE streaming and keep the task alive until cancelled.
                """
                try:
                    # NOTE: `startAutoRefresh` is synchronous and starts a background thread.
                    self.startAutoRefresh(self._api_host, self._client_key, sse_handler)
                    
                    # Wait indefinitely until the task is cancelled - basically saying "Keep this service 'active' until someone cancels me."
                    # reconnection logic is handled inside SSEClient's thread
                    await asyncio.Future()
                except asyncio.CancelledError:
                    # Normal shutdown flow
                    raise
                except Exception:
                    logger.exception("Unexpected error in SSE lifecycle task")
                finally:
                    try:
                        # stopAutoRefresh blocks joining a thread, so it needs to be run in executor
                        # to avoid blocking the async event loop
                        await main_loop.run_in_executor(
                            None, 
                            lambda: self.stopAutoRefresh(timeout=10)
                        )
                    except Exception:
                        logger.exception("Failed to stop SSE auto-refresh")

            # Start a task that owns the SSE lifecycle and cleanup.
            self._refresh_task = asyncio.create_task(_maintain_sse_connection())

    async def _start_http_refresh(self, interval: int = 60) -> None:
        """Enhanced HTTP polling with backoff"""
        if self._refresh_task:
            return

        async def refresh_loop() -> None:
            try:
                while not self._stop_event.is_set():
                    async with self.refresh_operation() as should_refresh:
                        if should_refresh:
                            try:
                                response = await self.load_features_async(
                                    api_host=self._api_host,
                                    client_key=self._client_key,
                                    decryption_key=self._decryption_key,
                                    ttl=self._cache_ttl
                                )
                                if response is not None:
                                    await self._handle_feature_update(response)
                                # On success, reset backoff and use normal interval
                                self._backoff.reset()
                                try:
                                    await asyncio.sleep(interval)
                                except asyncio.CancelledError:
                                    # Allow cancellation during sleep
                                    raise
                            except Exception as e:
                                # On failure, use backoff delay
                                delay = self._backoff.next_delay()
                                logger.error(f"Refresh failed, next attempt in {delay:.2f}s: {str(e)}")
                                traceback.print_exc()
                                try:
                                    await asyncio.sleep(delay)
                                except asyncio.CancelledError:
                                    # Allow cancellation during sleep
                                    raise
            except asyncio.CancelledError:
                # Clean exit on cancellation
                raise
            finally:
                # Ensure we're marked as stopped
                self._stop_event.set()

        self._refresh_task = asyncio.create_task(refresh_loop())

    async def start_feature_refresh(self, strategy: FeatureRefreshStrategy, callback=None):
        """Initialize feature refresh based on strategy"""
        self._refresh_callback = callback
        
        if strategy == FeatureRefreshStrategy.SERVER_SENT_EVENTS:
            await self._start_sse_refresh()
        else:
            await self._start_http_refresh()

    async def stop_refresh(self) -> None:
        """Clean shutdown of refresh tasks"""
        self._stop_event.set()
        # Ensure any SSE background thread is stopped as well.
        try:
            self.stopAutoRefresh(timeout=10)
        except Exception:
            # Best-effort cleanup; task cancellation below will proceed.
            logger.exception("Error stopping SSE auto-refresh")
        if self._refresh_task:
            # Cancel the task
            self._refresh_task.cancel()
            try:
                # Wait for it to actually finish
                await self._refresh_task
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error(f"Error during refresh task cleanup: {e}")
            finally:
                self._refresh_task = None
                self._backoff.reset()
        self._stop_event.clear()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.stop_refresh()
    
    async def load_features_async(
        self, api_host: str, client_key: str, decryption_key: str = "", ttl: int = 60
    ) -> Optional[Dict]:
        # Use stored values when called internally
        if api_host == self._api_host and client_key == self._client_key:
            decryption_key = self._decryption_key
            ttl = self._cache_ttl
        return await super().load_features_async(api_host, client_key, decryption_key, ttl)

class GrowthBookClient:
    def __init__(
        self,
        options: Optional[Union[Dict[str, Any], Options]] = None
    ):  
        self.options = (
            options if isinstance(options, Options)
            else Options(**options) if options
            else Options()
        )
        
        # Thread-safe tracking state
        self._tracked: Dict[str, bool] = {}  # Access only within async context
        self._tracked_lock = threading.Lock()
        
        # Thread-safe subscription management
        self._subscriptions: Set[Callable[[Experiment, Result], None]] = set()
        self._subscriptions_lock = threading.Lock()

        # Add sticky bucket cache
        self._sticky_bucket_cache: Dict[str, Dict[str, Any]] = {
            'attributes': {},
            'assignments': {}
        }
        self._sticky_bucket_cache_lock = False
        
        # Plugin support
        self._tracking_plugins: List[Any] = self.options.tracking_plugins or []
        self._initialized_plugins: List[Any] = []

        self._features_repository = (
            EnhancedFeatureRepository(
                self.options.api_host or "https://cdn.growthbook.io",
                self.options.client_key or "",
                self.options.decryption_key or "",
                self.options.cache_ttl,
                self.options.http_connect_timeout,
                self.options.http_read_timeout
            )
            if self.options.client_key
            else None
        )
        
        self._global_context: Optional[GlobalContext] = None
        self._context_lock = asyncio.Lock()
        
        # Initialize plugins
        self._initialize_plugins()


    def _track(self, experiment: Experiment, result: Result, user_context: UserContext) -> None:
        """Thread-safe tracking implementation"""
        if not self.options.on_experiment_viewed:
            return

        # Create unique key for this tracking event
        key = (
            result.hashAttribute
            + str(result.hashValue)
            + experiment.key
            + str(result.variationId)
        )

        with self._tracked_lock:
            if not self._tracked.get(key):
                try:
                    self.options.on_experiment_viewed(experiment, result, user_context)
                    self._tracked[key] = True
                except Exception:
                    logger.exception("Error in tracking callback")

    def subscribe(self, callback: Callable[[Experiment, Result], None]) -> Callable[[], None]:
        """Thread-safe subscription management"""
        with self._subscriptions_lock:
            self._subscriptions.add(callback)
            def unsubscribe():
                with self._subscriptions_lock:
                    self._subscriptions.discard(callback)
            return unsubscribe

    def _fire_subscriptions(self, experiment: Experiment, result: Result) -> None:
        """Thread-safe subscription notifications"""
        with self._subscriptions_lock:
            subscriptions = self._subscriptions.copy()

        for callback in subscriptions:
            try:
                callback(experiment, result)
            except Exception:
                logger.exception("Error in subscription callback")


    async def set_features(self, features: dict) -> None:
        await self._feature_update_callback({"features": features})
        
    
    async def _refresh_sticky_buckets(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
        """Refresh sticky bucket assignments only if attributes have changed"""
        if not self.options.sticky_bucket_service:
            return {}

        # Use compare-and-swap pattern
        while not self._sticky_bucket_cache_lock:
            if attributes == self._sticky_bucket_cache['attributes']:
                return self._sticky_bucket_cache['assignments']
            
            self._sticky_bucket_cache_lock = True
            try:
                assignments = self.options.sticky_bucket_service.get_all_assignments(attributes)
                self._sticky_bucket_cache['attributes'] = attributes.copy()
                self._sticky_bucket_cache['assignments'] = assignments
                return assignments
            finally:
                self._sticky_bucket_cache_lock = False
        
        # Fallback return for edge case where loop condition is never satisfied
        return {}

    async def initialize(self) -> bool:
        """Initialize client with features and start refresh"""
        if not self._features_repository:
            logger.error("No features repository available")
            return False

        try:
            # Initial feature load
            initial_features = await self._features_repository.load_features_async(
                self.options.api_host or "https://cdn.growthbook.io", 
                self.options.client_key or "", 
                self.options.decryption_key or "", 
                self.options.cache_ttl
            )
            if not initial_features:
                logger.error("Failed to load initial features")
                return False

            # Create global context with initial features
            await self._feature_update_callback(initial_features)
            
            # Set up callback for future updates
            self._features_repository.add_callback(self._feature_update_callback)
            
            # Start feature refresh
            refresh_strategy = self.options.refresh_strategy or FeatureRefreshStrategy.STALE_WHILE_REVALIDATE
            await self._features_repository.start_feature_refresh(refresh_strategy)
            return True
            
        except Exception as e:
            logger.error(f"Initialization failed: {str(e)}", exc_info=True)
            traceback.print_exc()
            return False

    async def _feature_update_callback(self, features_data: Dict[str, Any]) -> None:
        """Handle feature updates and manage global context"""
        if not features_data:
            logger.warning("Warning: Received empty features data")
            return

        async with self._context_lock:
            features = {}

            for key, feature in features_data.get("features", {}).items():
                if isinstance(feature, Feature):
                    features[key] = feature
                else:
                    features[key] = Feature(
                        rules=feature.get("rules", []),
                        defaultValue=feature.get("defaultValue", None),
                    )

            if self._global_context is None:
                # Initial creation of global context
                self._global_context = GlobalContext(
                        options=self.options,
                        features=features,
                        saved_groups=features_data.get("savedGroups", {})
                )
            else:
                # Update existing global context
                self._global_context.features = features
                self._global_context.saved_groups = features_data.get("savedGroups", {})

    async def __aenter__(self):
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def create_evaluation_context(self, user_context: UserContext) -> EvaluationContext:
        """Create evaluation context for feature evaluation"""
        if self._global_context is None:
            raise RuntimeError("GrowthBook client not properly initialized")
            
        # Get sticky bucket assignments if needed
        sticky_assignments = await self._refresh_sticky_buckets(user_context.attributes)
        
        # update user context with sticky bucket assignments
        user_context.sticky_bucket_assignment_docs = sticky_assignments

        return EvaluationContext(
            user=user_context,
            global_ctx=self._global_context,
            stack=StackContext(evaluated_features=set())
        )

    async def eval_feature(self, key: str, user_context: UserContext) -> FeatureResult:
        """Evaluate a feature with proper async context management"""
        async with self._context_lock:
            context = await self.create_evaluation_context(user_context)
            result = core_eval_feature(key=key, evalContext=context, tracking_cb=self._track)
            # Call feature usage callback if provided
            if self.options.on_feature_usage:
                try:
                    self.options.on_feature_usage(key, result, user_context)
                except Exception:
                    logger.exception("Error in feature usage callback")
            return result

    async def is_on(self, key: str, user_context: UserContext) -> bool:
        """Check if a feature is enabled with proper async context management"""
        async with self._context_lock:
            context = await self.create_evaluation_context(user_context)
            result = core_eval_feature(key=key, evalContext=context, tracking_cb=self._track)
            # Call feature usage callback if provided
            if self.options.on_feature_usage:
                try:
                    self.options.on_feature_usage(key, result, user_context)
                except Exception:
                    logger.exception("Error in feature usage callback")
            return result.on
    
    async def is_off(self, key: str, user_context: UserContext) -> bool:
        """Check if a feature is set to off with proper async context management"""
        async with self._context_lock:
            context = await self.create_evaluation_context(user_context)
            result = core_eval_feature(key=key, evalContext=context, tracking_cb=self._track)
            # Call feature usage callback if provided
            if self.options.on_feature_usage:
                try:
                    self.options.on_feature_usage(key, result, user_context)
                except Exception:
                    logger.exception("Error in feature usage callback")
            return result.off
    
    async def get_feature_value(self, key: str, fallback: Any, user_context: UserContext) -> Any:
        async with self._context_lock:
            context = await self.create_evaluation_context(user_context)
            result = core_eval_feature(key=key, evalContext=context, tracking_cb=self._track)
            # Call feature usage callback if provided
            if self.options.on_feature_usage:
                try:
                    self.options.on_feature_usage(key, result, user_context)
                except Exception:
                    logger.exception("Error in feature usage callback")
            return result.value if result.value is not None else fallback

    async def run(self, experiment: Experiment, user_context: UserContext) -> Result:
        """Run experiment with tracking"""
        async with self._context_lock:
            context = await self.create_evaluation_context(user_context)
            result = run_experiment(
                experiment=experiment, 
                evalContext=context,
                tracking_cb=self._track
            )
            # Fire subscriptions synchronously
            self._fire_subscriptions(experiment, result)
            return result
        
    async def close(self) -> None:
        """Clean shutdown with proper cleanup"""
        if self._features_repository:
            await self._features_repository.stop_refresh()

        # Clear tracking and subscription state
        with self._tracked_lock:
            self._tracked.clear()
        with self._subscriptions_lock:
            self._subscriptions.clear()
        # Clear context
        async with self._context_lock:
            self._global_context = None
            
        # Cleanup plugins
        self._cleanup_plugins()

    @property
    def user_agent_suffix(self) -> Optional[str]:
        """Get the suffix appended to the User-Agent header"""
        return feature_repo.user_agent_suffix
        
    @user_agent_suffix.setter
    def user_agent_suffix(self, value: Optional[str]) -> None:
        """Set a suffix to be appended to the User-Agent header"""
        feature_repo.user_agent_suffix = value

    def _initialize_plugins(self) -> None:
        """Initialize all tracking plugins with this GrowthBookClient instance."""
        for plugin in self._tracking_plugins:
            try:
                if hasattr(plugin, 'initialize'):
                    # Plugin is a class instance with initialize method
                    plugin.initialize(self)
                    self._initialized_plugins.append(plugin)
                    logger.debug(f"Initialized plugin: {plugin.__class__.__name__}")
                elif callable(plugin):
                    # Plugin is a callable function
                    plugin(self)
                    self._initialized_plugins.append(plugin)
                    logger.debug(f"Initialized callable plugin: {plugin.__name__}")
                else:
                    logger.warning(f"Plugin {plugin} is neither callable nor has initialize method")
            except Exception as e:
                logger.error(f"Failed to initialize plugin {plugin}: {e}")

    def _cleanup_plugins(self) -> None:
        """Cleanup all initialized plugins."""
        for plugin in self._initialized_plugins:
            try:
                if hasattr(plugin, 'cleanup'):
                    plugin.cleanup()
                    logger.debug(f"Cleaned up plugin: {plugin.__class__.__name__}")
            except Exception as e:
                logger.error(f"Error cleaning up plugin {plugin}: {e}")
        self._initialized_plugins.clear()