from __future__ import annotations

import json
import uuid
from typing import TYPE_CHECKING, Any, Optional

from mockmesh.core.logger import get_logger
from mockmesh.interceptors.responses import FakeAzureResponse

if TYPE_CHECKING:
    from mockmesh.storage.manager import StorageManager

_log = get_logger("interceptors.azure.servicebus")


class _CtxMixin:
    def __enter__(self):
        return self

    def __exit__(self, *_: Any) -> None:
        self.close()  # type: ignore[attr-defined]


class _MockServiceBusReceivedMessage:
    def __init__(self, body_bytes: bytes, nosql_key: str,
                 queue_key: str, storage: Any) -> None:
        self._body_bytes  = body_bytes
        self._nosql_key   = nosql_key
        self._queue_key   = queue_key
        self._storage     = storage
        self.message_id   = nosql_key
        self.sequence_number = 0

    @property
    def body(self):
        yield self._body_bytes


class _MockServiceBusSender(_CtxMixin):
    def __init__(self, queue_name: str, storage: Any) -> None:
        self._queue   = queue_name
        self._storage = storage

    def send_messages(self, message: Any, timeout: Any = None) -> None:
        msgs = message if isinstance(message, (list, tuple)) else [message]
        for msg in msgs:
            body   = _extract_sb_body(msg)
            msg_id = str(uuid.uuid4())
            self._storage.nosql.put(
                f"sb_{self._queue}", msg_id,
                {"MessageId": msg_id, "Body": body, "status": "queued"},
            )
        _log.info("ServiceBus send_messages (mock)",
                  extra={"mm_queue": self._queue, "mm_count": len(msgs)})

    def close(self) -> None:
        pass


class _MockServiceBusReceiver(_CtxMixin):
    def __init__(self, queue_name: str, storage: Any,
                 max_wait_time: Any = None) -> None:
        self._queue   = queue_name
        self._storage = storage

    def receive_messages(self, max_message_count: int = 1,
                         max_wait_time: Any = None) -> list:
        q_key   = f"sb_{self._queue}"
        records = self._storage.nosql.query(q_key, status="queued")
        result  = []
        for rec in records[:max_message_count]:
            body_str   = rec["value"].get("Body", "")
            body_bytes = body_str.encode("utf-8") if isinstance(body_str, str) else body_str
            self._storage.nosql.put(
                q_key, rec["key"],
                {**rec["value"], "status": "locked"},
            )
            result.append(_MockServiceBusReceivedMessage(
                body_bytes=body_bytes,
                nosql_key=rec["key"],
                queue_key=q_key,
                storage=self._storage,
            ))
        _log.info("ServiceBus receive_messages (mock)",
                  extra={"mm_queue": self._queue, "mm_count": len(result)})
        return result

    def complete_message(self, msg: "_MockServiceBusReceivedMessage") -> None:
        self._storage.nosql.put(
            msg._queue_key, msg._nosql_key,
            {**self._storage.nosql.get(msg._queue_key, msg._nosql_key),
             "status": "completed"},
        )

    def close(self) -> None:
        pass


class _MockServiceBusClient(_CtxMixin):
    def __init__(self, storage: Any) -> None:
        self._storage = storage

    def get_queue_sender(self, queue_name: str, **_: Any) -> _MockServiceBusSender:
        return _MockServiceBusSender(queue_name, self._storage)

    def get_queue_receiver(self, queue_name: str,
                           max_wait_time: Any = None, **_: Any) -> _MockServiceBusReceiver:
        return _MockServiceBusReceiver(queue_name, self._storage, max_wait_time)

    def close(self) -> None:
        pass


def _extract_sb_body(msg: Any) -> str:
    body_attr = getattr(msg, "body", None)
    if isinstance(body_attr, str):
        return body_attr
    if isinstance(body_attr, bytes):
        return body_attr.decode("utf-8", errors="replace")

    raw_amqp = getattr(msg, "_raw_amqp_message", None)
    if raw_amqp is not None:
        amqp_body = getattr(raw_amqp, "body", None)
        if isinstance(amqp_body, bytes):
            return amqp_body.decode("utf-8", errors="replace")
        if isinstance(amqp_body, str):
            return amqp_body
        if amqp_body is not None and hasattr(amqp_body, "__iter__"):
            try:
                chunks = []
                for chunk in amqp_body:
                    chunks.append(
                        bytes(chunk).decode("utf-8", errors="replace")
                        if isinstance(chunk, (bytes, bytearray)) else str(chunk)
                    )
                if chunks:
                    return "".join(chunks)
            except Exception:
                pass

    if body_attr is not None and hasattr(body_attr, "__iter__"):
        try:
            chunks = []
            for chunk in body_attr:
                chunks.append(
                    bytes(chunk).decode("utf-8", errors="replace")
                    if isinstance(chunk, (bytes, bytearray)) else str(chunk)
                )
            return "".join(chunks)
        except Exception:
            pass

    return ""


def handle_servicebus(
    sm: "StorageManager",
    parsed: dict,
    request: Any,
    method: str,
    url: str,
    default_responses: dict,
) -> Optional[FakeAzureResponse]:
    from mockmesh.interceptors.azure.utils import _get_body, _enc

    entity = parsed["container"]
    q_key  = "sb_" + entity
    tpl    = default_responses.get("servicebus", {})

    if method == "POST" and "messages" in url:
        raw     = _get_body(request)
        msg_id  = str(uuid.uuid4())
        content = raw.decode("utf-8", errors="replace") if raw else ""
        sm.nosql.put(q_key, msg_id,
                     {"MessageId": msg_id, "Body": content, "status": "queued"})
        status = tpl.get("POST", {}).get("status", 201)
        return FakeAzureResponse(
            status, {"Location": "/" + entity + "/messages/" + msg_id}, b"")

    if method in ("GET", "DELETE") and "messages" in url:
        records = sm.nosql.query(q_key, status="queued")
        if records:
            rec = records[0]
            sm.nosql.put(q_key, rec["key"], {**rec["value"], "status": "locked"})
            body_bytes = rec["value"].get("Body", "").encode("utf-8")
            bp_tpl = dict(tpl.get("broker_properties_template", {}))
            broker = {**bp_tpl,
                      "MessageId": rec["key"],
                      "LockToken": str(uuid.uuid4())}
            status = tpl.get("GET", {}).get("status", 200)
            return FakeAzureResponse(
                status, {"BrokerProperties": json.dumps(broker)}, body_bytes)
        status = tpl.get("NO_MESSAGE", {}).get("status", 204)
        return FakeAzureResponse(status, {}, b"")

    return None
