import asyncio
import os
from dataclasses import dataclass
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 autobahn.wamp.types import CallOptions, PublishOptions, SubscribeOptions, RegisterOptions
from autobahn.wamp.request import Subscription, Registration
from enum import Enum

# Import Pydantic types for validation
try:
    from .types import CrossbarCallParams, SubscriptionParams
except ImportError:
    from ironflock.types import CrossbarCallParams, SubscriptionParams


@dataclass
class SubscriptionEntry:
    """Tracks a subscription with the metadata needed for resubscription on reconnect."""
    topic: str
    handler: Callable
    options: Optional[SubscribeOptions] = None
    # None when resubscription failed and entry is preserved for retry
    sub: Optional[Subscription] = None


@dataclass
class RegistrationEntry:
    """Tracks a registration with the metadata needed for re-registration on reconnect."""
    topic: str
    endpoint: Callable
    options: Optional[RegisterOptions] = None
    # None when re-registration failed and entry is preserved for retry
    reg: Optional[Registration] = None

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

def getSerialNumber(serial_number: Optional[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

STUDIO_DEV_WS_URI = "wss://cbw.ironflock.dev/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.ironflock.dev": STUDIO_DEV_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[SubscriptionEntry] = []
        self.registrations: List[RegistrationEntry] = []
        self.connection_drop_wait_time: int = 6000  # milliseconds
        self.realm: Optional[str] = None
        self.serial_number: Optional[str] = None
        self._connection_event: Optional[asyncio.Event] = None
        self._is_connected = False
        self._session_wait_timeout = 10.0  # seconds - long enough for reconnection
        self._first_connection_timeout = 10.0  # seconds - timeout for initial connection
        self._max_connection_retries = 5  # maximum retries for first connection
        self._max_reconnect_delay = 10  # seconds - max backoff for reconnection
        self._connection_error: Optional[Exception] = None
        self._component_task: Optional[asyncio.Task] = None
        self._should_reconnect = False
        self._reconnect_task: Optional[asyncio.Task] = None
        
    @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
        _outer = self
        
        class AppSession(ApplicationSession):
            def onConnect(self):
                print('onConnect called')
                assert _outer.realm is not None
                self.join(_outer.realm, ['wampcra'], _outer.serial_number)
                
            def onChallenge(self, challenge) -> str:
                print(f'challenge requested for {challenge.method}')
                if challenge.method == "wampcra":
                    signature = auth.compute_wcs(
                        _outer.serial_number, challenge.extra["challenge"]
                    )
                    return signature.decode('ascii') if isinstance(signature, bytes) else str(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 event for connection signaling (can be reset unlike Future)
        self._connection_event = asyncio.Event()
        self._connection_error = None
        
    async def _on_open(self, session: ApplicationSession, details):
        """Handle session open event"""
        self.session = session
        self._is_connected = True
        self._connection_error = None
        await self._resubscribe_all()
        await self._reregister_all()
        print(f"Connection to IronFlock app realm '{session._realm}' established")
        
        # Signal that connection is established
        if self._connection_event:
            self._connection_event.set()
            
    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}")
        
        # Store error for connection waiters
        self._connection_error = Exception(f"Connection closed: {details.reason}")
        
        # Reset event so next connection attempt can wait
        if self._connection_event:
            self._connection_event.clear()
            
    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}")
        
        # Reset event so next connection attempt can wait
        if self._connection_event:
            self._connection_event.clear()
        
        # Trigger automatic reconnection if enabled
        if self._should_reconnect:
            if self._reconnect_task is None or self._reconnect_task.done():
                self._reconnect_task = asyncio.create_task(self._reconnect_loop())

    async def _session_wait(self) -> None:
        """Wait for session to be available"""
        if self.session and self._is_connected:
            return
            
        # If we have a connection event, wait on it with timeout
        if self._connection_event:
            try:
                await asyncio.wait_for(
                    self._connection_event.wait(),
                    timeout=self._session_wait_timeout
                )
            except asyncio.TimeoutError:
                raise TimeoutError('Timeout waiting for session')
        else:
            # No connection event means configure() wasn't called or connection not started
            # Wait with timeout using sleep-based approach
            await asyncio.sleep(self._session_wait_timeout)
            raise TimeoutError('Timeout waiting for session')
        
        # Double-check we have a valid session
        if not self.session or not self._is_connected:
            if self._connection_error:
                raise self._connection_error
            raise RuntimeError('Session not available')
            
    @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 with timeout and retry logic"""
        if not self.component:
            raise ValueError("Must call configure() before start()")

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

        last_error: Optional[Exception] = None
        
        for attempt in range(1, self._max_connection_retries + 1):
            try:
                # Reset connection state for this attempt
                self._connection_error = None
                if self._connection_event:
                    self._connection_event.clear()
                
                # Stop any existing component task
                if self._component_task and not self._component_task.done():
                    self._component_task.cancel()
                    try:
                        await self._component_task
                    except asyncio.CancelledError:
                        pass
                
                # Start the component in a background task
                self._component_task = asyncio.create_task(self._run_component())
                
                # Wait for connection with timeout
                if self._connection_event:
                    try:
                        await asyncio.wait_for(
                            self._connection_event.wait(),
                            timeout=self._first_connection_timeout
                        )
                        
                        # Check if we actually connected or got an error
                        if self._is_connected and self.session:
                            print(f"Successfully connected on attempt {attempt}")
                            # Enable automatic reconnection after initial connection
                            self._should_reconnect = True
                            return
                        elif self._connection_error:
                            raise self._connection_error
                            
                    except asyncio.TimeoutError:
                        print(f"Connection attempt {attempt}/{self._max_connection_retries} timed out after {self._first_connection_timeout}s")
                        last_error = TimeoutError(f"Connection timed out after {self._first_connection_timeout}s")
                        
                        # Stop the component before retrying
                        await self._stop_component()
                        
                        if attempt < self._max_connection_retries:
                            retry_delay = min(2 ** attempt, self._max_reconnect_delay)
                            print(f"Retrying in {retry_delay}s...")
                            await asyncio.sleep(retry_delay)
                        continue
                        
            except Exception as e:
                last_error = e
                print(f"Connection attempt {attempt}/{self._max_connection_retries} failed: {e}")
                
                # Stop the component before retrying
                await self._stop_component()
                
                if attempt < self._max_connection_retries:
                    retry_delay = min(2 ** attempt, self._max_reconnect_delay)
                    print(f"Retrying in {retry_delay}s...")
                    await asyncio.sleep(retry_delay)
                continue
        
        # All retries exhausted
        raise ConnectionError(
            f"Failed to connect after {self._max_connection_retries} attempts. Last error: {last_error}"
        )
    
    async def _reconnect_loop(self) -> None:
        """Reconnect loop with exponential backoff, triggered after disconnect"""
        delay = 1
        attempt = 0
        while self._should_reconnect and not self._is_connected:
            attempt += 1
            print(f"Reconnection attempt {attempt}, retrying in {delay}s...")
            await asyncio.sleep(delay)
            
            if not self._should_reconnect:
                break
            
            try:
                # Stop old component task
                await self._stop_component()
                
                # Reset connection state
                self._connection_error = None
                if self._connection_event:
                    self._connection_event.clear()
                
                # Start fresh component in background
                self._component_task = asyncio.create_task(self._run_component())
                
                # Wait for connection
                if self._connection_event:
                    await asyncio.wait_for(
                        self._connection_event.wait(),
                        timeout=self._first_connection_timeout
                    )
                    
                    if self._is_connected and self.session:
                        print(f"Reconnected successfully after {attempt} attempt(s)")
                        return
                    elif self._connection_error:
                        raise self._connection_error
                        
            except asyncio.CancelledError:
                return
            except Exception as e:
                print(f"Reconnection attempt {attempt} failed: {e}")
            
            # Exponential backoff capped at _max_reconnect_delay
            delay = min(delay * 2, self._max_reconnect_delay)
    
    async def _run_component(self) -> None:
        """Run the component in the background"""
        if not self.component:
            return
        try:
            await self.component.start()
        except asyncio.CancelledError:
            pass
        except Exception as e:
            print(f"Component error: {e}")
            self._connection_error = e
            if self._connection_event:
                self._connection_event.set()  # Signal to unblock waiters
    
    async def _stop_component(self) -> None:
        """Stop the component and cleanup"""
        try:
            if self._component_task and not self._component_task.done():
                self._component_task.cancel()
                try:
                    await asyncio.wait_for(self._component_task, timeout=5.0)
                except (asyncio.CancelledError, asyncio.TimeoutError):
                    pass
            
            if self.component:
                try:
                    await asyncio.wait_for(self.component.stop(), timeout=5.0)
                except asyncio.TimeoutError:
                    print("Component stop timed out")
                except Exception as e:
                    print(f"Error stopping component: {e}")
        except Exception as e:
            print(f"Error during component cleanup: {e}")
        finally:
            self._is_connected = False
            self.session = None
            
    async def stop(self) -> None:
        """Stop the connection"""
        self._should_reconnect = False
        
        # Cancel any running reconnect loop
        if self._reconnect_task and not self._reconnect_task.done():
            self._reconnect_task.cancel()
            try:
                await self._reconnect_task
            except asyncio.CancelledError:
                pass
            self._reconnect_task = None
        
        await self._stop_component()
        
        if self.component:
            # Check if component is properly initialized before stopping
            if hasattr(self.component, '_done_f') and self.component._done_f is not None:
                try:
                    await self.component.stop()
                except Exception:
                    pass
        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[CallOptions] = None
    ) -> Any:
        """Call a remote procedure"""
        # Validate parameters using Pydantic
        try:
            params = CrossbarCallParams(
                topic=topic,
                args=args or [],
                kwargs=kwargs or {},
                options=options
            )
        except Exception as e:
            raise ValueError(f"Invalid call parameters: {e}")
        
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        result = await self.session.call(params.topic, *(params.args or []), **(params.kwargs or {}), options=options)  # type: ignore[misc]
        return result
        
    async def subscribe(
        self, 
        topic: str, 
        handler: Callable,
        options: Optional[SubscribeOptions] = None
    ) -> Optional[Any]:
        """Subscribe to a topic"""
        # Validate parameters using Pydantic
        try:
            params = SubscriptionParams(
                topic=topic,
                options=options
            )
        except Exception as e:
            raise ValueError(f"Invalid subscription parameters: {e}")
        
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        sub = await self.session.subscribe(handler, params.topic, options=options)  # type: ignore[misc]
        if sub:
            self.subscriptions.append(SubscriptionEntry(topic=topic, handler=handler, options=options, sub=sub))
        return sub
        
    async def publish(
        self, 
        topic: str, 
        args: Optional[List[Any]] = None,
        kwargs: Optional[Dict[str, Any]] = None,
        options: Optional[PublishOptions] = None
    ) -> Optional[Any]:
        """Publish to a topic"""
        # Validate parameters using Pydantic
        try:
            params = CrossbarCallParams(  # Using CrossbarCallParams as it has the same structure
                topic=topic,
                args=args or [],
                kwargs=kwargs or {},
                options=options
            )
        except Exception as e:
            raise ValueError(f"Invalid publish parameters: {e}")
        
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        result = await self.session.publish(params.topic, *(params.args or []), options=options, **(params.kwargs or {}))  # type: ignore[misc]
        return result
        
    async def unsubscribe(self, subscription: Any) -> None:
        """Unsubscribe from a subscription. Accepts a Subscription, SubscriptionEntry, or dict."""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
        
        # Find matching entry
        entry = None
        for e in self.subscriptions:
            if e is subscription or e.sub is subscription:
                entry = e
                break
        
        if entry:
            self.subscriptions.remove(entry)
            if entry.sub:
                entry.sub.unsubscribe()
            
    async def unsubscribe_func(
        self, 
        topic: str, 
        handler: Optional[Callable] = None
    ) -> bool:
        """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 = [e for e in self.subscriptions if e.topic == topic and e.handler == handler]
                    
        for entry in to_remove:
            try:
                await self.unsubscribe(entry)
            except Exception as e:
                print(f"Failed to unsubscribe from {topic}: {e}")
                return False
        return True

    async def unsubscribe_topic(self, topic: str) -> bool:
        """Unsubscribe all handlers from a topic"""
        matching = [e for e in self.subscriptions if e.topic == topic]
        
        if not matching:
            return True
            
        for entry in matching:
            try:
                await self.unsubscribe(entry)
            except Exception as e:
                print(f"Failed to unsubscribe from {topic}: {e}")
        return True
        
    async def register(
        self, 
        topic: str, 
        endpoint: Callable,
        options: Optional[RegisterOptions] = None
    ) -> Optional[Any]:
        """Register a procedure"""
        if options is None:
            options = RegisterOptions(force_reregister=True)
        elif options.force_reregister is None:
            options.force_reregister = True

        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
            
        reg = await self.session.register(endpoint, topic, options=options)  # type: ignore[misc]
        if reg:
            self.registrations.append(RegistrationEntry(topic=topic, endpoint=endpoint, options=options, reg=reg))
        return reg
        
    async def unregister(self, registration: Any) -> None:
        """Unregister a procedure. Accepts a Registration, RegistrationEntry, or dict."""
        await self._session_wait()
        if not self.session:
            raise RuntimeError("No active session")
        
        entry = None
        for e in self.registrations:
            if e is registration or e.reg is registration:
                entry = e
                break
        
        if entry:
            self.registrations.remove(entry)
            if entry.reg:
                entry.reg.unregister()
            
    async def _resubscribe_all(self) -> None:
        """Resubscribe to all topics after reconnection using stored metadata"""
        if not self.subscriptions:
            return
            
        print(f'Resubscribing to {len(self.subscriptions)} subscription(s)')
        
        # Save and clear — subscribe() will re-populate the list with fresh entries
        old_entries = self.subscriptions.copy()
        self.subscriptions.clear()
        
        for entry in old_entries:
            try:
                await self.subscribe(entry.topic, entry.handler, entry.options)
            except Exception as e:
                print(f"Failed to resubscribe to {entry.topic}: {e}")
                # Preserve entry (without stale sub) so next reconnect can retry
                self.subscriptions.append(SubscriptionEntry(topic=entry.topic, handler=entry.handler, options=entry.options))

    async def _reregister_all(self) -> None:
        """Re-register all procedures after reconnection using stored metadata"""
        if not self.registrations:
            return
            
        print(f'Re-registering {len(self.registrations)} procedure(s)')
        
        old_entries = self.registrations.copy()
        self.registrations.clear()
        
        for entry in old_entries:
            try:
                await self.register(entry.topic, entry.endpoint, entry.options)
            except Exception as e:
                print(f"Failed to re-register {entry.topic}: {e}")
                self.registrations.append(RegistrationEntry(topic=entry.topic, endpoint=entry.endpoint, options=entry.options))