from __future__ import annotations

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

from mockmesh.interceptors.azure.utils import _enc, _get_body
from mockmesh.interceptors.responses import FakeAzureResponse

if TYPE_CHECKING:
    from mockmesh.storage.manager import StorageManager

_JSON_HDR = {"Content-Type": "application/json"}


class _CtxMixin:
    def __enter__(self):
        return self

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


class _MockEventData:
    def __init__(self, body: str = "") -> None:
        self.body = body
        self.properties: dict = {}
        self.system_properties: dict = {}


class _MockEventDataBatch:
    def __init__(self) -> None:
        self._events: list[_MockEventData] = []

    def add(self, event_data: Any) -> None:
        self._events.append(event_data)


class _MockEventHubProducerClient(_CtxMixin):
    def __init__(self, hub_name: str, storage: Any) -> None:
        self._hub = hub_name
        self._storage = storage

    def create_batch(self) -> _MockEventDataBatch:
        return _MockEventDataBatch()

    def send_batch(self, batch: Any) -> None:
        events = batch._events if hasattr(batch, "_events") else [batch]
        q_key = f"eh_{self._hub}"
        for ev in events:
            msg_id = str(uuid.uuid4())
            if hasattr(ev, "body_as_str"):
                body = ev.body_as_str()
            elif isinstance(ev.body, str):
                body = ev.body
            else:
                body = str(ev.body)
            self._storage.nosql.put(q_key, msg_id, {
                "MessageId": msg_id, "Body": body, "status": "queued",
            })

    def close(self) -> None:
        pass


class _MockEventHubConsumerClient(_CtxMixin):
    def __init__(self, hub_name: str, consumer_group: str = "$Default",
                 storage: Any = None) -> None:
        self._hub = hub_name
        self._consumer_group = consumer_group
        self._storage = storage

    def receive_batch(self, partition_id: str = "0",
                      max_batch_size: int = 1) -> list:
        q_key = f"eh_{self._hub}"
        records = self._storage.nosql.query(q_key, status="queued")
        result = []
        for rec in records[:max_batch_size]:
            self._storage.nosql.put(
                q_key, rec["key"], {**rec["value"], "status": "locked"})
            ev = _MockEventData(rec["value"].get("Body", ""))
            ev.system_properties = {"x-opt-sequence-number": 0}
            result.append(ev)
        return result

    def close(self) -> None:
        pass


def handle_eventhubs(
    sm: "StorageManager", parsed: dict, request: object,
    method: str, url: str, tpl: dict,
) -> Optional[FakeAzureResponse]:
    hub_name = parsed["container"]
    namespace = parsed.get("account", "default")
    q_key = f"eh_{namespace}_{hub_name}"

    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": "/" + hub_name + "/messages/" + msg_id}, b"")

    if method == "POST" and "/consumergroups" in url:
        resp = dict(tpl.get("consumergroup", {"name": "$Default"}))
        return FakeAzureResponse(200, _JSON_HDR, _enc(resp))

    if method == "GET" and "/consumergroups" in url:
        resp = dict(tpl.get("consumergroup", {"name": "$Default"}))
        return FakeAzureResponse(200, _JSON_HDR, _enc(resp))

    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)
        return FakeAzureResponse(204, {}, b"")

    if method == "GET":
        resp = dict(tpl.get("GET", {}))
        resp.setdefault("name", hub_name)
        return FakeAzureResponse(200, _JSON_HDR, _enc(resp))

    if method == "PUT":
        resp = dict(tpl.get("PUT", {}))
        resp.setdefault("name", hub_name)
        return FakeAzureResponse(200, _JSON_HDR, _enc(resp))

    if method == "DELETE":
        sm.nosql.delete(q_key, hub_name)
        return FakeAzureResponse(200, {}, b"")

    return None
