"""
methods to create, serialize and deserialize DataMessage objects.
"""

import base64
import json
import re
from datetime import datetime, timezone
from io import BytesIO
from typing import Any, Dict, List

import amqp.adapter.file_handler
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.pydantic_serializer import python_type_to_json
from amqp.adapter.serializer import (
    map_as_string,
    map_with_list_as_string,
    parse_map_list_string,
    parse_map_string,
)
from amqp.model.model import DataMessage, DataMessageMagicByte
from amqp.model.snowflake_id import SnowflakeId
from amqp.router.utils import sanitize_as_rabbitmq_name


class DataMessageFactory:
    SNOWFLAKE_ID_BITS = 80
    SEQUENCE_BITS = 12
    MACHINE_ID_BITS = 24
    TIMESTAMP_BITS = SNOWFLAKE_ID_BITS - SEQUENCE_BITS - MACHINE_ID_BITS

    MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1
    SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1

    _instance = None
    snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp()

    def __init__(self, machine_id: int):
        self.machine_id = (machine_id & self.MACHINE_ID_MASK) << self.SEQUENCE_BITS
        self.snowflake_sequence = 0
        self.snowflake_last_timestamp = 0

    @classmethod
    def get_instance(cls, machine_id: int):
        """
        return reusable factory instance
        :param machine_id: a generator ID, should be unique for each component to prevent duplicate IDs
        :return: factory instance
        """
        if cls._instance is None:
            cls._instance = DataMessageFactory(machine_id)
        return cls._instance

    def generate_snowflake_id(self) -> SnowflakeId:
        """
        Generate an unique ID following the SnowflakeID structure, see https://en.wikipedia.org/wiki/Snowflake_ID
        :return: 80-bit wide ID
        """
        timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)

        if timestamp < self.snowflake_last_timestamp:
            raise RuntimeError("Clock moved backwards. Refusing to generate id")

        if timestamp == self.snowflake_last_timestamp:
            self.snowflake_sequence = (self.snowflake_sequence + 1) & self.SEQUENCE_MASK
            if self.snowflake_sequence == 0:
                while timestamp == self.snowflake_last_timestamp:
                    timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)
                self.snowflake_last_timestamp = timestamp
        else:
            self.snowflake_sequence = 0
            self.snowflake_last_timestamp = timestamp

        time = timestamp - 1000 * int(DataMessageFactory.snowflake_epoch)
        return SnowflakeId(
            time >> (64 - self.SEQUENCE_BITS - self.MACHINE_ID_BITS),
            (
                (time << (self.SEQUENCE_BITS + self.MACHINE_ID_BITS))
                | self.machine_id
                | self.snowflake_sequence
            )
            & 0xFFFFFFFFFFFFFFFF,
        )

    def create_request_message(
        self,
        method: str,
        domain: str,
        path: str,
        headers: Dict[str, List[str]],
        message,
    ) -> DataMessage:
        """
        Encapsulate / hide the system level inputs, create the AMQ message from business relevant input
        :param method: a method requested (e.g. HTTP method like GET, POST in case of REST request)
        :param domain: service host name
        :param path: path at the service
        :param headers: HTTP headers
        :param message: payload
        :return: DataMessage object
        """
        serializable_headers = headers.copy() if headers else {}
        payload = base64.b64encode(message.encode("utf-8") if isinstance(message, str) else message)
        trace_info = {}
        return DataMessage(
            DataMessageMagicByte.HTTP_REQUEST.value,
            self.generate_snowflake_id(),
            method,
            domain,
            path,
            serializable_headers,
            trace_info,
            payload,
        )

    def safe_rsplit(self, fq_method_name, default="_"):
        """
        Safely split a fully qualified method name into package, class, and method

        Args:
            fq_method_name (str): Fully qualified method name
            default (str, optional): Default value for missing tokens. Defaults to ''.

        Returns:
            tuple: (package, class_name, method)
        """
        if not fq_method_name:
            return default, default, default
        parts = fq_method_name.rsplit(".", 2)
        # Pad with default values if insufficient tokens
        while len(parts) < 3:
            parts.insert(0, default)
        return tuple(parts)

    def create_rpc_message(
        self,
        fq_method_name: str,
        swarm_id: str,
        args: Dict[str, Any] = None,
    ) -> DataMessage:
        """
        Create a RPC message with the given parameters.
        :param fq_method_name: fully qualified method name (e.g. 'com.example.ServiceClass.method_name')
        :param args: arguments whose serialized form will form the request body
        :return: DataMessage object
        """
        package, class_name, method = self.safe_rsplit(fq_method_name)
        headers = {"Content-Type": ["application/octet-stream"], "Return-Address": [swarm_id]}
        trace_info = {}
        body = python_type_to_json(args).encode(encoding="utf-8") if args else b""
        print(
            f"Creating RPC message for method: {fq_method_name}, args: {args} body: {body.decode('utf-8') if body else 'None'}"
        )

        return DataMessage(
            DataMessageMagicByte.RPC_REQUEST.value,
            self.generate_snowflake_id(),
            package,
            class_name,
            method,
            headers,
            trace_info,
            body,
        )

    @staticmethod
    def amq_message_routing_key(message: DataMessage) -> str:
        """
        Constructs the message specific routing key.
        :param message: DataMessage input
        :return: routing key (as string) compliant to RabbitMQ allowed character set.
        """
        return sanitize_as_rabbitmq_name(f"{message.domain()}.{message.path()}")

    @staticmethod
    def get_timestamp_from_id(_id: SnowflakeId) -> int:
        """
        Helper method - retrieve the timestamp from the ID
        :param _id: SnowflakeID (includes encoded timestamp)
        :return: the message timestamp
        """
        bits = _id.get_bits()
        low = bits[0] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
        high = bits[1] << (
            64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS
        )
        return (high | low) + DataMessageFactory.snowflake_epoch * 1000

    @staticmethod
    def to_string(message: DataMessage) -> str:
        """
        Human-readable representation.
        :param message: DataMessage input
        :return: as string
        """
        return (
            f"DataMessage{{"
            f"id={message.id()}, "
            f"method='{message.method()}', "
            f"domain='{message.domain()}', "
            f"path='{message.path()}', "
            f"headers={{{map_with_list_as_string(message.headers())}}}, "
            f"traceInfo={{{map_as_string(message.trace_info())}}}, "
            f"base64body={message.base64body().decode()}"
            f"}}"
        )

    @staticmethod
    def serialize(message: DataMessage) -> bytes:
        """
        Converts DataMessage to byte array
        :param message: message to serialize
        :return: messages as bytes
        """
        id_bytes = str(message.id()).encode()
        prolog = (
            f"|m={message.method()}|d={message.domain()}|p={message.path()}|h={{{map_with_list_as_string(message.headers())}}}|t={{{map_as_string(message.trace_info())}}}|b="
        ).encode("utf-8")
        header_length = 1 + len(id_bytes) + len(prolog)
        prolog_len = header_length.to_bytes(2, byteorder="big")
        msg_length = 2 + header_length + len(message.base64body())
        baos = BytesIO(bytearray(msg_length))
        baos.write(prolog_len)
        baos.write(message.magic().encode("utf-8"))
        baos.write(id_bytes)
        baos.write(prolog)
        baos.write(message.base64body())
        return baos.getvalue()

    @staticmethod
    def from_bytes(input_bytes: bytes) -> DataMessage:
        """
        Builds the DataMessage from its serialized (byte array) form.
        :param input_bytes: serialized message
        :return: DataMessage instance
        """
        prolog_len = (input_bytes[0] << 8) + input_bytes[1]
        metadata = input_bytes[2 : prolog_len + 2].decode()

        pattern = re.compile(
            r"([AR])"
            r"([0-9A-Fa-f]+)\|"
            r"m=([^|]*)\|"
            r"d=([^|]*)\|"
            r"p=([^|]*)\|"
            r"h=(\{[^}]*\})\|"
            r"t=(\{[^}]*\})\|"
            r"b="
        )
        match = pattern.match(metadata)
        if not match:
            raise ValueError(
                f"DataMessage format does not match expected format. "
                f"Input string [{metadata}] does not match pattern: {pattern.pattern}"
            )

        id_str = match.group(2)
        method = match.group(3)
        domain = match.group(4)
        path = match.group(5)
        headers_str = match.group(6)
        trace_info_str = match.group(7)
        body_b64 = input_bytes[prolog_len + 2 :]

        headers = parse_map_list_string(headers_str)
        trace_info = parse_map_string(trace_info_str)

        return DataMessage(
            match.group(1),  # Magic byte (R for RPC, A for HTTP)
            SnowflakeId.from_hex(id_str),
            method,
            domain,
            path,
            headers,
            trace_info,
            body_b64,
        )

    @staticmethod
    async def from_stream(
        stream: bytes,
        rabbitmq_url: str,
        loop,
        file_handler: FileHandler = StreamingFileHandler(),
    ) -> DataMessage | None:
        """
        Builds the DataMessage from its serialized (byte array) form.
        :param stream: serialized message
        :param rabbitmq_url: RabbitMQ connection URL
        :param loop: event loop
        :param file_handler: file handler for downloading the message
        :return: DataMessage instance
        """
        # async implementation
        _message_data = json.loads(stream.decode())
        _req_info = _message_data.get("req-info", "")
        if _req_info:
            (body, eof) = await file_handler.download_buffer(
                loop=loop, rabbitmq_url=rabbitmq_url, queue_name=_req_info
            )
            if eof == amqp.adapter.file_handler.END_OF_FILE:
                return DataMessageFactory.from_bytes(body)
        return None