import base64
import json
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List

from aio_pika.abc import AbstractRobustConnection

from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId
from amqp.router.utils import sanitize_as_rabbitmq_name

RS = ":"


"""
Define the structure of messages used in Clever Micro
"""


class DataMessageMagicByte(Enum):
    HTTP_REQUEST = "A"
    RPC_REQUEST = "R"


@dataclass
class CleverMicroMessage:
    DEFAULT_CONTENT_TYPE = ["application/octet-stream"]

    """
    The common structure of a message that transfers the user request to the Service and back. All data related messages
    conform to this base structure:

    magic - identifies the message type (REST, RPC, DataResponse)
    id - unique message ID
    m_field - method name (for REST style calls) or service name (for RPC style calls) or response code for DataResponse
    d_field - domain name (for REST style calls) or module name (for RPC style calls) or  error code for DataResponse
    p_field - path (for REST style calls) or method name (for RPC style calls) or error cause for DataResponse
    headers - headers of the message (e.g. Content-Type, Accept, etc.)
    trace_info - trace information (e.g. trace ID, span ID, etc. for Jaeger trace monitor)
    base64body - the body of the message encoded in base64
    """

    def __init__(
        self,
        magic: str,
        id: SnowflakeId,
        m_field: str,
        d_field: str,
        p_field: str,
        headers: Dict[str, List[str]],
        trace_info: Dict[str, str],
        base64body: bytes,
    ):
        self._magic = magic
        self._id: SnowflakeId = id
        self._m_field = m_field
        self._d_field = d_field
        self._p_field = p_field
        self._headers: Dict[str, List[str]] = headers
        self._trace_info: Dict[str, str] = trace_info
        self._base64body: bytes = base64body

    def magic(self):
        return self._magic

    def id(self):
        return self._id

    def m_field(self) -> str:
        return self._m_field

    def d_field(self) -> str:
        return self._d_field

    def p_field(self) -> str:
        return self._p_field

    def headers(self) -> Dict[str, List[str]]:
        return self._headers

    def trace_info(self) -> Dict[str, str]:
        return self._trace_info

    def base64body(self) -> bytes:
        return self._base64body

    def body(self) -> bytes:
        return base64.b64decode(self.base64body())

    def content_type(self) -> str:
        return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0]


class DataMessage(CleverMicroMessage):
    """
    The class represents a structure of a message that transfers the user request to the Service.
    It may also be used to make a call/request between services.
    The response to this call is in format of DataResponse class.

    See DataMessageFactory for code to instantiate, serialize and deserialize the class
    """

    def __init__(
        self,
        magic: str,
        id: SnowflakeId,
        method: str,
        domain: str,
        path: str,
        headers: Dict[str, List[str]],
        trace_info: Dict[str, str],
        base64body: bytes,
    ):
        super().__init__(magic, id, method, domain, path, headers, trace_info, base64body)

    def method(self) -> str:
        return self.m_field()

    def domain(self) -> str:
        return self.d_field()

    def path(self) -> str:
        return self.p_field()

    def routing_key(self) -> str:
        return sanitize_as_rabbitmq_name(f"{self.d_field()}.{self.p_field()}")


class MultipartDataMessage(DataMessage):
    def __init__(self, message: DataMessage, extra_data: dict):
        super().__init__(
            message.magic(),
            message.id(),
            message.method(),
            message.domain(),
            message.path(),
            message.headers(),
            message.trace_info(),
            message.base64body(),
        )
        self.extra_data = extra_data


@dataclass
class AMQMessage(DataMessage):
    """
    Expand the core DataMessage type to include the channel, which is required to create a dedicated queue
    in case the response should initiate file download.
    """

    _connection: AbstractRobustConnection = field(repr=False)  # Don't show in repr

    def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection):
        # Initialize the parent DataMessage with all properties from the input data_message
        super().__init__(
            magic=data_message.magic(),
            id=data_message.id(),
            method=data_message.method(),
            domain=data_message.domain(),
            path=data_message.path(),
            headers=data_message.headers(),
            trace_info=data_message.trace_info(),
            base64body=data_message.base64body(),
        )
        self._connection = connection
        self.extra_data = data_message.extra_data if hasattr(data_message, "extra_data") else {}

    @property
    def connection(self) -> AbstractRobustConnection:
        """Get the AMQP connection associated with this message"""
        return self._connection


class DataResponse(CleverMicroMessage):
    """
    The response message to REST & RPC style calls (note: the request is made with DataMessage)

    See DataResponseFactory for code to instantiate, serialize and deserialize the class
    """

    def __init__(
        self,
        magic: str,
        id: SnowflakeId,
        response_code: str,
        error: str,
        error_cause: str,
        headers: Dict[str, List[str]],
        trace_info: Dict[str, str],
        base64body: bytes,
    ):
        super().__init__(
            magic,
            id,
            response_code,
            error,
            error_cause,
            headers,
            trace_info,
            base64body,
        )

    def response_code(self) -> str:
        return self.m_field()

    def error(self) -> str:
        return self.d_field()

    def error_cause(self) -> str:
        return self.p_field()


@dataclass(frozen=True)
class AMQRoute:
    """
    Configuration item representing one AMQRoute. Adapter uses this to declare the Exchange
    and/or queue and bind the key as routing expression.

    :param component_name: a service/component name that this route points to (owner of the route)
    :param exchange: RabbitMQ exchange name
    :param queue: RabbitMQ queue name
    :param key: RabbitMQ TOPIC Exchange routing key
    :param timeout: Timeout in milliseconds, how long to wait for reply, or not waiting if less then 1
    :param valid_until: optional timestamp (millis from epoch) after which the route is ignored
    """

    component_name: str
    exchange: str
    queue: str
    key: str
    timeout: int
    valid_until: int

    FORMAT = "Component-Name:Exchange-Name:Queue-Name:Routing-Key:Timeout:Valid-Until"

    def __str__(self):
        return (
            f"{self.component_name}{RS}{self.exchange}{RS}{self.queue}{RS}{self.key}"
            f"{RS}{self.timeout}{RS}{self.valid_until}"
        )

    def __eq__(self, other):
        if isinstance(other, AMQRoute):
            return (
                self.component_name == other.component_name
                and self.queue == other.queue
                and self.key == other.key
            )
        return False

    @property
    def primary_key(self):
        """
        Build the unique identifier of the route
        :return primary key
        """
        return f"{self.component_name}{RS}{self.queue}{RS}{self.key}"

    def as_string(self) -> str:
        return self.__str__()


# Ensure backward compatibility
AMQResponse = DataResponse


@dataclass(frozen=True)
class ServiceMessage:
    """
    The CleverMicro service message that is used to communicate the internal state of the AMQ channel.
    In v0.2 it supports service discovery / route registration and Backpressure status notification.
    """

    id: SnowflakeId = None
    payload_type: int = 0
    message_type: ServiceMessageType = ServiceMessageType.UNKNOWN
    recipient_name: str = None
    reply_to: str = None
    trace_info: Dict[str, str] = field(default_factory=dict)
    message: str = None

    def is_valid(self) -> bool:
        return self.message_type is not None


class ScalingRequestAlert(Enum):
    OVERLOAD = "OVERLOAD"
    IDLE = "IDLE"
    UPDATE = "UPDATE"
    SCALE_UP = "SCALE_UP"
    SCALE_DOWN = "SCALE_DOWN"


@dataclass(frozen=True)
class ScalingRequest:
    service_id: str
    task_id: str
    max_availability: int
    current_availability: int
    request_type: ScalingRequestAlert
    version: int = 1

    def to_json(self) -> str:
        """Converts the object to a JSON string matching the Java structure"""
        return json.dumps(
            {
                "version": self.version,
                "serviceId": self.service_id,
                "taskId": self.task_id,
                "maxAvailability": self.max_availability,
                "currentAvailability": self.current_availability,
                "requestType": self.request_type.value,  # Using .value for Enum serialization
            },
            indent=2,
        )

    @classmethod
    def from_json(cls, json_str: str) -> "ScalingRequest":
        """Creates a ScalingRequest from a JSON string"""
        data = json.loads(json_str)
        return cls(
            service_id=data["serviceId"],
            task_id=data["taskId"],
            max_availability=data["maxAvailability"],
            current_availability=data["currentAvailability"],
            request_type=ScalingRequestAlert(data["requestType"]),  # Convert string to Enum
        )


class AMQErrorMessage:
    """
    The error response message to RPC style call (note: the request is made with DataMessage)
    """

    def __init__(
        self,
        response_code: int,
        error: str,
        detail: str,
        content_type: str = "application/json",
    ):
        self.response_code = response_code
        self.error = error
        self.detail = detail
        self.content_type = content_type

    def serialize(self) -> bytes:
        return json.dumps(
            {
                "response_code": self.response_code,
                "error": self.error,
                "detail": self.detail,
            }
        ).encode("utf-8")