import websockets, json, asyncio, hmac, hashlib, os, re, html, datetime, async_timeout, time
import redis.asyncio as aioredis
from mysql import connector
import traceback
from dotenv import load_dotenv
from websockets.frames import CloseCode
import logging

logging.basicConfig(level=logging.INFO)

load_dotenv("/opt/.env")

class WebsocketServer:
    def __init__(self, host, port):
        self.shutdown_event = asyncio.Event()
        self.__host = host
        self.__port = port

        self.__secure_phrase = os.environ.get("WEBSOCKET_SECURE_PHRASE")
        self.__redis_pass = os.environ.get("REDIS_PASS")

        if not self.__secure_phrase:
            raise ValueError("Missing security variable!")
        
        if not self.__redis_pass:
            raise ValueError("Missing db security variable!")

        self.__clients = {}
        self.__available_channels = [
            'orders',
            'transactions',
            'test_data',
            'candles'
        ]

        self.__allowed_test_data_params = [
            'amount',
            'final_amount',
            'status',
            'equity',
            'position',
            'average_price',
            'timestamp',
            'market_price',
            'fee',
            'start_date'
        ]

    def __calculate_signature(self, user_id, timestamp, nonce, secret_phrase):
        return hmac.new(
            bytes(user_id, "latin-1"),
            msg=bytes(f'{timestamp}\n{nonce}\n{secret_phrase}', "latin-1"),
            digestmod=hashlib.sha256
        ).hexdigest().lower()

    def validate_auth_signature(self, input_signature=None, nonce=None, timestamp=None, user_id=None):
        if not input_signature or not nonce or not timestamp or not user_id:
            return False

        if not re.fullmatch(r"[a-f0-9]{32}", nonce):
            return False

        try:
            timestamp_value = int(timestamp)
            user_id_value = int(user_id)
        except (TypeError, ValueError):
            return False

        if user_id_value <= 0 or abs(int(time.time() * 1000) - timestamp_value) > 30_000:
            return False

        correct_signature = self.__calculate_signature(user_id, timestamp, nonce, self.__secure_phrase)

        return hmac.compare_digest(input_signature, correct_signature)

    async def authenticate(self, auth_str, redis):
        if not auth_str:
            return False
        auth_data = html.unescape(auth_str).split("&")
        if len(auth_data) != 4 or not self.validate_auth_signature(*auth_data):
            return False

        signature, nonce, timestamp, user_id = auth_data
        reserved = await redis.set(
            f"ws_auth_nonce:{user_id}:{nonce}", "1", nx=True, ex=60
        )
        return int(user_id) if reserved else False

    def _owns_test_sync(self, user_id, test_id):
        connection = connector.connect(
            host=os.environ.get('DB_HOST'),
            database=os.environ.get('DB_NAME'),
            user=os.environ.get('DB_USERNAME'),
            password=os.environ.get('DB_PASS'),
        )
        try:
            with connection.cursor() as cursor:
                cursor.execute(
                    "SELECT 1 FROM bee_tests WHERE id = %s AND owner_id = %s LIMIT 1",
                    (test_id, user_id),
                )
                return cursor.fetchone() is not None
        finally:
            connection.close()

    async def owns_test(self, user_id, test_id):
        return await asyncio.to_thread(self._owns_test_sync, user_id, test_id)

    def validate_input(self, user_input):
        if not isinstance(user_input, (str, list, dict, int, float, bool, type(None))):
            return None

        if isinstance(user_input, str):
            return re.sub(r'[^A-Za-z0-9ĂÂÎȘȚăâîșțЁёА-я\-@\.:\/№\\_°%,\s+]', "", html.escape(user_input.strip()))
        if isinstance(user_input, list):
            return [self.validate_input(item) for item in user_input]
        if isinstance(user_input, dict):
            return {key: self.validate_input(value) for key, value in user_input.items()}
        return user_input

    # async def connect_to_redis(self, websocket, data):
    #     channel = f"test.*.{data.get('test_id')}"
    #     redis = await aioredis.from_url(
    #         "redis://127.0.0.1",
    #         password=self.__redis_pass,
    #         encoding="utf-8",
    #         decode_responses=True
    #     )

    #     try:
    #         pubsub = redis.pubsub()
    #         await pubsub.psubscribe(channel)

    #         while True:
    #             message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
    #             if message:
    #                 ws_channel = message.get("channel").split(".")[1]

    #                 if ws_channel in data.get("channels"):
    #                     message_data = json.loads(message['data'])
    #                     result = {param: message_data.get(param) for param in data.get("test_params", []) if param in self.__allowed_test_data_params}

    #                     await websocket.send(json.dumps({
    #                         'result': {
    #                             'channel': ws_channel,
    #                             'data': result
    #                         }
    #                     }))
    #                     # await asyncio.sleep(0.01)
    #     except asyncio.CancelledError:
    #         print(f"{datetime.datetime.now()} | Subscription to {channel} canceled")
    #     except websockets.exceptions.ConnectionClosedError:
    #         print(f"{datetime.datetime.now()} | Closer Error")
    #         raise
    #     except websockets.exceptions.ConnectionClosedOK:
    #         print(f"{datetime.datetime.now()} | Closed Ok")
    #         raise
    #     finally:
    #         print(f"{datetime.datetime.now()} | Close subscription for {channel}")
    #         await pubsub.punsubscribe(channel)
    #         await pubsub.close()
    #         await redis.close()

    async def reader(self, channel: aioredis.client.PubSub, websocket, data: dict):

        print(f"Reader Listening on channel: {channel}")

        while True:
            try:
                async with async_timeout.timeout(1):
                    message = await channel.get_message(ignore_subscribe_messages=True)
                    if message is not None:
                        
                        # print(f"(Reader) Message Received: {message}")
                        # await asyncio.sleep(0.001)
                        ws_channel = message.get("channel").split(".")[1]

                        if ws_channel in data.get("channels"):

                            # if message["data"].decode() == STOPWORD:
                            #     print("(Reader) STOP")
                            #     break

                            message_data = result = json.loads(message['data'])

                            if data.get("test_params"):
                                for param in data.get("test_params"):
                                    if param in self.__allowed_test_data_params:
                                        result[param] = message_data.get(param)

                            await websocket.send(json.dumps({
                                'result': {
                                    'channel': ws_channel,
                                    'data': result
                                }
                            }))

                    await asyncio.sleep(0.01)
            except asyncio.TimeoutError:
                pass

    async def register(self, websocket):

        client_address = websocket.request_headers.get('X-Forwarded-For')
        client_key = id(websocket)

        redis = await aioredis.from_url(
            "redis://127.0.0.1",
            password=self.__redis_pass,
            encoding="utf-8",
            decode_responses=True
        )
        pubsub = redis.pubsub()

        self.__clients[client_key] = {
            "websocket": websocket,
            "redis": redis,
            "pubsub": pubsub,
            "tasks": {}
        }

        print("\n-------------- NEW CONNECTION --------------\n")
        print(f"{datetime.datetime.now()} | Client: {client_address}")

        try:
            authenticated_user_id = None
            while True:
                data = json.loads(await websocket.recv())
                if authenticated_user_id is None:
                    authenticated_user_id = await self.authenticate(data.get('auth'), redis)
                    if not authenticated_user_id:
                        raise ValueError("Invalid Authentication")
                data = self.validate_input(data)

                if 'channels' not in data:
                    raise ValueError("Missing Parameter: channels")
                
                for channel in data.get('channels'):
                    if channel not in self.__available_channels:
                        raise ValueError(f"Invalid Channel: {channel}")

                if 'test_id' not in data:
                    raise ValueError("Missing Parameter: test_id")

                test_id = int(data.get('test_id'))
                if not await self.owns_test(authenticated_user_id, test_id):
                    raise ValueError("Invalid Test")

                if "test_data" in data.get('channels') and "test_params" not in data:
                    raise ValueError("Missing Parameter: test_params")
                
                channel = f"test.*.{test_id}"

                if channel not in self.__clients[client_key]['tasks']:
                    await pubsub.psubscribe(channel)

                    print(f"Subscribed on channel: {channel}")

                    task = asyncio.create_task(self.reader(pubsub, websocket, data))

                    self.__clients[client_key]['tasks'][channel] = task

        except websockets.exceptions.ConnectionClosed:
            print(f"{datetime.datetime.now()} | Client {client_address} closed the connection")
        except websockets.exceptions.ConnectionClosedError:
            print(f"{datetime.datetime.now()} | Client {client_address} dropped the connection unexpectedly")
        except websockets.exceptions.ConnectionClosedOK:
            print(f"{datetime.datetime.now()} | Client {client_address} closed the connection gracefully")
        except websockets.exceptions.InvalidHandshake as e:
            print(f"{datetime.datetime.now()} | InvalidHandshake connection error: {e}")
        except TimeoutError as e:
            print(f"{datetime.datetime.now()} | Timeout receiving data from WebSocket: {e}")
        except ValueError as ve:
            print(f"{datetime.datetime.now()} | ValueError: {ve}")
            traceback.print_exc()
            await websocket.close(code=CloseCode.INVALID_DATA, reason=str(ve))
            await websocket.wait_closed()
        except Exception as e:
            print(f"{datetime.datetime.now()} | Error: {e}")
            traceback.print_exc()
            await websocket.close(code=CloseCode.INTERNAL_ERROR, reason="Internal Server Error")
            await websocket.wait_closed()
        finally:
            print(f"{datetime.datetime.now()} | Initiate Disconnect")

            client = self.__clients.get(client_key)
            if client and len(client['tasks']):
                for channel, task in client['tasks'].items():
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass
                    except websockets.exceptions.ConnectionClosed:
                        pass

            await pubsub.close()
            await redis.close()

            self.__clients.pop(client_key, None)
            print("\n---------------------------------\n")

    async def start_server(self):
        self.__server = await websockets.serve(self.register, self.__host, self.__port)
        try:
            os.makedirs('/run/razgar', exist_ok=True)
            with open('/run/razgar/bee-ws.ready', 'w', encoding='utf-8') as marker:
                marker.write('ready\\n')
        except OSError as error:
            raise RuntimeError('BEE websocket readiness marker could not be written') from error
        await self.shutdown_event.wait()

    async def shutdown(self):
        self.shutdown_event.set()
        """Gracefully shut down the WebSocket server."""
        print(f"{datetime.datetime.now()} | Shutting down the WebSocket server")

        # Close all active client connections
        if len(self.__clients):
            for client_key, client_info in list(self.__clients.items()):
                websocket = client_info["websocket"]

                # Cancel all tasks associated with this client
                if len(client_info['tasks']):
                    for channel, task in client_info["tasks"].items():
                        task.cancel()
                        try:
                            await task  # Await task to ensure any necessary cleanup is performed
                        except asyncio.CancelledError:
                            pass

                await client_info["pubsub"].close()
                await client_info["redis"].close()

                # Close the WebSocket connection gracefully
                try:
                    await websocket.close(code=CloseCode.GOING_AWAY, reason="Server shutting down")
                    await websocket.wait_closed()
                except Exception as e:
                    print(f"{datetime.datetime.now()} | Error closing WebSocket for client {client_key}: {e}")

                # Remove client from the dictionary
                self.__clients.pop(client_key, None)

        # Stop the WebSocket server
        if self.__server:
            self.__server.close()
            await self.__server.wait_closed()

        try:
            os.unlink('/run/razgar/bee-ws.ready')
        except FileNotFoundError:
            pass

        print(f"{datetime.datetime.now()} | WebSocket server shutdown complete")


def main():
    server = WebsocketServer(host="localhost", port=8765)
    loop = asyncio.get_event_loop()

    try:
        loop.run_until_complete(server.start_server())
    except KeyboardInterrupt:
        print("\nKeyboard interrupt received, shutting down...")
    finally:
        loop.run_until_complete(server.shutdown())
        loop.close()

if __name__ == "__main__":
    main()