"""
Kafka sink for extractor.

Publishes each record as a JSON message to a Kafka topic.
Useful for streaming pipelines where downstream consumers (Flink, Spark, etc.)
process extractor output in real time.

Auth (env vars):
    EXTRACTOR_KAFKA_BROKERS          (comma-separated, default: localhost:9092)
    EXTRACTOR_KAFKA_SECURITY_PROTOCOL (PLAINTEXT, SASL_SSL, etc.)
    EXTRACTOR_KAFKA_SASL_MECHANISM   (PLAIN, SCRAM-SHA-256, etc.)
    EXTRACTOR_KAFKA_SASL_USERNAME
    EXTRACTOR_KAFKA_SASL_PASSWORD

Install: pip install 'coresdk-extractor[kafka]'
"""

from __future__ import annotations

import json
import logging
import os
from typing import Any

from .base import BaseSink
from .registry import register_sink

logger = logging.getLogger(__name__)


class KafkaSink(BaseSink):
    def __init__(
        self,
        brokers: str | None = None,
        topic: str = "extractor_elements",
        batch_size: int = 100,
    ) -> None:
        self.brokers = brokers or os.environ.get("EXTRACTOR_KAFKA_BROKERS", "localhost:9092")
        self.topic = topic
        self.batch_size = batch_size
        self._producer = None
        self._buffer: list[dict[str, Any]] = []
        self._total_inserted = 0

    def connect(self) -> None:
        try:
            from confluent_kafka import Producer
        except ImportError as e:
            raise ImportError(
                "confluent-kafka not installed. Run: pip install 'coresdk-extractor[kafka]'"
            ) from e
        conf: dict[str, str] = {"bootstrap.servers": self.brokers}
        proto = os.environ.get("EXTRACTOR_KAFKA_SECURITY_PROTOCOL")
        if proto:
            conf["security.protocol"] = proto
        mechanism = os.environ.get("EXTRACTOR_KAFKA_SASL_MECHANISM")
        if mechanism:
            conf["sasl.mechanism"] = mechanism
        user = os.environ.get("EXTRACTOR_KAFKA_SASL_USERNAME")
        if user:
            conf["sasl.username"] = user
        pwd = os.environ.get("EXTRACTOR_KAFKA_SASL_PASSWORD")
        if pwd:
            conf["sasl.password"] = pwd
        self._producer = Producer(conf)
        logger.info("KafkaSink connected: %s topic=%s", self.brokers, self.topic)

    def feed(self, record: dict[str, Any]) -> None:
        self._buffer.append(record)
        if len(self._buffer) >= self.batch_size:
            self.write_batch(self._buffer)
            self._buffer = []

    def write_batch(self, records: list[dict[str, Any]]) -> None:
        if not records or self._producer is None:
            return
        for r in records:
            key = r.get("source_sha256", "")
            value = json.dumps(r, ensure_ascii=False).encode()
            self._producer.produce(self.topic, key=key.encode(), value=value)
        self._producer.flush()
        self._total_inserted += len(records)

    def flush(self) -> None:
        if self._buffer:
            self.write_batch(self._buffer)
            self._buffer = []
        if self._producer:
            self._producer.flush()

    def close(self) -> None:
        if self._producer is not None:
            self._producer.flush()
            if hasattr(self._producer, "close"):
                self._producer.close()
        self._producer = None

    @property
    def total_inserted(self) -> int:
        return self._total_inserted


register_sink("kafka", KafkaSink)
