"""Catalog statistics collection for EarthCatalog ingestion pipeline.

This module provides comprehensive statistics tracking during STAC item ingestion,
enabling detailed insights into catalog composition, spatial/temporal distribution,
data quality, and processing performance without requiring post-hoc queries.

Key Features:
    - HyperLogLog for memory-efficient approximate unique counting
    - Spatial distribution tracking (items per cell, hotspots)
    - Temporal distribution (items per year/month)
    - Overhead metrics (spanning items, duplication ratio)
    - Data quality metrics (null geometries, missing datetimes)
    - Processing performance metrics

Memory Efficiency:
    Uses HyperLogLog algorithm for unique counting, requiring only ~1.5KB
    to track billions of unique items with ~2% error rate. This allows
    accurate overhead calculation even for 100M+ item catalogs.

Thread Safety:
    The IngestionStatistics class is designed for single-threaded use within
    each worker. For distributed processing, each worker maintains its own
    statistics instance, which are merged at the end of processing.

Example:
    >>> stats = IngestionStatistics()
    >>> for item in items:
    ...     stats.record_item(item, tiles=["abc123"], is_spanning=False)
    >>> summary = stats.get_summary()
    >>> print(f"Unique items: {summary['unique_granules']}")
"""

import hashlib
import math
import time
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any


class HyperLogLog:
    """Memory-efficient probabilistic cardinality estimator.

    Implements the HyperLogLog algorithm for approximate distinct counting.
    Uses only ~1.5KB of memory to estimate cardinality with ~2% error rate,
    regardless of the number of unique elements.

    The algorithm works by:
    1. Hashing each element to a 64-bit value
    2. Partitioning hash space into 2^p buckets (default p=14 = 16384 buckets)
    3. Tracking the maximum number of leading zeros seen in each bucket
    4. Using harmonic mean of bucket estimates for final cardinality

    Attributes:
        p: Precision parameter (default 14, giving 16384 registers)
        m: Number of registers (2^p)
        registers: Array of maximum leading zero counts per bucket

    Example:
        >>> hll = HyperLogLog(precision=14)
        >>> for item_id in item_ids:
        ...     hll.add(item_id)
        >>> estimated_unique = hll.count()
        >>> print(f"Approximately {estimated_unique} unique items")

    Error Rate:
        Standard error is approximately 1.04 / sqrt(m)
        With p=14 (m=16384): error ≈ 0.81% (typically <2% in practice)

    Memory Usage:
        Uses 1 byte per register = 2^p bytes
        p=14: 16KB, p=12: 4KB, p=10: 1KB
    """

    def __init__(self, precision: int = 14):
        """Initialize HyperLogLog with specified precision.

        Args:
            precision: Number of bits for bucket addressing (4-16).
                      Higher precision = more accuracy but more memory.
                      Default 14 gives ~0.8% error with 16KB memory.
        """
        if not 4 <= precision <= 16:
            raise ValueError("Precision must be between 4 and 16")

        self.p = precision
        self.m = 1 << precision  # 2^p registers
        self.registers = [0] * self.m

        # Alpha constant for bias correction (depends on m)
        if self.m >= 128:
            self.alpha = 0.7213 / (1 + 1.079 / self.m)
        elif self.m == 64:
            self.alpha = 0.709
        elif self.m == 32:
            self.alpha = 0.697
        else:
            self.alpha = 0.673

    def _hash(self, value: str) -> int:
        """Hash a string value to a 64-bit integer."""
        # Use MD5 for speed and reasonable distribution (not for security)
        h = hashlib.md5(value.encode("utf-8"), usedforsecurity=False).digest()
        return int.from_bytes(h[:8], "big")

    def _leading_zeros(self, value: int, max_bits: int = 64) -> int:
        """Count leading zeros in the binary representation."""
        if value == 0:
            return max_bits
        return max_bits - value.bit_length()

    def add(self, value: str) -> None:
        """Add an element to the HyperLogLog.

        Args:
            value: String value to add (typically an item ID)
        """
        h = self._hash(value)

        # Use first p bits for bucket index
        bucket = h >> (64 - self.p)

        # Use remaining bits for leading zero count
        remaining = h & ((1 << (64 - self.p)) - 1)
        zeros = self._leading_zeros(remaining, 64 - self.p) + 1

        # Update register if this is a new maximum
        self.registers[bucket] = max(self.registers[bucket], zeros)

    def count(self) -> int:
        """Estimate the cardinality (number of unique elements).

        Returns:
            Estimated number of unique elements added.
        """
        # Compute harmonic mean of 2^register values
        indicator = sum(2.0 ** (-r) for r in self.registers)
        estimate = self.alpha * self.m * self.m / indicator

        # Apply small range correction (linear counting)
        if estimate <= 2.5 * self.m:
            zeros = self.registers.count(0)
            if zeros > 0:
                estimate = self.m * math.log(self.m / zeros)

        # Apply large range correction
        elif estimate > (1 << 32) / 30:
            estimate = -(1 << 32) * math.log(1 - estimate / (1 << 32))

        return int(estimate)

    def merge(self, other: "HyperLogLog") -> None:
        """Merge another HyperLogLog into this one.

        Useful for combining statistics from multiple workers.

        Args:
            other: Another HyperLogLog instance with same precision
        """
        if self.p != other.p:
            raise ValueError("Cannot merge HyperLogLogs with different precision")

        for i in range(self.m):
            self.registers[i] = max(self.registers[i], other.registers[i])


@dataclass
class IngestionStatistics:
    """Comprehensive statistics collector for STAC ingestion pipeline.

    Tracks all relevant metrics during ingestion with minimal memory overhead.
    Designed to be used within a single worker process and merged after
    distributed processing completes.

    Categories of statistics tracked:
        - Core counts: stored references, unique granules (via HyperLogLog)
        - Overhead: spanning items, duplication metrics
        - Spatial: items per cell, hotspots, bounding box
        - Temporal: items per year/month, temporal extent
        - Quality: null geometries, missing datetimes, geometry types
        - Missions: items per dataset/collection
        - Processing: timing, failures, throughput

    Example:
        >>> stats = IngestionStatistics()
        >>> stats.start_processing()
        >>> for item in items:
        ...     tiles = grid.tiles_for_geometry(item['geometry'])
        ...     stats.record_item(item, tiles, is_spanning=len(tiles)>1)
        >>> stats.finish_processing()
        >>> summary = stats.get_summary()
    """

    # HyperLogLog for unique counting (precision 14 = ~0.8% error, 16KB memory)
    unique_ids: HyperLogLog = field(default_factory=lambda: HyperLogLog(precision=14))

    # Core counts
    stored_references: int = 0
    items_routed_to_global: int = 0

    # Spanning metrics
    spanning_items_count: int = 0
    total_tiles_for_spanning: int = 0  # Sum of tiles for all spanning items
    max_tiles_per_item: int = 0
    tiles_per_spanning_histogram: dict[int, int] = field(default_factory=lambda: defaultdict(int))

    # Spatial distribution
    items_per_cell: dict[str, int] = field(default_factory=lambda: defaultdict(int))
    bbox_min_lon: float = 180.0
    bbox_max_lon: float = -180.0
    bbox_min_lat: float = 90.0
    bbox_max_lat: float = -90.0

    # Temporal distribution (year -> month -> count)
    items_per_year_month: dict[str, dict[str, int]] = field(
        default_factory=lambda: defaultdict(lambda: defaultdict(int))
    )
    earliest_datetime: datetime | None = None
    latest_datetime: datetime | None = None

    # Data quality
    null_geometry_count: int = 0
    missing_datetime_count: int = 0
    geometry_types: dict[str, int] = field(default_factory=lambda: defaultdict(int))

    # Mission breakdown
    items_per_mission: dict[str, int] = field(default_factory=lambda: defaultdict(int))

    # Processing metrics
    start_time: float | None = None
    end_time: float | None = None
    urls_processed: int = 0
    urls_failed: int = 0

    # File statistics (populated during consolidation)
    file_sizes: list[int] = field(default_factory=list)
    items_per_file: list[int] = field(default_factory=list)

    # Consolidation metrics
    duplicates_removed: int = 0
    new_items: int = 0
    existing_items: int = 0

    def start_processing(self) -> None:
        """Mark the start of processing for timing metrics."""
        self.start_time = time.time()

    def finish_processing(self) -> None:
        """Mark the end of processing for timing metrics."""
        self.end_time = time.time()

    def record_url_processed(self, success: bool = True) -> None:
        """Record a URL processing attempt.

        Args:
            success: Whether the URL was successfully processed
        """
        self.urls_processed += 1
        if not success:
            self.urls_failed += 1

    def record_item(
        self,
        item: dict[str, Any],
        tiles: list[str],
        is_spanning: bool,
        routed_to_global: bool = False,
        mission: str | None = None,
    ) -> None:
        """Record statistics for a single STAC item.

        Args:
            item: STAC item dictionary
            tiles: List of tile IDs this item maps to
            is_spanning: Whether item spans multiple tiles
            routed_to_global: Whether item was routed to global partition
            mission: Extracted mission/dataset name
        """
        item_id = item.get("id", "")

        # Track unique IDs via HyperLogLog
        if item_id:
            self.unique_ids.add(item_id)

        # Count stored reference (this counts each tile assignment)
        if routed_to_global:
            self.stored_references += 1
            self.items_routed_to_global += 1
        else:
            # Item stored once per tile it intersects
            self.stored_references += len(tiles) if tiles else 1

        # Track spanning metrics
        if is_spanning and not routed_to_global:
            self.spanning_items_count += 1
            tile_count = len(tiles)
            self.total_tiles_for_spanning += tile_count
            self.max_tiles_per_item = max(self.max_tiles_per_item, tile_count)

            # Bucket into histogram (1, 2, 3-5, 6-10, 11-20, 21-50, 50+)
            if tile_count <= 2:
                bucket = tile_count
            elif tile_count <= 5:
                bucket = 5
            elif tile_count <= 10:
                bucket = 10
            elif tile_count <= 20:
                bucket = 20
            elif tile_count <= 50:
                bucket = 50
            else:
                bucket = 100  # 50+
            self.tiles_per_spanning_histogram[bucket] += 1

        # Track spatial distribution
        for tile in tiles:
            self.items_per_cell[tile] += 1

        # Update bounding box from geometry
        geometry = item.get("geometry")
        if geometry:
            self._update_bbox_from_geometry(geometry)
            geom_type = geometry.get("type", "Unknown")
            self.geometry_types[geom_type] += 1
        else:
            self.null_geometry_count += 1

        # Track temporal distribution
        props = item.get("properties", {})
        dt_str = props.get("datetime")
        if dt_str:
            self._record_datetime(dt_str)
        else:
            self.missing_datetime_count += 1

        # Track mission
        if mission:
            self.items_per_mission[mission] += 1

    def _update_bbox_from_geometry(self, geometry: dict[str, Any]) -> None:
        """Update bounding box from GeoJSON geometry."""
        try:
            # Extract coordinates based on geometry type
            geom_type = geometry.get("type", "")
            coords = geometry.get("coordinates", [])

            if not coords:
                return

            # Flatten coordinates to list of [lon, lat] pairs
            flat_coords = self._flatten_coordinates(coords, geom_type)

            for lon, lat in flat_coords:
                self.bbox_min_lon = min(self.bbox_min_lon, lon)
                self.bbox_max_lon = max(self.bbox_max_lon, lon)
                self.bbox_min_lat = min(self.bbox_min_lat, lat)
                self.bbox_max_lat = max(self.bbox_max_lat, lat)
        except (TypeError, ValueError, KeyError, IndexError):
            # Silently ignore malformed geometries that can't be processed
            pass

    def _flatten_coordinates(self, coords: Any, geom_type: str) -> list[tuple[float, float]]:
        """Flatten nested coordinate arrays to list of (lon, lat) tuples."""
        result = []

        if geom_type == "Point":
            if len(coords) >= 2:
                result.append((coords[0], coords[1]))
        elif geom_type == "LineString":
            for coord in coords:
                if len(coord) >= 2:
                    result.append((coord[0], coord[1]))
        elif geom_type == "Polygon":
            for ring in coords:
                for coord in ring:
                    if len(coord) >= 2:
                        result.append((coord[0], coord[1]))
        elif geom_type == "MultiPoint":
            for coord in coords:
                if len(coord) >= 2:
                    result.append((coord[0], coord[1]))
        elif geom_type == "MultiLineString":
            for line in coords:
                for coord in line:
                    if len(coord) >= 2:
                        result.append((coord[0], coord[1]))
        elif geom_type == "MultiPolygon":
            for polygon in coords:
                for ring in polygon:
                    for coord in ring:
                        if len(coord) >= 2:
                            result.append((coord[0], coord[1]))

        return result

    def _record_datetime(self, dt_str: str) -> None:
        """Record datetime for temporal statistics.

        Uses datetime.fromisoformat() for fast parsing of ISO 8601 strings.
        Falls back gracefully for non-standard formats.
        """
        try:
            # Fast path: datetime.fromisoformat() is ~100x faster than pd.to_datetime()
            # Handle 'Z' suffix which fromisoformat doesn't support directly
            if dt_str.endswith("Z"):
                dt_str = dt_str[:-1] + "+00:00"
            dt = datetime.fromisoformat(dt_str)

            # Update temporal extent
            if self.earliest_datetime is None or dt < self.earliest_datetime:
                self.earliest_datetime = dt
            if self.latest_datetime is None or dt > self.latest_datetime:
                self.latest_datetime = dt

            # Track by year/month
            year = str(dt.year)
            month = f"{dt.month:02d}"
            self.items_per_year_month[year][month] += 1

        except (ValueError, TypeError):
            self.missing_datetime_count += 1

    def record_file(self, file_size_bytes: int, item_count: int) -> None:
        """Record statistics for a written file.

        Args:
            file_size_bytes: Size of the written file in bytes
            item_count: Number of items in the file
        """
        self.file_sizes.append(file_size_bytes)
        self.items_per_file.append(item_count)

    def record_consolidation(self, new_items: int, existing_items: int, duplicates_removed: int) -> None:
        """Record consolidation statistics.

        Args:
            new_items: Number of new items added
            existing_items: Number of items that existed before
            duplicates_removed: Number of duplicate items removed
        """
        self.new_items += new_items
        self.existing_items += existing_items
        self.duplicates_removed += duplicates_removed

    def merge(self, other: "IngestionStatistics") -> None:
        """Merge statistics from another instance.

        Used to combine statistics from multiple workers in distributed processing.

        Args:
            other: Another IngestionStatistics instance to merge
        """
        # Merge HyperLogLog
        self.unique_ids.merge(other.unique_ids)

        # Merge counts
        self.stored_references += other.stored_references
        self.items_routed_to_global += other.items_routed_to_global
        self.spanning_items_count += other.spanning_items_count
        self.total_tiles_for_spanning += other.total_tiles_for_spanning
        self.max_tiles_per_item = max(self.max_tiles_per_item, other.max_tiles_per_item)

        # Merge histograms
        for bucket, count in other.tiles_per_spanning_histogram.items():
            self.tiles_per_spanning_histogram[bucket] += count

        # Merge spatial distribution
        for cell, count in other.items_per_cell.items():
            self.items_per_cell[cell] += count

        # Merge bbox
        self.bbox_min_lon = min(self.bbox_min_lon, other.bbox_min_lon)
        self.bbox_max_lon = max(self.bbox_max_lon, other.bbox_max_lon)
        self.bbox_min_lat = min(self.bbox_min_lat, other.bbox_min_lat)
        self.bbox_max_lat = max(self.bbox_max_lat, other.bbox_max_lat)

        # Merge temporal distribution
        for year, months in other.items_per_year_month.items():
            for month, count in months.items():
                self.items_per_year_month[year][month] += count

        # Merge temporal extent
        if other.earliest_datetime:
            if self.earliest_datetime is None or other.earliest_datetime < self.earliest_datetime:
                self.earliest_datetime = other.earliest_datetime
        if other.latest_datetime:
            if self.latest_datetime is None or other.latest_datetime > self.latest_datetime:
                self.latest_datetime = other.latest_datetime

        # Merge quality metrics
        self.null_geometry_count += other.null_geometry_count
        self.missing_datetime_count += other.missing_datetime_count
        for geom_type, count in other.geometry_types.items():
            self.geometry_types[geom_type] += count

        # Merge mission counts
        for mission, count in other.items_per_mission.items():
            self.items_per_mission[mission] += count

        # Merge processing metrics
        self.urls_processed += other.urls_processed
        self.urls_failed += other.urls_failed

        # Merge file statistics
        self.file_sizes.extend(other.file_sizes)
        self.items_per_file.extend(other.items_per_file)

        # Merge consolidation metrics
        self.new_items += other.new_items
        self.existing_items += other.existing_items
        self.duplicates_removed += other.duplicates_removed

    def get_summary(self) -> dict[str, Any]:
        """Generate comprehensive statistics summary for schema output.

        Returns:
            Dictionary containing all collected statistics formatted for
            inclusion in the catalog schema JSON.
        """
        unique_granules = self.unique_ids.count()

        # Calculate derived metrics
        duplication_ratio = self.stored_references / unique_granules if unique_granules > 0 else 1.0
        overhead_percentage = (
            ((self.stored_references - unique_granules) / unique_granules) * 100 if unique_granules > 0 else 0.0
        )
        spanning_percentage = (self.spanning_items_count / unique_granules) * 100 if unique_granules > 0 else 0.0
        avg_tiles_per_spanning = (
            self.total_tiles_for_spanning / self.spanning_items_count if self.spanning_items_count > 0 else 0.0
        )

        # Calculate spatial statistics
        cell_counts = list(self.items_per_cell.values())
        spatial_stats = self._calculate_distribution_stats(cell_counts)

        # Get hotspot cells (top 10 by count)
        sorted_cells = sorted(self.items_per_cell.items(), key=lambda x: x[1], reverse=True)[:10]
        hotspot_cells = [{"cell": cell, "count": count} for cell, count in sorted_cells]

        # Calculate file statistics
        file_size_stats = self._calculate_distribution_stats(self.file_sizes)
        items_per_file_stats = self._calculate_distribution_stats(self.items_per_file)

        # Calculate processing metrics
        duration = self.end_time - self.start_time if self.start_time and self.end_time else 0.0
        items_per_second = unique_granules / duration if duration > 0 else 0.0

        # Format temporal distribution
        temporal_distribution = {}
        for year in sorted(self.items_per_year_month.keys()):
            year_data: dict[str, Any] = {"total": 0, "months": {}}
            for month in sorted(self.items_per_year_month[year].keys()):
                count = self.items_per_year_month[year][month]
                year_data["months"][month] = count
                year_data["total"] += count
            temporal_distribution[year] = year_data

        # Build summary
        summary = {
            # Core counts
            "unique_granules": unique_granules,
            "stored_references": self.stored_references,
            "total_partitions": len(self.items_per_cell),
            "total_files": len(self.items_per_file) if self.items_per_file else len(self.items_per_cell),
            # Overhead metrics
            "overhead": {
                "spanning_items": self.spanning_items_count,
                "spanning_percentage": round(spanning_percentage, 2),
                "duplication_ratio": round(duplication_ratio, 3),
                "overhead_percentage": round(overhead_percentage, 2),
                "avg_tiles_per_spanning_item": round(avg_tiles_per_spanning, 2),
                "max_tiles_per_item": self.max_tiles_per_item,
                "tiles_distribution": dict(self.tiles_per_spanning_histogram),
            },
            # Global partition metrics
            "global_partition": {
                "items_routed_to_global": self.items_routed_to_global,
                "percentage_global": round(
                    (self.items_routed_to_global / unique_granules) * 100 if unique_granules > 0 else 0.0, 2
                ),
            },
            # Spatial distribution
            "spatial": {
                "bbox": (
                    [
                        round(self.bbox_min_lon, 6),
                        round(self.bbox_min_lat, 6),
                        round(self.bbox_max_lon, 6),
                        round(self.bbox_max_lat, 6),
                    ]
                    if self.bbox_min_lon <= self.bbox_max_lon
                    else None
                ),
                "cells_with_data": len(self.items_per_cell),
                "items_per_cell": spatial_stats,
                "hotspot_cells": hotspot_cells,
            },
            # Temporal distribution
            "temporal": {
                "earliest": (self.earliest_datetime.isoformat() + "Z" if self.earliest_datetime else None),
                "latest": (self.latest_datetime.isoformat() + "Z" if self.latest_datetime else None),
                "years_with_data": sorted(self.items_per_year_month.keys()),
                "distribution": temporal_distribution,
            },
            # Data quality
            "quality": {
                "null_geometries": self.null_geometry_count,
                "missing_datetime": self.missing_datetime_count,
                "geometry_types": dict(self.geometry_types),
            },
            # Mission breakdown
            "missions": dict(self.items_per_mission),
            # File statistics
            "files": {
                "total_count": len(self.items_per_file) if self.items_per_file else 0,
                "total_size_bytes": sum(self.file_sizes) if self.file_sizes else None,
                "size_stats": file_size_stats if self.file_sizes else None,
                "items_per_file": items_per_file_stats if self.items_per_file else None,
            },
            # Processing metrics
            "processing": {
                "run_timestamp": (datetime.now(UTC).isoformat().replace("+00:00", "Z")),
                "duration_seconds": round(duration, 2),
                "urls_processed": self.urls_processed,
                "urls_failed": self.urls_failed,
                "success_rate": round(
                    (
                        ((self.urls_processed - self.urls_failed) / self.urls_processed) * 100
                        if self.urls_processed > 0
                        else 100.0
                    ),
                    2,
                ),
                "items_per_second": round(items_per_second, 2),
                "new_items": self.new_items,
                "existing_items": self.existing_items,
                "duplicates_removed": self.duplicates_removed,
            },
        }

        return summary

    def _calculate_distribution_stats(self, values: list[int] | list[float]) -> dict[str, float] | None:
        """Calculate distribution statistics for a list of values.

        Args:
            values: List of numeric values

        Returns:
            Dictionary with min, max, mean, median, std_dev, or None if empty
        """
        if not values:
            return None

        sorted_values = sorted(values)
        n = len(sorted_values)

        mean = sum(sorted_values) / n

        # Calculate median
        if n % 2 == 0:
            median = (sorted_values[n // 2 - 1] + sorted_values[n // 2]) / 2
        else:
            median = sorted_values[n // 2]

        # Calculate standard deviation
        variance = sum((x - mean) ** 2 for x in sorted_values) / n
        std_dev = variance**0.5

        return {
            "min": sorted_values[0],
            "max": sorted_values[-1],
            "mean": round(mean, 2),
            "median": round(median, 2),
            "std_dev": round(std_dev, 2),
        }
