# -*- coding: utf-8 -*-
"""
Kafka 适配器
基于 aiokafka 异步客户端
支持：生产者/消费者、消费组、分区、偏移量管理
"""

import asyncio
import json
import logging
from typing import Optional

from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
from aiokafka.errors import KafkaError

from uas.comm.base import CommBaseAdapter, ConnectionConfig, MessageHandler
from uas.comm.message import Message

logger = logging.getLogger(__name__)


class KafkaAdapterComm(CommBaseAdapter):
    """
    Kafka 通信适配器
    URL 格式: kafka://[user:pass@]host[:port][/group_id][?options]
    示例:     kafka://broker:9092/my-group?auto_offset_reset=earliest&acks=all
    """

    def __init__(self, config: ConnectionConfig):
        super().__init__(config)
        self._producer: Optional[AIOKafkaProducer] = None
        self._consumers: dict[str, AIOKafkaConsumer] = {}
        self._consume_tasks: dict[str, asyncio.Task] = {}

    @property
    def _bootstrap_servers(self) -> str:
        return f"{self.config.host}:{self.config.port}"

    @property
    def _group_id(self) -> str:
        return self.config.database or "default-group"

    # ── 生命周期 ──────────────────────────────────────────────────

    async def connect(self) -> None:
        opts = self.config.options
        self._producer = AIOKafkaProducer(
            bootstrap_servers = self._bootstrap_servers,
            acks              = opts.get("acks", "all"),
            compression_type  = opts.get("compression_type", None),
            value_serializer  = lambda v: json.dumps(v, ensure_ascii=False).encode(),
        )
        await self._producer.start()
        self._running = True
        self.logger.info(f"✅ Kafka Producer 已连接: {self._bootstrap_servers}")

    async def disconnect(self) -> None:
        self._running = False
        # 停止所有消费任务
        for topic, task in self._consume_tasks.items():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
        # 停止所有消费者
        for consumer in self._consumers.values():
            await consumer.stop()
        # 停止生产者
        if self._producer:
            await self._producer.stop()
        self.logger.info("🔌 Kafka 连接已关闭")

    # ── 核心功能 ──────────────────────────────────────────────────

    async def publish(self, message: Message) -> bool:
        if not self._producer:
            raise RuntimeError("Kafka 未连接")
        try:
            await self._producer.send_and_wait(
                message.topic,
                value   = message.to_dict(),
                headers = [(k, v.encode()) for k, v in message.headers.items()],
            )
            self.logger.debug(f"📤 Kafka 发布 topic={message.topic!r}")
            return True
        except KafkaError as e:
            self.logger.error(f"Kafka 发布失败: {e}")
            return False

    async def subscribe(self, topic: str, handler: MessageHandler) -> None:
        opts = self.config.options
        self._register_handler(topic, handler)

        if topic not in self._consumers:
            consumer = AIOKafkaConsumer(
                topic,
                bootstrap_servers    = self._bootstrap_servers,
                group_id             = self._group_id,
                auto_offset_reset    = opts.get("auto_offset_reset", "latest"),
                enable_auto_commit   = opts.get("enable_auto_commit", "true") == "true",
                value_deserializer   = lambda v: json.loads(v.decode()),
            )
            await consumer.start()
            self._consumers[topic] = consumer
            self._consume_tasks[topic] = asyncio.create_task(
                self._consume_loop(topic, consumer)
            )
        self.logger.info(f"📥 订阅 Kafka topic: {topic!r} group={self._group_id}")

    async def unsubscribe(self, topic: str) -> None:
        if topic in self._consume_tasks:
            self._consume_tasks[topic].cancel()
            del self._consume_tasks[topic]
        if topic in self._consumers:
            await self._consumers[topic].stop()
            del self._consumers[topic]
        self._handlers.pop(topic, None)
        self.logger.info(f"🚫 取消订阅 Kafka topic: {topic!r}")

    # ── 内部消费循环 ──────────────────────────────────────────────

    async def _consume_loop(self, topic: str, consumer: AIOKafkaConsumer):
        self.logger.info(f"🔄 Kafka 消费循环启动: {topic!r}")
        try:
            async for record in consumer:
                try:
                    message = Message.from_dict(record.value)
                    await self._dispatch(message)
                except Exception as e:
                    self.logger.error(f"消息处理异常: {e}", exc_info=True)
        except asyncio.CancelledError:
            self.logger.info(f"🛑 Kafka 消费循环停止: {topic!r}")

    async def health_check(self) -> bool:
        try:
            return self._producer is not None and self._running
        except Exception:
            return False