import asyncio
import time
import os
from typing import Optional, List, Dict, Any, Callable
from autobahn.asyncio.wamp import ApplicationSession
from autobahn.asyncio.component import Component
from autobahn.wamp import auth
from enum import Enum

class Stage(Enum):
    """Stage enumeration for different deployment environments"""
    DEVELOPMENT = "DEV"
    PRODUCTION = "PROD"

def getSerialNumber(serial_number: str = None) -> str:
    """Get device serial number from parameter or environment variable"""
    if serial_number is None:
        s_num = os.environ.get("DEVICE_SERIAL_NUMBER")
        if s_num is None:
            raise Exception("ENV Variable 'DEVICE_SERIAL_NUMBER' is not set!")
    else:
        s_num = serial_number
    return s_num

DATAPODS_WS_URI = "wss://cbw.datapods.io/ws-ua-usr"
STUDIO_WS_URI_OLD = "wss://cbw.record-evolution.com/ws-ua-usr"
STUDIO_WS_URI = "wss://cbw.ironflock.com/ws-ua-usr"
LOCALHOST_WS_URI = "ws://localhost:8080/ws-ua-usr"

socketURIMap = {
    "https://studio.datapods.io": DATAPODS_WS_URI,
    "https://studio.record-evolution.com": STUDIO_WS_URI_OLD,
    "https://studio.ironflock.com": STUDIO_WS_URI,
    "http://localhost:8086": LOCALHOST_WS_URI,
    "http://host.docker.internal:8086": LOCALHOST_WS_URI
}

class CrossbarConnection:
    """
    Python version of the TypeScript CrossbarConnection class.
    Manages WAMP connections with automatic reconnection, subscription management, and session handling.
    """
    
    def __init__(self):
        self.session: Optional[ApplicationSession] = None
        self.component: Optional[Component] = None
        self.connection_options: Optional[Dict[str, Any]] = None
        self.subscriptions: List[Any] = []  # List of autobahn subscriptions
        self.registrations: List[Any] = []  # List of autobahn registrations
        self.connection_drop_wait_time: int = 6000  # milliseconds
        self.realm: Optional[str] = None
        self.serial_number: Optional[str] = None
        self._first_connection_future: Optional[asyncio.Future] = None
        self._is_connected = False
        self._session_wait_timeout = 6.0  # seconds
        
    @staticmethod
    def getWebSocketURI():
        reswarm_url = os.environ.get("RESWARM_URL")
        if not reswarm_url:
            return STUDIO_WS_URI
        return socketURIMap.get(reswarm_url)
        
    async def configure(
        self, 
        swarm_key: int, 
        app_key: int, 
        stage: Stage,
        cburl: Optional[str] = None,
        serial_number: Optional[str] = None
    ):
        """
        Configure the crossbar connection with realm and WAMPCRA authentication details.
        
        Args:
            swarm_key: The swarm key identifier
            app_key: The application key identifier  
            stage: The deployment stage (development, production)
            cburl: Optional crossbar URL (defaults to environment variable RESWARM_URL mapping)
            serial_number: Optional device serial number (defaults to DEVICE_SERIAL_NUMBER env var)
        """
        self.realm = f"realm-{swarm_key}-{app_key}-{stage.value.lower()}"
        self.serial_number = getSerialNumber(serial_number)
        
        # Get URL from environment if not provided
        if cburl is None:
            cburl = self.getWebSocketURI()
            
        self.connection_options = {
            'url': cburl,
            'realm': self.realm,
            'authmethods': ['wampcra'],
            'authid': self.serial_number,
            'max_retries': -1,  # infinite retries
            'max_retry_delay': 2,
            'initial_retry_delay': 1,
            'serializers': ['msgpack']
        }
        
        # Create the component with WAMPCRA authentication
        transports = [{
            'type': 'websocket',
            'url': cburl,
            'serializers': ['msgpack']
        }]
        
        # Custom session class for WAMPCRA authentication
        class AppSession(ApplicationSession):
            def onConnect(session_self):
                print('onConnect called')
                session_self.join(self.realm, ['wampcra'], self.serial_number)
                
            def onChallenge(session_self, challenge):
                print(f'challenge requested for {challenge.method}')
                if challenge.method == "wampcra":
                    signature = auth.compute_wcs(
                        self.serial_number, challenge.extra["challenge"]
                    )
                    return signature
                raise Exception(f"Invalid authmethod {challenge.method}")
        
        self.component = Component(
            transports=transports,
            realm=self.realm,
            session_factory=AppSession
        )
        
        # Set up event handlers
        self.component.on('join', self._on_open)
        self.component.on('leave', self._on_close)
        self.component.on('disconnect', self._on_disconnect)
        
        # Create future for first connection
        self._first_connection_future = asyncio.Future()
        
    async def _on_open(self, session: ApplicationSession, details):
        """Handle session open event"""
        self.session = session
        self._is_connected = True
        await self._resubscribe_all()
        print(f"Connection to IronFlock app realm '{session._realm}' established")
        
        if self._first_connection_future and not self._first_connection_future.done():
            self._first_connection_future.set_result(None)
            
    async def _on_close(self, session: ApplicationSession, details):
        """Handle session close event"""
        self.session = None
        self._is_connected = False
        print(f"Connection to IronFlock app realm {self.realm} closed: {details.reason}")
        
        if self._first_connection_future and not self._first_connection_future.done():
            self._first_connection_future.set_exception(
                Exception(f"Connection closed: {details.reason}")
            )
            
    async def _on_disconnect(self, session: ApplicationSession, was_clean: bool):
        """Handle disconnect event"""
        self.session = None
        self._is_connected = False
        print(f"Disconnected from IronFlock app realm {self.realm}, clean: {was_clean}")

    async def _session_wait(self) -> None:
        """Wait for session to be available"""
        start_time = time.time()
        while not self.session or not self._is_connected:
            if time.time() - start_time > self._session_wait_timeout:
                raise TimeoutError('Timeout waiting for session')
            await asyncio.sleep(0.2)
            
    @property
    def is_open(self) -> bool:
        """Check if the connection is open"""
        return self._is_connected and self.session is not None
        
    async def start(self) -> None:
        """Start the connection"""
        if not self.component:
            raise ValueError("Must call configure() before start()")

        print(f'Starting connection for IronFlock app realm {self.realm}')

        # Start the component (non-blocking in autobahn asyncio)
        self.component.start()
        
        # Wait for first connection to be established
        if self._first_connection_future:
            await self._first_connection_future
            
    async def stop(self) -> None:
        """Stop the connection"""
        if self.component:
            await self.component.stop()
        self._is_connected = False
        self.session = None
        
    async def call(
        self, 
        topic: str, 
        args: Optional[List[Any]] = None, 
        kwargs: Optional[Dict[str, Any]] = None,
        options: Optional[Dict[str, Any]] = None
    ) -> Any:
        """Call a remote procedure"""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        args = args or []
        kwargs = kwargs or {}
        
        result = await self.session.call(topic, *args, **kwargs, options=options)
        return result
        
    async def subscribe(
        self, 
        topic: str, 
        handler: Callable,
        options: Optional[Dict[str, Any]] = None
    ) -> Optional[Any]:
        """Subscribe to a topic"""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        subscription = await self.session.subscribe(handler, topic, options=options)
        if subscription:
            self.subscriptions.append(subscription)
        return subscription
        
    async def publish(
        self, 
        topic: str, 
        args: Optional[List[Any]] = None,
        kwargs: Optional[Dict[str, Any]] = None,
        options: Optional[Dict[str, Any]] = None
    ) -> Optional[Any]:
        """Publish to a topic"""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        args = args or []
        kwargs = kwargs or {}
        
        result = await self.session.publish(topic, *args, options=options, **kwargs)
        return result
        
    async def unsubscribe(self, subscription: Any) -> None:
        """Unsubscribe from a subscription"""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        if subscription in self.subscriptions:
            self.subscriptions.remove(subscription)
            await self.session.unsubscribe(subscription)
            
    async def unsubscribe_func(
        self, 
        topic: str, 
        handler: Optional[Callable] = None
    ) -> None:
        """Unsubscribe a specific function from a topic"""
        await self._session_wait()
        
        if handler is None:
            return await self.unsubscribe_topic(topic)
            
        # Find and remove subscriptions matching topic and handler
        to_remove = []
        for subscription in self.subscriptions:
            if hasattr(subscription, 'topic') and subscription.topic == topic:
                if hasattr(subscription, 'handler') and subscription.handler == handler:
                    to_remove.append(subscription)
                    
        for subscription in to_remove:
            try:
                await self.unsubscribe(subscription)
            except Exception as e:
                print(f"Failed to unsubscribe from {topic}: {e}")
                
    async def unsubscribe_topic(self, topic: str) -> bool:
        """Unsubscribe all handlers from a topic"""
        matching_subs = [
            sub for sub in self.subscriptions 
            if hasattr(sub, 'topic') and sub.topic == topic
        ]
        
        if not matching_subs:
            return True
            
        # Unsubscribe all matching subscriptions
        tasks = [self.unsubscribe(sub) for sub in matching_subs]
        await asyncio.gather(*tasks, return_exceptions=True)
        return True
        
    async def register(
        self, 
        topic: str, 
        endpoint: Callable,
        options: Optional[Dict[str, Any]] = None
    ) -> Optional[Any]:
        """Register a procedure"""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        registration = await self.session.register(endpoint, topic, options=options)
        if registration:
            self.registrations.append(registration)
        return registration
        
    async def unregister(self, registration: Any) -> None:
        """Unregister a procedure"""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        if registration in self.registrations:
            self.registrations.remove(registration)
            await self.session.unregister(registration)
            
    async def _resubscribe_all(self) -> None:
        """Resubscribe to all topics after reconnection"""
        if not self.subscriptions:
            return
            
        print(f'Resubscribing to {len(self.subscriptions)} subscriptions')
        
        # Store current subscriptions and clear the list
        old_subscriptions = self.subscriptions.copy()
        self.subscriptions.clear()
        
        # Resubscribe to all topics
        tasks = []
        for old_sub in old_subscriptions:
            if hasattr(old_sub, 'topic') and hasattr(old_sub, 'handler'):
                task = self.subscribe(old_sub.topic, old_sub.handler)
                tasks.append(task)
                
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)