from __future__ import annotations

import importlib
import importlib
from typing import TYPE_CHECKING, Any, Callable, Optional

import confluent_kafka
import kafka as kafka_python

from mockmesh.core.fallback import FallbackConfig
from mockmesh.core.logger import get_logger
from mockmesh.plugins.base import BasePlugin

if TYPE_CHECKING:
    from mockmesh.response.engine import ResponseEngine

_log = get_logger("interceptors.kafka")


_active_kafka_rules: list = []


def _match_kafka_rule(topic: str) -> dict:
    for rule in _active_kafka_rules:
        t = rule.get("match", {}).get("topic", "*")
        if t in ("*", topic):
            return rule
    return {}


class _MockConfluentMessage:

    def __init__(self, topic: str, value: Any, offset: int = 0, partition: int = 0) -> None:
        self._topic = topic
        self._value = value if isinstance(value, bytes) else str(value).encode()
        self._offset = offset
        self._partition = partition
        self._error = None
        self._key = b"mock-key"

    def topic(self) -> str:
        return self._topic

    def value(self) -> bytes:
        return self._value

    def key(self) -> bytes:
        return self._key

    def offset(self) -> int:
        return self._offset

    def partition(self) -> int:
        return self._partition

    def error(self) -> None:
        return None


class _MockConfluentProducer:
    def __init__(self, config: dict, response_engine: "ResponseEngine", **kwargs: Any) -> None:
        self._re = response_engine
        self._config = config

    def produce(
        self,
        topic: str,
        value: Any = None,
        key: Any = None,
        callback: Optional[Callable] = None,
        **kwargs: Any,
    ) -> None:
        rule = _match_kafka_rule(topic)
        pr = rule.get("response", {}).get("produce") or rule.get("produce_response", {"error": None, "offset": 0, "partition": 0})
        _log.info(
            "Intercepted Kafka produce (confluent)",
            extra={"mm_topic": topic, "mm_offset": pr.get("offset")},
        )
        if callback:
            msg = _MockConfluentMessage(topic, value or b"", pr.get("offset", 0), pr.get("partition", 0))
            callback(pr.get("error"), msg)

    def flush(self, timeout: float = -1) -> int:
        return 0

    def poll(self, timeout: float = 0) -> int:
        return 0


class _MockConfluentConsumer:
    def __init__(self, config: dict, response_engine: "ResponseEngine", **kwargs: Any) -> None:
        self._re = response_engine
        self._topics: list[str] = []
        self._running = True

    def subscribe(self, topics: list[str], **kwargs: Any) -> None:
        self._topics = list(topics)

    def poll(self, timeout: float = 1.0) -> Optional[_MockConfluentMessage]:
        if not self._topics:
            return None
        topic = self._topics[0]
        rule = _match_kafka_rule(topic)
        cr = rule.get("response", {}).get("consume") or rule.get("consume_response", {})
        if not cr:
            return None
        _log.info("Intercepted Kafka poll (confluent)", extra={"mm_topic": topic})
        return _MockConfluentMessage(
            topic=cr.get("topic", topic),
            value=cr.get("value", "{}").encode(),
            offset=cr.get("offset", 0),
            partition=cr.get("partition", 0),
        )

    def commit(self, **kwargs: Any) -> None:
        pass

    def close(self) -> None:
        self._running = False


class _MockKafkaProducer:
    def __init__(self, response_engine: "ResponseEngine", **kwargs: Any) -> None:
        self._re = response_engine

    class _FutureResult:
        def __init__(self, topic: str, offset: int = 0, partition: int = 0):
            self.topic = topic
            self.offset = offset
            self.partition = partition

    class _FakeFuture:
        def __init__(self, result: Any) -> None:
            self._result = result

        def get(self, timeout: Optional[float] = None) -> Any:
            return self._result

    def send(self, topic: str, value: Any = None, key: Any = None, **kwargs: Any) -> Any:
        rule = _match_kafka_rule(topic)
        pr = rule.get("response", {}).get("produce") or rule.get("produce_response", {"offset": 0, "partition": 0})
        _log.info(
            "Intercepted Kafka send (kafka-python)",
            extra={"mm_topic": topic},
        )
        result = self._FutureResult(topic, pr.get("offset", 0), pr.get("partition", 0))
        return self._FakeFuture(result)

    def flush(self, timeout: Optional[float] = None) -> None:
        pass

    def close(self, timeout: Optional[float] = None) -> None:
        pass


class _MockKafkaConsumer:
    def __init__(self, *topics: str, response_engine: "ResponseEngine", **kwargs: Any) -> None:
        self._re = response_engine
        self._topics = list(topics)

    def __iter__(self):
        for topic in self._topics:
            rule = _match_kafka_rule(topic)
            cr = rule.get("response", {}).get("consume") or rule.get("consume_response", {})
            if cr:
                yield _MockKafkaRecord(
                    topic=cr.get("topic", topic),
                    value=cr.get("value", "{}"),
                    offset=cr.get("offset", 0),
                    partition=cr.get("partition", 0),
                )

    def close(self) -> None:
        pass


class _MockKafkaRecord:
    def __init__(self, topic: str, value: Any, offset: int, partition: int) -> None:
        self.topic = topic
        self.value = value.encode() if isinstance(value, str) else value
        self.offset = offset
        self.partition = partition
        self.key = b"mock-key"


class KafkaInterceptor(BasePlugin):

    name = "kafka"

    def __init__(self, response_engine: "ResponseEngine") -> None:
        super().__init__(response_engine)
        self._originals: dict[str, Any] = {}

    def activate(
        self,
        config: dict,
        *,
        storage: "Any | None" = None,
        provider_defaults: Optional[dict] = None,
        fallback_config: Optional[FallbackConfig] = None,
        custom_handlers: Optional[dict] = None,
        **_: Any,
    ) -> None:
        self._fallback_config = fallback_config
        self._custom_handlers = custom_handlers if custom_handlers is not None else {}
        global _active_kafka_rules
        if provider_defaults is not None:
            kafka_rules = provider_defaults.get("kafka")
            if kafka_rules is not None:
                _active_kafka_rules = kafka_rules
        self._patch_confluent()
        self._patch_kafka_python()
        self._active = True
        _log.info("Kafka interceptor activated")

    def deactivate(self) -> None:
        if not self._active:
            return
        for key, original in self._originals.items():
            module_path, attr = key.rsplit(".", 1)
            try:
                mod = importlib.import_module(module_path)
                setattr(mod, attr, original)
            except Exception:
                pass
        self._originals.clear()
        self._active = False
        _log.info("Kafka interceptor deactivated")

    def _patch_confluent(self) -> None:
        re = self.response_engine

        class _Producer(_MockConfluentProducer):
            def __init__(self, conf: dict, **kwargs: Any):
                super().__init__(conf, re, **kwargs)

        class _Consumer(_MockConfluentConsumer):
            def __init__(self, conf: dict, **kwargs: Any):
                super().__init__(conf, re, **kwargs)

        self._originals["confluent_kafka.Producer"] = confluent_kafka.Producer
        self._originals["confluent_kafka.Consumer"] = confluent_kafka.Consumer
        confluent_kafka.Producer = _Producer  # type: ignore
        confluent_kafka.Consumer = _Consumer  # type: ignore
        _log.debug("Patched confluent_kafka.Producer / Consumer")

    def _patch_kafka_python(self) -> None:
        re = self.response_engine

        class _KProducer(_MockKafkaProducer):
            def __init__(self, **kwargs: Any):
                super().__init__(re, **kwargs)

        class _KConsumer(_MockKafkaConsumer):
            def __init__(self, *topics: str, **kwargs: Any):
                super().__init__(*topics, response_engine=re, **kwargs)

        self._originals["kafka.KafkaProducer"] = kafka_python.KafkaProducer
        self._originals["kafka.KafkaConsumer"] = kafka_python.KafkaConsumer
        kafka_python.KafkaProducer = _KProducer  # type: ignore
        kafka_python.KafkaConsumer = _KConsumer  # type: ignore
        _log.debug("Patched kafka.KafkaProducer / KafkaConsumer")
