import threading
from collections import defaultdict
from typing import Any, Dict, List, Optional


def escape_label(value: str) -> str:
    return value.replace(chr(92), chr(92)*2).replace('"', chr(92)+'"').replace(chr(10), chr(92)+'n')


class SimpleHistogram:
    """
    Lightweight, thread-safe Prometheus-compatible histogram.
    """

    def __init__(self, name: str, help_text: str, buckets: List[float]):
        self.name = name
        self.help_text = help_text
        self.buckets = sorted(buckets)
        self._counts: Dict[float, int] = {b: 0 for b in self.buckets}
        self._count = 0
        self._sum = 0.0
        self._lock = threading.Lock()

    def observe(self, value: float) -> None:
        with self._lock:
            self._count += 1
            self._sum += value
            for b in self.buckets:
                if value <= b:
                    self._counts[b] += 1

    def to_prometheus(self) -> str:
        lines = [
            f"# HELP {self.name} {self.help_text}",
            f"# TYPE {self.name} histogram",
        ]
        with self._lock:
            for b in self.buckets:
                lines.append(f'{self.name}_bucket{{le="{b}"}} {self._counts[b]}')
            lines.append(f'{self.name}_bucket{{le="+Inf"}} {self._count}')
            lines.append(f"{self.name}_count {self._count}")
            lines.append(f"{self.name}_sum {self._sum:.4f}")
        return "\n".join(lines) + "\n"


class WorkerMetrics:
    """
    Thread-safe operational metrics collector for the Processing Worker.
    Provides data for /metrics (Prometheus) and /internal/stats (Ingestion API aggregator).
    """

    def __init__(self):
        self._lock = threading.Lock()
        self._processed_total: Dict[tuple, int] = defaultdict(int)
        self._dlq_total: Dict[str, int] = defaultdict(int)
        self._in_flight_buffer_size: int = 0
        self._last_flush_at: Optional[str] = None
        self._consumer_lag: Dict[str, int] = defaultdict(int)

        # Histograms
        self.kafka_to_ch_latency = SimpleHistogram(
            name="cdr_kafka_to_clickhouse_latency_seconds",
            help_text="End-to-end latency from Kafka produce to ClickHouse write",
            buckets=[0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0],
        )
        self.batch_flush_duration = SimpleHistogram(
            name="cdr_batch_flush_duration_seconds",
            help_text="ClickHouse batch flush duration in seconds",
            buckets=[0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0],
        )
        self.batch_size_at_flush = SimpleHistogram(
            name="cdr_batch_size_at_flush",
            help_text="Batch sizes at flush time",
            buckets=[10, 50, 100, 250, 500, 1000],
        )

    def record_processed(
        self, service_type: str, radio_access_type: str, count: int = 1
    ) -> None:
        with self._lock:
            self._processed_total[(service_type, radio_access_type)] += count

    def record_dlq(self, reason: str, count: int = 1) -> None:
        with self._lock:
            self._dlq_total[reason] += count

    def set_in_flight_buffer(self, size: int) -> None:
        with self._lock:
            self._in_flight_buffer_size = size

    def set_last_flush_at(self, timestamp_iso: str) -> None:
        with self._lock:
            self._last_flush_at = timestamp_iso

    def set_consumer_lag(self, partition: int, lag: int) -> None:
        with self._lock:
            self._consumer_lag[str(partition)] = lag

    def get_internal_stats(self) -> Dict[str, Any]:
        """
        Returns structured dictionary for GET /internal/stats.
        """
        with self._lock:
            total_processed = sum(self._processed_total.values())
            total_dlq = sum(self._dlq_total.values())
            in_flight = self._in_flight_buffer_size
            last_flush = self._last_flush_at

            # Breakdowns
            by_service: Dict[str, int] = defaultdict(int)
            by_rat: Dict[str, int] = defaultdict(int)
            for (svc, rat), cnt in self._processed_total.items():
                by_service[svc] += cnt
                by_rat[rat] += cnt

            dlq_reasons = dict(self._dlq_total)

        return {
            "pipeline": {
                "total_processed": total_processed,
                "total_dlq": total_dlq,
                "in_flight_buffer": in_flight,
                "last_flush_at": last_flush,
            },
            "breakdown": {
                "by_service_type": dict(by_service),
                "by_radio_access_type": dict(by_rat),
                "dlq_by_reason": dlq_reasons,
            },
        }

    def get_prometheus_metrics(self) -> str:
        """
        Formats all worker metrics into standard Prometheus exposition format.
        """
        lines = []

        # 1. cdr_processed_total
        lines.append("# HELP cdr_processed_total Total CDRs successfully written to ClickHouse")
        lines.append("# TYPE cdr_processed_total counter")
        with self._lock:
            for (svc, rat), count in self._processed_total.items():
                lines.append(f'cdr_processed_total{{service_type="{escape_label(svc)}",radio_access_type="{escape_label(rat)}"}} {count}')

            # 2. cdr_dlq_total
            lines.append("# HELP cdr_dlq_total Total CDRs routed to DLQ by reason")
            lines.append("# TYPE cdr_dlq_total counter")
            for reason, count in self._dlq_total.items():
                lines.append(f'cdr_dlq_total{{reason="{escape_label(reason)}"}} {count}')

            # 3. Gauges
            lines.append("# HELP cdr_in_flight_buffer_size Current in-memory buffer size")
            lines.append("# TYPE cdr_in_flight_buffer_size gauge")
            lines.append(f"cdr_in_flight_buffer_size {self._in_flight_buffer_size}")

            lines.append("# HELP cdr_kafka_consumer_lag Kafka consumer lag per partition")
            lines.append("# TYPE cdr_kafka_consumer_lag gauge")
            for part, lag in self._consumer_lag.items():
                lines.append(f'cdr_kafka_consumer_lag{{partition="{part}"}} {lag}')

        # 4. Histograms
        lines.append(self.kafka_to_ch_latency.to_prometheus().strip())
        lines.append(self.batch_flush_duration.to_prometheus().strip())
        lines.append(self.batch_size_at_flush.to_prometheus().strip())

        return "\n".join(lines) + "\n"


worker_metrics = WorkerMetrics()