########################################################################
#
# Copyright (C) 2021, 2023
# Associated Universities, Inc. Washington DC, USA.
#
# This script is free software; you can redistribute it and/or modify it
# under the terms of the GNU Library General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This library is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Library General Public
# License for more details.
#
# You should have received a copy of the GNU Library General Public License
# along with this library; if not, write to the Free Software Foundation,
# Inc., 675 Massachusetts Ave, Cambridge, MA 02139, USA.
#
# Correspondence concerning AIPS++ should be adressed as follows:
#        Internet email: casa-feedback@nrao.edu.
#        Postal address: AIPS++ Project Office
#                        National Radio Astronomy Observatory
#                        520 Edgemont Road
#                        Charlottesville, VA 22903-2475 USA
#
########################################################################
'''Implementation of a general purpose event manager between Python and
JavaScript via ``websockets``. This provides a mechanism like the
``ImagePipe``. The difference is that ``ImagePipe`` is tuned for use
with CASA images but ``DataPipe`` can have generic messages.'''

import inspect
import threading
import asyncio
import traceback
import time
import json
from uuid import uuid4

from bokeh.models.sources import DataSource
from bokeh.util.compiler import TypeScript
from bokeh.core.properties import Tuple, String, Int, Instance, Nullable, Bool
from bokeh.models.callbacks import Callback

from ...utils import serialize, deserialize, is_interactive_jupyter, is_colab
from ..state import casalib_url, cubevisjs_url
from .. import BokehInit

class DataPipe(DataSource,BokehInit):
    """This class allows for communication between Python and the JavaScript implementation
    running in a browser. It allows Python code to send a message to JavaScript and register
    a callback which will receive the result. JavaScript code can do the same to request
    information from Python. Generally messages are expected to be dictionaries in the
    Python domain and objects in the JavaScript domain. UUIDs are used to keep messages in
    sync, and messages sent with the same UUID are queued until the currently pending
    message reply is recieved. A class can use multipe UUIDs to control queuing behavior.

    Attributes
    ----------
    address: tuple of string and int
        the string is the IP address for the network that should be used and the
        integer is the port number, see ``cubevis.utils.find_ws_address``
    init_script: JavaScript
        this javascript is run when this DataPipe object is initialized. init_script
        is used to run caller JavaScript which needs to be run at initialization time.
        This is optional and does not need to be set.
    """

    init_script = Nullable(Instance(Callback), help="""
    JavaScript to be run during initialization of an instance of an DataPipe object.
    """)

    backend_ip = String(default="127.0.0.1", help="the IP address that the Python server binds to locally" )
    backend_port = Int( help="the port that the Python server listens on" )
    backend_url = String( help="the full URL used by the browser to connect to the Python server" )

    instance_id = String( help="Unique ID for each DataPipe object" )

    conflict_check = Bool( default=True, help="Perform check to avoid reuse of URL for GUI. Not needed in the Jupyter context" )

    # Class-level session tracking to prevent multiple connections
    _active_sessions = {}  # session_id -> {'websocket': ws, 'timestamp': time, 'datapipe': instance}
    _session_lock = threading.Lock()

    ###################################################################
    ### filled from cubevis.bokeh.state._initialize._order_bokeh_js ###
    ###################################################################
    #__javascript__ = [ casalib_url( ), cubevisjs_url( ) ]

    def __init__( self, *args, abort=None, address=None, **kwargs ):

        if 'conflict_check' not in kwargs:
            kwargs['conflict_check'] = not is_interactive_jupyter( )
        if 'instance_id' not in kwargs:
            kwargs['instance_id'] = str(uuid4( ))

        if address is not None:
            kwargs['backend_ip'] = address[0]
            kwargs['backend_port'] = address[1]

        super( ).__init__( *args, **kwargs )

        self.__send_queue = { }
        self.__pending = { }
        self.__incoming_callbacks = { }
        self.__websocket = None
        self.__lock = threading.Lock( )
        self.__session = None
        self.__abort = abort

        if self.__abort is not None and not callable(self.__abort):
                raise RuntimeError(f'abort function must be callable ({type(self.__abort)} is not)')

        # Fetch URL in Colab
        if is_colab( ):
            from google.colab.output import eval_js
            # Colab maps the internal port to a secure external URL
            external_https = eval_js(f"google.colab.kernel.proxyPort({self.backend_port})")
            print(f"Colab Proxy URL for {self.backend_port} is: {external_https}")
            self.backend_url = external_https.replace("https://", "wss://")
        else:
            # Standard local/remote access
            self.backend_url = f"ws://{self.backend_ip}:{self.backend_port}"

    def __enqueue_send( self, ident, msg, callback ):
        ### it is assumed that this is called AFTER the lock has been aquired
        if ident in self.__send_queue:
            self.__send_queue[ident].insert(0, { 'cb': callback, 'msg': msg })
        else:
            self.__send_queue[ident] = [ { 'cb': callback, 'msg': msg } ]
    def __dequeue_send( self, ident ):
        ### it is assumed that this is called AFTER the lock has been aquired
        if ident in self.__send_queue and self.__send_queue[ident]:
            return self.__send_queue[ident].pop( )
        return None
    async def __put_pending( self, ident, callback ):
        ### it is assumed that this is called AFTER the lock has been aquired
        ## info about request sent to javascript waits in this queue
        ## until the javascript reply is received...
        if ident in self.__pending:
            if self.__websocket is not None:
                await self.__websocket.send( { 'id': '', 'message': 'queueing callback, but already one callback waiting', 'direction': 'error' } )
        else:
            self.__pending[ident] = callback

    def __get_pending( self, ident ):
        ### it is assumed that this is called AFTER the lock has been aquired
        ## when reply arrives for a request to javascript
        ## the callback is retrieved from the waiting queue
        if ident in  self.__pending:
            result = self.__pending[ident]
            del self.__pending[ident]
            return result
        return None

    @classmethod
    async def __is_websocket_alive(cls, websocket):
        """Check if a websocket connection is still alive"""
        try:
            # Try to ping the websocket with a short timeout
            await asyncio.wait_for(websocket.ping(), timeout=2.0)
            return True
        except (asyncio.TimeoutError, ConnectionError, Exception):
            return False

    @classmethod
    def __cleanup_dead_sessions(cls):
        """Remove dead sessions from tracking"""
        current_time = time.time()
        dead_sessions = []

        for session_id, session_info in cls._active_sessions.items():
            # Remove sessions older than 5 minutes (stale cleanup)
            if current_time - session_info['timestamp'] > 300:
                dead_sessions.append(session_id)

        for session_id in dead_sessions:
            del cls._active_sessions[session_id]

    async def __handle_session_conflict(self, websocket, existing_session_info, new_session_id):
        """Handle session conflict by checking if existing connection is alive"""
        existing_ws = existing_session_info['websocket']
        existing_datapipe = existing_session_info['datapipe']

        # Check if existing websocket is still alive
        if await self.__is_websocket_alive(existing_ws):
            # Existing connection is alive - send conflict message to BOTH connections
            conflict_msg = {
                'id': 'session_conflict',
                'message': {
                    'type': 'session_conflict',
                    'error': 'Multiple windows/tabs detected. Please use only one browser window.',
                    'action': 'close_duplicate'
                },
                'direction': 'error'
            }

            try:
                # Send to existing connection
                await existing_ws.send(serialize(conflict_msg))

                # Send to new connection
                await websocket.send(serialize(conflict_msg))

                # Close the new connection
                await websocket.close(code=1008, reason='Session conflict')

                # Call abort on the existing DataPipe instance
                if existing_datapipe.__abort is not None:
                    err = RuntimeError(f"Session conflict detected: New connection attempted with session {new_session_id}")
                    existing_datapipe.__abort(err)

                return False  # Reject new connection

            except Exception as e:
                print(f"Error handling session conflict: {e}")
                # If we can't communicate, treat existing as dead
                pass

        # Existing connection is dead - replace it
        return True  # Allow new connection

    def register( self, ident, callback ):
        """Register a callback to handle all requests coming from JavaScript. The
        callback will be called whenever a request arrives.

        Parameters
        ----------
        ident: string
            UUID which is associated with the messages that should be delivered
            to this callback
        callback: (string) => dictionary
            Callback which receives the message from Python as its sole parameter.
            The return value of this callback is delivered to the JavaScript code
            as the ''reply'' to to JavaScript in response to the ''request'' contained
            in the message.
        """
        with self.__lock:
            self.__incoming_callbacks[ident] = callback

    async def send( self, ident, message, callback ):
        """Send a `message` to JavaScript identified by `ident`. Once the reply is
        received, the `callback` will be called with the reply message.

        Parameters
        ----------
        ident: string
            UUID to associate with this message. It is used to keep track of the callback
            that should be called when the reply is received.
        message: dictionary
            This dictionary contains the request for the JavaScript code.
        callback: (string) => void
            Callback which receives the message that JavaScript generates in response to
            the request contained in the `message` parameter.
        """
        with self.__lock:
            if self.__websocket is not None:
                msg = { 'id': ident, 'message': message, 'direction': 'p2j' }
                if ident in self.__pending:
                    self.__enqueue_send( ident, msg, callback )
                else:
                    if ident in self.__send_queue and self.__send_queue[ident]:
                        self.__enqueue_send( ident, msg, callback )
                        existing = self.__dequeue_send(ident)
                        self.__put_pending(ident, existing['cb'])
                        await self.__websocket.send(serialize( existing['msg'] ))
                    else:
                        await self.__put_pending(ident, callback)
                        await self.__websocket.send(serialize( msg ))

    async def process_messages( self, websocket ):
        """Process messages related to image display updates.

        Parameters
        ----------
        websocket: websocket object
            Websocket serve function passes in the websocket object to use.
        path:
            Websocket serve provides a second parameter.
        """
        try:
            self.__websocket = websocket
            session_established = False

            print(f"✓ WebSocket server received connection on port {self.backend_port}")

            async for message in websocket:
                msg = deserialize(message)
                if 'session' not in msg:
                    await self.__websocket.close( )
                    err = RuntimeError(f'session not in: {msg}')
                    if self.__abort is not None:
                        self.__abort( err )
                    else:
                        raise err
                    return

                # Handle session initialization with conflict detection
                if not session_established:
                    new_session_id = msg['session']

                    with self._session_lock:
                        # Clean up any stale sessions first
                        self.__cleanup_dead_sessions()

                        # Check if session already exists
                        if new_session_id in self._active_sessions:
                            existing_session_info = self._active_sessions[new_session_id]

                            # Handle the conflict
                            if not await self.__handle_session_conflict(websocket, existing_session_info, new_session_id):
                                return  # Connection was rejected

                        # Register this session as active
                        self._active_sessions[new_session_id] = {
                            'websocket': websocket,
                            'timestamp': time.time(),
                            'datapipe': self
                        }
                        self.__session = new_session_id
                        session_established = True

                # Existing session validation (your original logic)
                elif self.__session != None and self.__session != msg['session']:
                    conflict_msg = {
                        'id': self.__session,
                        'message': {
                            'type': 'session_corruption',
                            'error': 'Session corruption detected',
                            'expected': self.__session,
                            'received': msg['session']
                        },
                        'direction': 'error'
                    }

                    await self.__websocket.send(serialize(conflict_msg))
                    await self.__websocket.close()

                    # Clean up from active sessions
                    with self._session_lock:
                        if self.__session in self._active_sessions:
                            del self._active_sessions[self.__session]

                    err = RuntimeError(f"session corruption: {msg['session']} does not equal {self.__session}")
                    if self.__abort is not None:
                        self.__abort( err )
                    else:
                        raise err
                    return

                with self.__lock:
                    ###
                    ### initialize session identifier (pre-state-corruption fixes)
                    ###
                    if self.__session == None:
                        self.__session = msg['session']

                    if msg['direction'] == 'p2j':
                        cb = self.__get_pending(msg['id'])
                        outgo = self.__dequeue_send(msg['id'])
                        if outgo is not None:
                            await websocket.send(serialize(outgo['msg']))
                            await self.__put_pending(msg['id'],outgo['cb'])
                        if cb is not None:
                            if inspect.isawaitable(cb):
                                await cb(msg['message'])
                            else:
                                cb(msg['message'])
                    elif msg['id'] != 'initialize':
                        ###
                        ### no cb/reply for initialization of session
                        ###
                        if msg['id'] not in self.__incoming_callbacks:
                            raise RuntimeError(f'incoming js request with no callback: {msg}')
                        result = self.__incoming_callbacks[msg['id']](msg['message'])
                        if inspect.isawaitable(result):
                            try:
                                return_message = "an exception occurred in creating the respone"
                                return_message = await result
                                await self.__websocket.send( serialize( { 'id': msg['id'],
                                                                          'message': return_message,
                                                                          'direction': msg['direction'] } ) )
                            except Exception as e:
                                trace_back = traceback.format_exc().replace('\n','\n                       ')
                                print('************************************************************************************************************************')
                                print( f'''EXCEPTION ENCOUNTERED: {repr(e)}''' )
                                print( f'''              MESSAGE: {repr(msg)}''' )
                                print( f'''          STACK TRACE: {trace_back}''' )
                                print('************************************************************************************************************************')
                                await self.__websocket.send( serialize( { 'id': msg['id'],
                                                                          'message': { 'error': "exception encountered",
                                                                                       'errant': str(return_message),
                                                                                       'exception': repr(e) },
                                                                          'direction': str(msg['direction']) } ) )
                        else:
                            try:
                                await self.__websocket.send( serialize( { 'id': msg['id'],
                                                                          'message': result,
                                                                          'direction': msg['direction'] } ) )
                            except Exception as e:
                                trace_back = traceback.format_exc().replace('\n','\n                       ')
                                print('************************************************************************************************************************')
                                print( f'''EXCEPTION ENCOUNTERED: {repr(e)}''' )
                                print( f'''              MESSAGE: {repr(msg)}''' )
                                print( f'''          STACK TRACE: {trace_back}''' )
                                print('************************************************************************************************************************')
                                await self.__websocket.send( serialize( { 'id': msg['id'],
                                                                          'message': { 'error': "exception encountered",
                                                                                       'errant': str(result),
                                                                                       'exception': repr(e) },
                                                                          'direction': str(msg['direction']) } ) )
        finally:
            # Clean up when connection closes
            if self.__session is not None:
                with self._session_lock:
                    if self.__session in self._active_sessions:
                        del self._active_sessions[self.__session]
            self.__websocket = None
