import asyncio
import json
import logging
from typing import List, Tuple, Dict, Any, Optional

try:
    # Type stubs might be missing for aiokafka so we ignore import errors in mypy
    from aiokafka import AIOKafkaConsumer, AIOKafkaProducer, TopicPartition # type: ignore
    import aiokafka.errors as kafka_errors # type: ignore
    KAFKA_AVAILABLE = True
except ImportError:
    KAFKA_AVAILABLE = False
    
from pspf.connectors.base import StreamingBackend
from pspf.state.store import StateStore
from pspf.utils.logging import get_logger

logger = get_logger("KafkaBackend")

class KafkaStreamBackend(StreamingBackend):
    """
    Kafka (or Redpanda) implementation of StreamingBackend.
    Requires `aiokafka` package.
    """
    def __init__(self, bootstrap_servers: str, topic: str, group_id: str, client_id: str, state_store: Optional[StateStore] = None) -> None:
        if not KAFKA_AVAILABLE:
            raise ImportError("aiokafka is required for KafkaStreamBackend. Install it with `pip install aiokafka`.")
            
        self.bootstrap_servers = bootstrap_servers
        self.topic = topic
        self.group_id = group_id
        self.client_id = client_id
        self.state_store = state_store
        self.retry_tracker_prefix = f"pspf:retries:{group_id}:{topic}:"
        
        self.consumer: Optional[AIOKafkaConsumer] = None
        self.producer: Optional[AIOKafkaProducer] = None
        self._connected = False

    @property
    def stream_key(self) -> str:
        return self.topic

    @property
    def group_name(self) -> str:
        return self.group_id

    def clone_with_topic(self, topic: str) -> "KafkaStreamBackend":
        """Create a new backend instance for a different topic, sharing the underlying connection."""
        return KafkaStreamBackend(
            bootstrap_servers=self.bootstrap_servers,
            topic=topic,
            group_id=self.group_id,
            client_id=self.client_id,
            state_store=self.state_store
        )

    async def connect(self) -> None:
        try:
            self.producer = AIOKafkaProducer(
                bootstrap_servers=self.bootstrap_servers,
                client_id=f"{self.client_id}-producer"
            )
            await self.producer.start()
            
            # Consumer will be started on demand or here? 
            # Best to start here to ensure connectivity
            self.consumer = AIOKafkaConsumer(
                self.topic,
                bootstrap_servers=self.bootstrap_servers,
                group_id=self.group_id,
                client_id=f"{self.client_id}-consumer",
                enable_auto_commit=False, # We Manually ACK
                auto_offset_reset="earliest"
            )
            await self.consumer.start()
            
            self._connected = True
            logger.info(f"Connected to Kafka at {self.bootstrap_servers}")
        except Exception as e:
            logger.error(f"Failed to connect to Kafka: {e}")
            raise

    async def close(self) -> None:
        if self.producer:
            await self.producer.stop()
        if self.consumer:
            await self.consumer.stop()
        self._connected = False
        logger.info("Closed Kafka connection")

    async def ping(self) -> bool:
        if not self._connected:
            raise ConnectionError("Not connected")
        # Lightweight check? Maybe metadata request?
        # aiokafka doesn't have a direct ping, but we can check assignment or partitions
        if self.consumer:
            await self.consumer.partitions_for_topic(self.topic)
        return True

    async def ensure_group_exists(self, start_id: str = "0") -> None:
        # Kafka handles group creation automatically on join
        pass

    async def add_event(self, data: Dict[str, Any], max_len: Optional[int] = None) -> str:
        if not self.producer:
            raise ConnectionError("Producer not connected")
            
        # Serialize
        payload = json.dumps(data).encode("utf-8")
        
        # Determine partition logic here? Or let producer duplicate.
        # Ideally we use event_id as key for ordering.
        key = data.get("event_id", "").encode("utf-8") if "event_id" in data else None
        
        try:
            record_metadata = await self.producer.send_and_wait(self.topic, payload, key=key)
            # Return "partition-offset" as ID
            return f"{record_metadata.partition}-{record_metadata.offset}"
        except Exception as e:
            logger.error(f"Failed to produce to Kafka: {e}")
            raise

    async def read_batch(self, count: int = 10, block_ms: int = 1000) -> List[Tuple[str, Dict[str, Any]]]:
        if not self.consumer:
            raise ConnectionError("Consumer not connected")
            
        try:
            # Poll for messages
            # timeout_ms in getmany
            result = await self.consumer.getmany(timeout_ms=block_ms, max_records=count)
            
            messages = []
            for tp, records in result.items():
                for record in records:
                    msg_id = f"{record.partition}-{record.offset}"
                    try:
                        data = json.loads(record.value.decode("utf-8"))
                        messages.append((msg_id, data))
                    except json.JSONDecodeError:
                        logger.warning(f"Skipping non-JSON message {msg_id}")
                        continue
                        
            return messages
        except Exception as e:
            logger.error(f"Error reading batch from Kafka: {e}")
            raise

    async def ack_batch(self, message_ids: List[str]) -> None:
        if not self.consumer:
            raise ConnectionError("Consumer not connected")
            
        # Kafka commits offsets, not individual IDs.
        # We need to find the highest offset per partition in this batch and commit that + 1.
        # Simple approach: We assume batch processing is successful.
        # AIOKafka commit() commits the *current* position? No, commit expects offsets.
        # This is tricky without tracking TopicPartition objects.
        # Limitation: We need to parse msg_id to get partition/offset back.
        
        offsets_to_commit: Dict[Any, Any] = {}
        
        for msg_id in message_ids:
            try:
                part_str, off_str = msg_id.split("-")
                partition = int(part_str)
                offset = int(off_str)
                
                # We reconstruct TopicPartition
                from aiokafka import TopicPartition # type: ignore
                tp = TopicPartition(self.topic, partition)
                
                # We want to commit offset + 1
                new_offset = offset + 1
                
                if tp not in offsets_to_commit or new_offset > offsets_to_commit[tp].offset:
                     # Struct to hold offset
                     # commit expects {TopicPartition: OffsetAndMetadata(offset, metadata)}
                     # or simple {TopicPartition: offset}
                     offsets_to_commit[tp] = new_offset
            except ValueError:
                logger.warning(f"Invalid message ID format for Kafka ACK: {msg_id}")
        
        if offsets_to_commit:
            await self.consumer.commit(offsets_to_commit) # type: ignore
            
            # Durable Retry Cleanup
            if self.state_store:
                for msg_id in message_ids:
                    retry_key = f"{self.retry_tracker_prefix}{msg_id}"
                    await self.state_store.delete(retry_key)

    async def claim_stuck_messages(self, min_idle_time_ms: int = 60000, count: int = 10) -> List[Tuple[str, Dict[str, Any]]]:
        # Kafka doesn't have PEL claiming mechanim like Redis Streams.
        # It relies on rebalancing if a consumer dies.
        # "Stuck" messages inside a live consumer are app logic.
        return []

    async def increment_retry_count(self, message_id: str) -> int:
        """
        Increment and return the retry count for a message.
        If a StateStore is available, it persists the count.
        """
        if self.state_store:
            retry_key = f"{self.retry_tracker_prefix}{message_id}"
            current_count = await self.state_store.get(retry_key, 0)
            new_count = current_count + 1
            await self.state_store.put(retry_key, new_count)
            return new_count
            
        if not hasattr(self, "_retry_counts"):
            self._retry_counts: Dict[str, int] = {}
        
        self._retry_counts[message_id] = self._retry_counts.get(message_id, 0) + 1
        return self._retry_counts[message_id]

    async def move_to_dlq(self, message_id: str, data: Dict[str, Any], error: str) -> None:
         # Produce to DLQ topic
         dlq_topic = f"{self.topic}-dlq"
         if self.producer:
             data["_error"] = error
             data["_original_id"] = message_id
             payload = json.dumps(data).encode("utf-8")
             await self.producer.send_and_wait(dlq_topic, payload)

             # Clear retry state
             if self.state_store:
                 retry_key = f"{self.retry_tracker_prefix}{message_id}"
                 await self.state_store.delete(retry_key)

    async def get_pending_info(self) -> Dict[str, Any]:
        """
        Retrieves lag info from Kafka by comparing committed offsets to end offsets.
        """
        if not self.consumer:
            return {"lag": 0, "pending": 0, "status": "disconnected"}
            
        try:
            partitions = self.consumer.assignment()
            if not partitions:
                # If no assignment yet, we can't calculate lag
                return {"lag": 0, "pending": 0, "status": "no_assignment"}
                
            # Get committed offsets
            committed_tasks = [self.consumer.committed(tp) for tp in partitions]
            committed_results = await asyncio.gather(*committed_tasks)
            committed_map = dict(zip(partitions, committed_results))
            
            # Get end offsets
            end_offsets = await self.consumer.end_offsets(partitions)
            
            total_lag = 0
            for tp in partitions:
                committed_offset = committed_map.get(tp)
                end_offset = end_offsets.get(tp, 0)
                
                if committed_offset is not None:
                    total_lag += max(0, end_offset - committed_offset)
                else:
                    # If never committed, we assume lag is the full partition if we don't know start
                    total_lag += end_offset
                    
            return {
                "lag": total_lag,
                "pending": total_lag,
                "partition_count": len(partitions),
                "status": "connected"
            }
        except Exception as e:
            logger.warning(f"Failed to calculate Kafka lag: {e}")
            return {"lag": 0, "pending": 0, "error": str(e)}

