from __future__ import annotations

import re
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional

import azure.core.pipeline.transport._requests_basic as _azure_transport
from azure.servicebus import ServiceBusClient as _ServiceBusClient
from azure.storage.blob import BlobClient as _BlobClient
from azure.eventhub import EventHubProducerClient as _EventHubProducerClient
from azure.eventhub import EventHubConsumerClient as _EventHubConsumerClient
from azure.storage.queue import QueueClient as _QueueClient
from azure.data.tables import TableServiceClient as _TableServiceClient
from azure.data.tables import TableClient as _TableClient
from azure.appconfiguration import AzureAppConfigurationClient as _AppConfigClient

from mockmesh.core.fallback import FallbackConfig, safe_intercept
from mockmesh.core.logger import get_logger
from mockmesh.interceptors.azure.blob import handle_blob
from mockmesh.interceptors.azure.cosmos import handle_cosmos
from mockmesh.interceptors.azure.keyvault import handle_keyvault
from mockmesh.interceptors.azure.servicebus import handle_servicebus, _MockServiceBusClient
from mockmesh.interceptors.azure.eventhubs import handle_eventhubs, _MockEventHubProducerClient, _MockEventHubConsumerClient
from mockmesh.interceptors.azure.storage_queue import handle_storage_queue, _MockQueueClient
from mockmesh.interceptors.azure.table import handle_table, _MockTableServiceClient, _MockTableClient
from mockmesh.interceptors.azure.appconfig import handle_appconfig, _MockAzureAppConfigClient
from mockmesh.interceptors.azure.utils import _enc
from mockmesh.interceptors.common import cloud_provider_for_url
from mockmesh.interceptors.loader import load_json
from mockmesh.interceptors.responses import FakeAzureResponse
from mockmesh.plugins.base import BasePlugin

if TYPE_CHECKING:
    from mockmesh.response.engine import ResponseEngine
    from mockmesh.storage.manager import StorageManager

_log = get_logger("interceptors.azure")

_DEFAULTS_JSON = Path(__file__).parent.parent.parent / "defaults" / "azure.json"


def _parse_azure_url(url: str) -> dict:
    out: dict = {"svc": "other", "container": "", "path": "", "url": url}

    if "blob.core.windows.net" in url:
        m = re.match(r"https://[^/]+\.blob\.core\.windows\.net/([^/?]+)(?:/([^?]*))?", url)
        out["svc"] = "blob"
        if m:
            out["container"] = m.group(1) or ""
            out["path"]      = (m.group(2) or "").rstrip("?")
        return out

    if "cosmos.azure.com" in url or "documents.azure.com" in url or ("/dbs/" in url and "colls" in url):
        out["svc"] = "cosmos"
        m = re.search(r"/colls/([^/]+)(?:/docs(?:/([^/?]*))?)?", url)
        if m:
            out["container"] = m.group(1) or ""
            out["path"]      = m.group(2) or ""
        return out

    if "database.windows.net" in url:
        out["svc"] = "sql"; return out

    if "redis.cache.windows.net" in url:
        out["svc"] = "redis"; return out

    if "servicebus.windows.net" in url and "eventhub" not in url.lower():
        out["svc"] = "servicebus"
        m = re.match(r"https://[^/]+\.servicebus\.windows\.net/([^/?]+)", url)
        if m:
            out["container"] = m.group(1) or ""
        return out

    if "servicebus.windows.net" in url and "eventhub" in url.lower():
        out["svc"] = "eventhubs"
        m = re.match(r"https://[^/]+\.servicebus\.windows\.net/([^/?]+)", url)
        if m:
            out["container"] = m.group(1) or ""
        return out

    if "queue.core.windows.net" in url:
        out["svc"] = "storage_queue"
        m = re.match(r"https://[^/]+\.queue\.core\.windows\.net/([^/?]+)", url)
        if m:
            out["container"] = m.group(1) or ""
        return out

    if "table.core.windows.net" in url or "table.cosmos.azure.com" in url:
        out["svc"] = "table"
        m = re.match(r"https://[^/]+\.(?:table\.core\.windows\.net|table\.cosmos\.azure\.com)/([^?(/]+)", url)
        if m:
            out["container"] = m.group(1) or ""
        return out

    if "azconfig.io" in url or "appconfig" in url.lower():
        out["svc"] = "appconfig"
        m = re.search(r"/kv/([^?/]*)", url)
        if m:
            out["container"] = m.group(1) or ""
        return out

    if "vault.azure.net" in url:
        out["svc"] = "keyvault"
        m = re.search(r"/secrets/([^/?/]+)(?:/([^/?]*))?", url)
        if m:
            out["container"] = m.group(1) or ""
            out["path"]      = m.group(2) or ""
        return out

    if "management.azure.com" in url or "management.core.windows.net" in url:
        out["svc"] = "arm"
        if "/managedClusters" in url:   out["svc"] = "aks"
        elif "/containerGroups" in url: out["svc"] = "aci"
        elif "/containerApps" in url:   out["svc"] = "containerapp"
        return out

    return out


def _is_auth_url(url: str, defaults: dict) -> bool:
    auth_hosts = defaults.get("_config", {}).get("auth", {}).get("hosts", [])
    return any(h in url for h in auth_hosts)


def _mock_auth_response(url: str, method: str, defaults: dict) -> dict:
    auth = defaults.get("_config", {}).get("auth", {})
    if "openid-configuration" in url or ".well-known" in url:
        return auth.get("openid_config", {})
    if "discovery" in url and "keys" in url:
        return auth.get("jwks", {})
    if "instance" in url:
        return auth.get("imds", {})
    return auth.get("token", {})


def _default_azure_body(svc: str, method: str, defaults: dict) -> dict:
    svc_map = defaults.get(svc, defaults.get("arm", {}))
    return svc_map.get(method.upper(), svc_map.get("GET", {"mocked": True, "service": svc}))


class _MockBlobProperties:
    def __init__(self, size: int, etag: str, content_type: str = "application/octet-stream") -> None:
        self.size              = size
        self.etag              = etag
        self.content_type      = content_type
        self.content_encoding  = None
        self.content_language  = None
        self.content_md5       = None
        self.cache_control     = None
        self.content_disposition = None
        self.content_settings  = self


class _MockStorageStreamDownloader:
    def __init__(self, data: bytes, content_type: str = "application/octet-stream",
                 etag: str = '"mock-etag"') -> None:
        self._data      = data
        self._pos       = 0
        self.size       = len(data)
        self.properties = _MockBlobProperties(len(data), etag, content_type)

    def readall(self) -> bytes:
        return self._data

    def read(self, length: int = -1) -> bytes:
        if length < 0:
            result, self._pos = self._data[self._pos:], len(self._data)
        else:
            result = self._data[self._pos: self._pos + length]
            self._pos += len(result)
        return result

    def chunks(self):
        yield self._data

    def __iter__(self):
        yield self._data


class AzureInterceptor(BasePlugin):
    name = "azure"

    def __init__(self, response_engine: "ResponseEngine") -> None:
        super().__init__(response_engine)
        self._original_send: Any             = None
        self._original_download_blob: Any    = None
        self._original_sb_from_conn_str: Any = None
        self._transport_cls: Any             = None
        self._blob_client_cls: Any           = None
        self._sb_client_cls: Any             = None
        self._originals: dict[str, Any]      = {}
        self._defaults: dict = load_json(_DEFAULTS_JSON, "map")

    def activate(
        self,
        config: dict,
        *,
        storage: Optional["StorageManager"] = None,
        provider_defaults: Optional[dict] = None,
        fallback_config: Optional[FallbackConfig] = None,
        custom_handlers: Optional[dict] = None,
        **_: Any,
    ) -> None:
        self.storage = storage
        self._fallback_config = fallback_config
        self._custom_handlers = custom_handlers if custom_handlers is not None else {}
        if provider_defaults is not None:
            self._defaults = provider_defaults.get("azure", self._defaults)
        self._transport_cls    = _azure_transport.RequestsTransport
        self._blob_client_cls  = _BlobClient
        self._sb_client_cls    = _ServiceBusClient
        self._patch()
        self._active = True
        _log.info("Azure interceptor activated",
                  extra={"mm_patched": "RequestsTransport.send + BlobClient.download_blob"
                                       " + ServiceBusClient.from_connection_string"})

    def deactivate(self) -> None:
        if not self._active:
            return
        self._unpatch()
        self._active = False
        _log.info("Azure interceptor deactivated")

    def _handle_azure_request(self, request: Any, url: str, method: str) -> Any:
        if _is_auth_url(url, self._defaults):
            _log.info("Azure intercepted (auth mock)", extra={"mm_url": url[:80]})
            return FakeAzureResponse(200, {"Content-Type": "application/json"},
                                     _enc(_mock_auth_response(url, method, self._defaults)), method)
        if self.storage and self.storage._ready:
            result = self._dispatch_storage(request, url, method)
            if result is not None:
                _log.info("Azure intercepted (storage)", extra={"mm_url": url[:80]})
                return result
        parsed = _parse_azure_url(url)
        _log.info("Azure intercepted (config mock)",
                  extra={"mm_svc": parsed["svc"], "mm_method": method, "mm_url": url[:80]})
        return FakeAzureResponse(200, {"Content-Type": "application/json"},
                                 _enc(_default_azure_body(parsed["svc"], method, self._defaults)), method)

    def _patch(self) -> None:
        interceptor = self
        self._original_send = self._transport_cls.send

        def _mock_send(transport_self: Any, request: Any, **kwargs: Any) -> Any:
            url:    str = getattr(request, "url",    "") or ""
            method: str = (getattr(request, "method", "GET") or "GET").upper()

            _provider = cloud_provider_for_url(url)
            if _provider and _provider != "azure":
                _log.debug(
                    "Azure interceptor skipping non-Azure URL",
                    extra={"mm_url": url[:80], "mm_provider": _provider},
                )
                return interceptor._original_send(transport_self, request, **kwargs)

            return safe_intercept(
                provider="azure",
                operation=f"{method} {url[:80]}",
                handler_fn=interceptor._handle_azure_request,
                handler_args=(request, url, method),
                fallback_config=interceptor._fallback_config,
                original_fn=interceptor._original_send,
                original_args=(transport_self, request),
                original_kwargs=kwargs,
                generic_response=FakeAzureResponse(
                    200, {"Content-Type": "application/json"},
                    _enc({"mocked": True}), method,
                ),
                custom_handlers=interceptor._custom_handlers,
            )

        self._transport_cls.send = _mock_send  # type: ignore

        if self._blob_client_cls is not None:
            self._original_download_blob = self._blob_client_cls.download_blob

            def _mock_download_blob(blob_self: Any, offset: Any = None,
                                    length: Any = None, **kwargs: Any) -> Any:
                url    = getattr(blob_self, "url", "") or ""
                parsed = _parse_azure_url(url)
                if (interceptor.storage and interceptor.storage._ready
                        and parsed["svc"] == "blob" and parsed["path"]):
                    obj = interceptor.storage.az_blob_get(parsed["container"], parsed["path"])
                    if obj:
                        return _MockStorageStreamDownloader(
                            data=obj["Body"],
                            content_type=obj.get("ContentType", "application/octet-stream"),
                            etag=obj.get("ETag", '"mock-etag"'),
                        )
                return _MockStorageStreamDownloader(b"")

            self._blob_client_cls.download_blob = _mock_download_blob  # type: ignore

        if self._sb_client_cls is not None:
            self._original_sb_from_conn_str = self._sb_client_cls.from_connection_string
            _sm = self.storage

            @classmethod  # type: ignore[misc]
            def _mock_from_conn_str(cls, conn_str: str, **kwargs: Any):
                return _MockServiceBusClient(_sm)

            self._sb_client_cls.from_connection_string = _mock_from_conn_str  # type: ignore

        self._patch_new_sdk_clients()

    def _patch_new_sdk_clients(self) -> None:
        _sm = self.storage

        self._originals["eh_producer_from_conn"] = _EventHubProducerClient.from_connection_string
        self._originals["eh_consumer_from_conn"] = _EventHubConsumerClient.from_connection_string

        @classmethod  # type: ignore[misc]
        def _mock_eh_producer_from_conn(cls, conn_str: str, eventhub_name: str = "", **kw: Any):
            return _MockEventHubProducerClient(eventhub_name or "default", _sm)

        @classmethod  # type: ignore[misc]
        def _mock_eh_consumer_from_conn(cls, conn_str: str, consumer_group: str = "$Default",
                                         eventhub_name: str = "", **kw: Any):
            return _MockEventHubConsumerClient(eventhub_name or "default", consumer_group, _sm)

        _EventHubProducerClient.from_connection_string = _mock_eh_producer_from_conn  # type: ignore
        _EventHubConsumerClient.from_connection_string = _mock_eh_consumer_from_conn  # type: ignore

        self._originals["queue_from_conn"] = _QueueClient.from_connection_string

        @classmethod  # type: ignore[misc]
        def _mock_queue_from_conn(cls, conn_str: str, queue_name: str, **kw: Any):
            return _MockQueueClient(queue_name, _sm)

        _QueueClient.from_connection_string = _mock_queue_from_conn  # type: ignore

        self._originals["table_svc_from_conn"] = _TableServiceClient.from_connection_string
        self._originals["table_from_conn"] = _TableClient.from_connection_string

        @classmethod  # type: ignore[misc]
        def _mock_table_svc_from_conn(cls, conn_str: str, **kw: Any):
            return _MockTableServiceClient(_sm)

        @classmethod  # type: ignore[misc]
        def _mock_table_from_conn(cls, conn_str: str, table_name: str, **kw: Any):
            return _MockTableClient(table_name, _sm)

        _TableServiceClient.from_connection_string = _mock_table_svc_from_conn  # type: ignore
        _TableClient.from_connection_string = _mock_table_from_conn  # type: ignore

        self._originals["appconfig_from_conn"] = _AppConfigClient.from_connection_string

        @classmethod  # type: ignore[misc]
        def _mock_appconfig_from_conn(cls, conn_str: str, **kw: Any):
            return _MockAzureAppConfigClient(_sm)

        _AppConfigClient.from_connection_string = _mock_appconfig_from_conn  # type: ignore

    def _unpatch_new_sdk_clients(self) -> None:
        restore_map = {
            "eh_producer_from_conn": (_EventHubProducerClient, "from_connection_string"),
            "eh_consumer_from_conn": (_EventHubConsumerClient, "from_connection_string"),
            "queue_from_conn": (_QueueClient, "from_connection_string"),
            "table_svc_from_conn": (_TableServiceClient, "from_connection_string"),
            "table_from_conn": (_TableClient, "from_connection_string"),
            "appconfig_from_conn": (_AppConfigClient, "from_connection_string"),
        }
        for key, (cls, attr) in restore_map.items():
            original = self._originals.pop(key, None)
            if original is not None:
                setattr(cls, attr, original)

    def _unpatch(self) -> None:
        if self._transport_cls and self._original_send:
            self._transport_cls.send = self._original_send
            self._original_send = None
        if self._blob_client_cls and self._original_download_blob:
            self._blob_client_cls.download_blob = self._original_download_blob
            self._original_download_blob = None
        if self._sb_client_cls and self._original_sb_from_conn_str:
            self._sb_client_cls.from_connection_string = self._original_sb_from_conn_str
            self._original_sb_from_conn_str = None
        self._unpatch_new_sdk_clients()

    def _dispatch_storage(self, request: Any, url: str, method: str) -> Optional[FakeAzureResponse]:
        sm     = self.storage
        parsed = _parse_azure_url(url)
        tpl    = self._defaults.get(parsed["svc"], {})
        defaults = self._defaults
        handlers = {
            "blob":          handle_blob,
            "cosmos":        handle_cosmos,
            "servicebus":    lambda sm, p, req, m, u, t: handle_servicebus(sm, p, req, m, u, defaults),
            "keyvault":      handle_keyvault,
            "eventhubs":     lambda sm, p, req, m, u, t: handle_eventhubs(sm, p, req, m, u, defaults),
            "storage_queue": handle_storage_queue,
            "table":         handle_table,
            "appconfig":     handle_appconfig,
        }
        handler = handlers.get(parsed["svc"])
        return handler(sm, parsed, request, method, url, tpl) if handler else None
