"""
FoundryMatch SaaS - Streaming Feature Extractor
================================================
This module provides a service for real-time feature extraction from an
event stream (e.g., Kafka). It's designed to enrich events with ML features
that can be used for online inference, real-time monitoring, or populating
a feature store.
"""

import json
import logging
import time
from datetime import datetime
from typing import Any, Dict, List, Optional

from kafka import KafkaConsumer, KafkaProducer
from prometheus_client import Counter, Histogram
from rapidfuzz import fuzz

from fmatch.saas.security.privacy_enhanced import EnhancedPrivacyFilter

log = logging.getLogger(__name__)

# --- Prometheus Metrics for Monitoring ---
STREAM_EVENTS_PROCESSED_TOTAL = Counter(
    "fmatch_stream_events_processed_total",
    "Total number of events processed by the streaming feature extractor.",
    ["topic", "status"],  # Status can be 'success' or 'error'
)

STREAM_PROCESSING_LATENCY_SECONDS = Histogram(
    "fmatch_stream_processing_latency_seconds",
    "Latency of processing a single event in the stream.",
    ["topic"],
)


class StreamingFeatureExtractor:
    """
    Consumes events, extracts features, and produces enriched events.
    """

    def __init__(
        self,
        bootstrap_servers: List[str],
        consumer_topic: str,
        producer_topic: str,
        group_id: str = "feature_extractors",
        privacy_filter: Optional[EnhancedPrivacyFilter] = None,
    ):
        """
        Initializes the streaming feature extractor.

        Args:
            bootstrap_servers: List of Kafka broker addresses.
            consumer_topic: The topic to consume raw events from.
            producer_topic: The topic to produce enriched events to.
            group_id: The Kafka consumer group ID.
            privacy_filter: An optional EnhancedPrivacyFilter to sanitize data.
        """
        self.consumer_topic = consumer_topic
        self.producer_topic = producer_topic
        self.privacy_filter = privacy_filter

        try:
            self.consumer = KafkaConsumer(
                consumer_topic,
                bootstrap_servers=bootstrap_servers,
                group_id=group_id,
                value_deserializer=lambda v: json.loads(v.decode("utf-8")),
                auto_offset_reset="latest",
            )
            self.producer = KafkaProducer(
                bootstrap_servers=bootstrap_servers,
                value_serializer=lambda v: json.dumps(v).encode("utf-8"),
            )
            log.info(
                f"StreamingFeatureExtractor initialized for topic '{consumer_topic}'"
            )
        except Exception as e:
            log.error(
                f"Failed to connect to Kafka at {bootstrap_servers}: {e}", exc_info=True
            )
            raise

        # Re-use feature functions from the training pipeline for consistency
        self.feature_functions = {
            "wratio": fuzz.WRatio,
            "partial_ratio": fuzz.partial_ratio,
            "token_set_ratio": fuzz.token_set_ratio,
            "jaro_winkler": fuzz.jaro_winkler_similarity,
        }

    def _extract_features(self, event: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Extracts features from a single event. This example assumes an event
        for a candidate match pair.
        """
        event_type = event.get("event_type")
        if event_type != "match_candidate_pair":
            return None  # Only process specific events

        source_record = event.get("source_record")
        ref_record = event.get("reference_record")
        mappings = event.get("mappings")

        if not all([source_record, ref_record, mappings]):
            log.warning(
                f"Skipping event {event.get('event_id')}: missing required fields."
            )
            return None

        features = {}
        for mapping in mappings:
            s_col, r_col = mapping.get("source"), mapping.get("ref")
            if not s_col or not r_col:
                continue

            s_val = str(source_record.get(s_col, "")).strip()
            r_val = str(ref_record.get(r_col, "")).strip()

            for func_name, func in self.feature_functions.items():
                feature_name = f"feat_{s_col}_{func_name}"
                score = func(s_val, r_val) if s_val and r_val else 0.0
                features[feature_name] = round(score / 100.0, 4)  # Normalize to 0-1

        # Add extracted features to the original event
        event["extracted_features"] = features
        event["features_extracted_at"] = datetime.utcnow().isoformat()
        return event

    def run(self):
        """Starts the main processing loop."""
        log.info(f"Starting to consume from topic '{self.consumer_topic}'...")
        for message in self.consumer:
            start_time = time.perf_counter()
            try:
                event = message.value

                # Apply privacy filter if configured
                if self.privacy_filter:
                    event = self.privacy_filter.filter_payload(event)

                enriched_event = self._extract_features(event)

                if enriched_event:
                    self.producer.send(self.producer_topic, enriched_event)
                    STREAM_EVENTS_PROCESSED_TOTAL.labels(
                        topic=self.consumer_topic, status="success"
                    ).inc()

            except json.JSONDecodeError:
                log.error(f"Failed to decode JSON message: {message.value[:100]}")
                STREAM_EVENTS_PROCESSED_TOTAL.labels(
                    topic=self.consumer_topic, status="error"
                ).inc()
            except Exception as e:
                log.error(f"Error processing event: {e}", exc_info=True)
                STREAM_EVENTS_PROCESSED_TOTAL.labels(
                    topic=self.consumer_topic, status="error"
                ).inc()
            finally:
                latency = time.perf_counter() - start_time
                STREAM_PROCESSING_LATENCY_SECONDS.labels(
                    topic=self.consumer_topic
                ).observe(latency)

    def close(self):
        """Closes Kafka connections gracefully."""
        log.info("Closing Kafka connections...")
        self.consumer.close()
        self.producer.flush()
        self.producer.close()
