"""
Asynchronous Avro consumer using confluent-kafka with asyncio.

Features:
- Non-blocking poll with async/await
- Suitable for FastAPI websocket streaming
- Background message fetching with queue

Usage:
    from investify_utils.kafka import AsyncAvroConsumer

    consumer = AsyncAvroConsumer(
        topic="my-topic",
        subject="my-topic-value",
        schema_registry_url="http://localhost:8081",
        bootstrap_servers="localhost:9092",
        group_id="my-consumer-group",
    )

    # In FastAPI websocket
    @app.websocket("/ws")
    async def websocket_endpoint(websocket: WebSocket):
        await websocket.accept()
        async for key, value in consumer:
            await websocket.send_json(value)
"""

import asyncio
import logging
from typing import AsyncIterator

from confluent_kafka import DeserializingConsumer, TopicPartition
from confluent_kafka.schema_registry import SchemaRegistryClient
from confluent_kafka.schema_registry.avro import AvroDeserializer
from confluent_kafka.serialization import StringDeserializer

logger = logging.getLogger(__name__)


class AsyncAvroConsumer:
    """
    Asynchronous Avro consumer for async frameworks.

    Uses a background thread for polling (confluent-kafka is not async-native)
    and an asyncio queue to bridge to async code.

    Args:
        topic: Kafka topic name or list of topics
        subject: Schema Registry subject name
        schema_registry_url: Schema Registry URL
        bootstrap_servers: Kafka bootstrap servers
        group_id: Consumer group ID
        seek_to_end: Start from latest offset (default: False)
        queue_size: Max messages to buffer (default: 100)
        **kwargs: Additional Kafka consumer config
    """

    def __init__(
        self,
        topic: str | list[str],
        subject: str,
        schema_registry_url: str,
        bootstrap_servers: str,
        group_id: str,
        seek_to_end: bool = False,
        queue_size: int = 100,
        **kwargs,
    ):
        self._schema_registry_url = schema_registry_url
        self._bootstrap_servers = bootstrap_servers
        self._subject = subject
        self._group_id = group_id
        self._topic = topic
        self._seek_to_end = seek_to_end
        self._queue_size = queue_size
        self._kwargs = kwargs

        self._consumer: DeserializingConsumer | None = None
        self._queue: asyncio.Queue | None = None
        self._poll_task: asyncio.Task | None = None
        self._running = False

    def _init_consumer(self) -> DeserializingConsumer:
        """Initialize the consumer."""
        schema_registry_client = SchemaRegistryClient({"url": self._schema_registry_url})
        registered_schema = schema_registry_client.get_latest_version(self._subject)
        schema_str = registered_schema.schema.schema_str

        avro_deserializer = AvroDeserializer(schema_registry_client, schema_str)

        consumer_config = {
            "bootstrap.servers": self._bootstrap_servers,
            "group.id": self._group_id,
            "key.deserializer": StringDeserializer("utf_8"),
            "value.deserializer": avro_deserializer,
            **self._kwargs,
        }
        consumer = DeserializingConsumer(consumer_config)

        topic_list = self._topic if isinstance(self._topic, list) else [self._topic]

        if self._seek_to_end:

            def seek_to_end_assign(c, partitions):
                for p in partitions:
                    high_offset = c.get_watermark_offsets(p)[1]
                    p.offset = high_offset
                c.assign(partitions)

            consumer.subscribe(topic_list, on_assign=seek_to_end_assign)
        else:
            consumer.subscribe(topic_list)

        return consumer

    async def start(self) -> None:
        """Start consuming messages in background."""
        if self._running:
            return

        self._consumer = self._init_consumer()
        self._queue = asyncio.Queue(maxsize=self._queue_size)
        self._running = True
        self._poll_task = asyncio.create_task(self._poll_loop())

    async def _poll_loop(self) -> None:
        """Background task for polling messages."""
        loop = asyncio.get_running_loop()

        while self._running:
            # Run blocking poll in thread executor
            msg = await loop.run_in_executor(None, self._consumer.poll, 0.1)

            if msg is None:
                continue

            if msg.error():
                logger.error(f"Consumer error: {msg.error()}")
                continue

            # Put message in queue (blocks if full)
            await self._queue.put((msg.key(), msg.value(), msg))

    async def poll(self, timeout: float = 1.0) -> tuple[str | None, dict | None, object] | None:
        """
        Poll for a single message.

        Args:
            timeout: Maximum time to wait in seconds

        Returns:
            Tuple of (key, value, raw_msg) or None if timeout
        """
        if not self._running:
            await self.start()

        try:
            return await asyncio.wait_for(self._queue.get(), timeout=timeout)
        except asyncio.TimeoutError:
            return None

    async def __aiter__(self) -> AsyncIterator[tuple[str | None, dict | None]]:
        """
        Async iterator for consuming messages.

        Yields:
            Tuple of (key, value) for each message
        """
        if not self._running:
            await self.start()

        while self._running:
            try:
                key, value, _ = await asyncio.wait_for(self._queue.get(), timeout=1.0)
                yield key, value
            except asyncio.TimeoutError:
                continue

    def commit(self) -> None:
        """Commit current offsets."""
        if self._consumer:
            self._consumer.commit()

    def commit_offsets(self, offsets: list[TopicPartition]) -> None:
        """Commit specific offsets."""
        if self._consumer:
            self._consumer.commit(offsets=offsets)

    async def close(self) -> None:
        """Stop consuming and close the consumer."""
        self._running = False

        if self._poll_task:
            self._poll_task.cancel()
            try:
                await self._poll_task
            except asyncio.CancelledError:
                pass
            self._poll_task = None

        if self._consumer:
            self._consumer.close()
            self._consumer = None

    async def __aenter__(self):
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()
