"""TelemetryCollector for privacy-preserving statistics collection.

This module collects anonymized statistics about compression patterns
to enable cross-user learning and improve compression over time.
"""

from __future__ import annotations

import hashlib
import json
import os
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .models import (
    AnonymizedToolStats,
    CompressionEvent,
    FieldDistribution,
    RetrievalStats,
    ToolSignature,
)


@dataclass
class TelemetryConfig:
    """Configuration for telemetry collection."""

    # Enable/disable telemetry
    enabled: bool = True

    # Storage
    storage_path: str | None = None  # Path to store telemetry data (None = in-memory only)
    auto_save_interval: int = 300  # Auto-save every N seconds (0 = disabled)

    # Privacy settings
    anonymize_tool_names: bool = True  # Hash tool names
    collect_field_names: bool = False  # If False, only collect field hashes
    collect_timing: bool = True  # Collect processing time

    # Aggregation settings
    max_events_in_memory: int = 10000  # Max events to keep in memory
    min_samples_for_recommendation: int = 10  # Min samples before making recommendations

    # Export settings
    include_field_distributions: bool = True  # Include detailed field stats in export
    include_recommendations: bool = True  # Include learned recommendations


class TelemetryCollector:
    """Collects and aggregates compression telemetry.

    Thread-safe collector that maintains anonymized statistics about
    compression patterns. Can be used to:
    - Understand what tool outputs look like (structurally)
    - Track which compression strategies work best
    - Learn optimal settings per tool type
    - Export data for cross-user aggregation

    Privacy guarantees:
    - No actual data values are stored
    - Tool names are hashed by default
    - Field names can be hashed
    - No user identifiers
    - No query content
    """

    def __init__(self, config: TelemetryConfig | None = None):
        """Initialize the telemetry collector.

        Args:
            config: Configuration options. Uses defaults if not provided.
        """
        self._config = config or TelemetryConfig()
        self._lock = threading.Lock()

        # Event storage
        self._events: list[CompressionEvent] = []

        # Aggregated stats per tool signature
        self._tool_stats: dict[str, AnonymizedToolStats] = {}

        # Retrieval tracking
        self._retrieval_stats: dict[str, RetrievalStats] = {}

        # Global counters
        self._total_compressions: int = 0
        self._total_retrievals: int = 0
        self._total_tokens_saved: int = 0

        # Auto-save tracking
        self._last_save_time: float = time.time()
        self._dirty: bool = False

        # Load existing data if storage path exists
        if self._config.storage_path:
            self._load_from_disk()

    def record_compression(
        self,
        items: list[dict[str, Any]],
        original_count: int,
        compressed_count: int,
        original_tokens: int,
        compressed_tokens: int,
        strategy: str,
        *,
        tool_name: str | None = None,
        strategy_reason: str | None = None,
        crushability_score: float | None = None,
        crushability_reason: str | None = None,
        kept_first_n: int = 0,
        kept_last_n: int = 0,
        kept_errors: int = 0,
        kept_anomalies: int = 0,
        kept_by_relevance: int = 0,
        kept_by_score: int = 0,
        processing_time_ms: float = 0.0,
    ) -> None:
        """Record a compression event.

        Args:
            items: Sample items from the original array (for structure analysis).
            original_count: Original number of items.
            compressed_count: Number of items after compression.
            original_tokens: Original token count.
            compressed_tokens: Compressed token count.
            strategy: Compression strategy used.
            tool_name: Optional tool name (will be hashed if configured).
            strategy_reason: Why this strategy was chosen.
            crushability_score: Crushability analysis score.
            crushability_reason: Crushability analysis reason.
            kept_first_n: Items kept from start.
            kept_last_n: Items kept from end.
            kept_errors: Error items kept.
            kept_anomalies: Anomalous items kept.
            kept_by_relevance: Items kept by relevance score.
            kept_by_score: Items kept by score field.
            processing_time_ms: Processing time in milliseconds.
        """
        if not self._config.enabled:
            return

        # Create tool signature from items
        signature = ToolSignature.from_items(items[:10])  # Sample first 10

        # Analyze field distributions
        field_distributions: list[FieldDistribution] = []
        if self._config.include_field_distributions and items:
            field_distributions = self._analyze_fields(items[:100])  # Sample 100

        # Calculate ratios
        compression_ratio = compressed_count / original_count if original_count > 0 else 0.0
        token_reduction = 1 - (compressed_tokens / original_tokens) if original_tokens > 0 else 0.0

        # Create event
        event = CompressionEvent(
            tool_signature=signature,
            original_item_count=original_count,
            compressed_item_count=compressed_count,
            compression_ratio=compression_ratio,
            original_tokens=original_tokens,
            compressed_tokens=compressed_tokens,
            token_reduction_ratio=token_reduction,
            strategy=strategy,
            strategy_reason=strategy_reason,
            crushability_score=crushability_score,
            crushability_reason=crushability_reason,
            field_distributions=field_distributions,
            kept_first_n=kept_first_n,
            kept_last_n=kept_last_n,
            kept_errors=kept_errors,
            kept_anomalies=kept_anomalies,
            kept_by_relevance=kept_by_relevance,
            kept_by_score=kept_by_score,
            timestamp=time.time(),
            processing_time_ms=processing_time_ms,
        )

        should_save = False
        with self._lock:
            # Store event
            self._events.append(event)
            if len(self._events) > self._config.max_events_in_memory:
                self._events = self._events[-self._config.max_events_in_memory :]

            # Update aggregated stats
            self._update_tool_stats(signature, event)

            # Update global counters
            self._total_compressions += 1
            self._total_tokens_saved += original_tokens - compressed_tokens
            self._dirty = True

            # Check if auto-save needed (don't actually save while holding lock)
            should_save = self._should_auto_save()

        # Auto-save outside lock to avoid blocking other operations
        if should_save:
            self.save()

    def record_retrieval(
        self,
        tool_signature_hash: str,
        retrieval_type: str,  # "full" or "search"
        query_fields: list[str] | None = None,
    ) -> None:
        """Record a retrieval event.

        This is called when an LLM retrieves compressed content, indicating
        the compression may have been too aggressive.

        Args:
            tool_signature_hash: Hash of the tool signature.
            retrieval_type: "full" (retrieved everything) or "search" (filtered).
            query_fields: Field names mentioned in search query (will be hashed).
        """
        if not self._config.enabled:
            return

        with self._lock:
            # Get or create retrieval stats
            if tool_signature_hash not in self._retrieval_stats:
                self._retrieval_stats[tool_signature_hash] = RetrievalStats(
                    tool_signature_hash=tool_signature_hash
                )

            stats = self._retrieval_stats[tool_signature_hash]
            stats.total_retrievals += 1

            if retrieval_type == "full":
                stats.full_retrievals += 1
            else:
                stats.search_retrievals += 1

            # Track queried fields (anonymized)
            if query_fields:
                for field_name in query_fields:
                    field_hash = self._hash_field_name(field_name)
                    stats.query_field_frequency[field_hash] = (
                        stats.query_field_frequency.get(field_hash, 0) + 1
                    )

            # Update global counter
            self._total_retrievals += 1
            self._dirty = True

            # Update tool stats with retrieval info
            if tool_signature_hash in self._tool_stats:
                self._tool_stats[tool_signature_hash].retrieval_stats = stats
                self._update_recommendations(tool_signature_hash)

    def get_stats(self) -> dict[str, Any]:
        """Get overall telemetry statistics.

        Returns:
            Dictionary with aggregated statistics.
        """
        with self._lock:
            return {
                "enabled": self._config.enabled,
                "total_compressions": self._total_compressions,
                "total_retrievals": self._total_retrievals,
                "total_tokens_saved": self._total_tokens_saved,
                "global_retrieval_rate": (
                    self._total_retrievals / self._total_compressions
                    if self._total_compressions > 0
                    else 0.0
                ),
                "tool_signatures_tracked": len(self._tool_stats),
                "events_in_memory": len(self._events),
                "avg_compression_ratio": self._calculate_avg_compression_ratio(),
                "avg_token_reduction": self._calculate_avg_token_reduction(),
            }

    def get_tool_stats(self, signature_hash: str) -> AnonymizedToolStats | None:
        """Get statistics for a specific tool signature.

        Args:
            signature_hash: The tool signature hash.

        Returns:
            AnonymizedToolStats if found, None otherwise.
        """
        with self._lock:
            return self._tool_stats.get(signature_hash)

    def get_all_tool_stats(self) -> dict[str, AnonymizedToolStats]:
        """Get statistics for all tracked tool signatures.

        Returns:
            Dictionary mapping signature hash to stats.
        """
        with self._lock:
            return dict(self._tool_stats)

    def get_recommendations(self, signature_hash: str) -> dict[str, Any] | None:
        """Get learned recommendations for a tool signature.

        Args:
            signature_hash: The tool signature hash.

        Returns:
            Recommendations dictionary if available, None otherwise.
        """
        with self._lock:
            stats = self._tool_stats.get(signature_hash)
            if not stats or stats.sample_size < self._config.min_samples_for_recommendation:
                return None

            return {
                "signature_hash": signature_hash,
                "recommended_min_items": stats.recommended_min_items,
                "recommended_preserve_fields": stats.recommended_preserve_fields,
                "skip_compression_recommended": stats.skip_compression_recommended,
                "confidence": stats.confidence,
                "based_on_samples": stats.sample_size,
                "retrieval_rate": (
                    stats.retrieval_stats.retrieval_rate if stats.retrieval_stats else None
                ),
            }

    def export_stats(self) -> dict[str, Any]:
        """Export all telemetry data for aggregation.

        This is the data that can be sent to a central server for
        cross-user learning (with user consent).

        Returns:
            Complete telemetry export.
        """
        with self._lock:
            export = {
                "version": "1.0",
                "export_timestamp": time.time(),
                "summary": {
                    "total_compressions": self._total_compressions,
                    "total_retrievals": self._total_retrievals,
                    "total_tokens_saved": self._total_tokens_saved,
                    "tool_signatures_tracked": len(self._tool_stats),
                },
                "tool_stats": {
                    sig_hash: stats.to_dict() for sig_hash, stats in self._tool_stats.items()
                },
            }

            if self._config.include_recommendations:
                export["recommendations"] = {
                    sig_hash: {
                        "recommended_min_items": stats.recommended_min_items,
                        "skip_compression_recommended": stats.skip_compression_recommended,
                        "confidence": stats.confidence,
                    }
                    for sig_hash, stats in self._tool_stats.items()
                    if stats.sample_size >= self._config.min_samples_for_recommendation
                }

            return export

    def import_stats(self, data: dict[str, Any]) -> None:
        """Import telemetry data from another source.

        This allows merging stats from multiple users for cross-user learning.

        Args:
            data: Exported telemetry data.
        """
        if not self._config.enabled:
            return

        with self._lock:
            # Import summary counters
            summary = data.get("summary", {})
            self._total_compressions += summary.get("total_compressions", 0)
            self._total_retrievals += summary.get("total_retrievals", 0)
            self._total_tokens_saved += summary.get("total_tokens_saved", 0)

            # Import tool stats
            tool_stats_data = data.get("tool_stats", {})
            for sig_hash, stats_dict in tool_stats_data.items():
                if sig_hash in self._tool_stats:
                    # Merge with existing
                    existing = self._tool_stats[sig_hash]
                    imported = AnonymizedToolStats.from_dict(stats_dict)
                    self._merge_tool_stats(existing, imported)
                else:
                    # Add new
                    self._tool_stats[sig_hash] = AnonymizedToolStats.from_dict(stats_dict)

            self._dirty = True

    def clear(self) -> None:
        """Clear all telemetry data. Mainly for testing."""
        with self._lock:
            self._events.clear()
            self._tool_stats.clear()
            self._retrieval_stats.clear()
            self._total_compressions = 0
            self._total_retrievals = 0
            self._total_tokens_saved = 0
            self._dirty = False

    def save(self) -> None:
        """Save telemetry data to disk."""
        if not self._config.storage_path:
            return

        with self._lock:
            # Build export data inline to avoid deadlock (export_stats also acquires lock)
            data = {
                "version": "1.0",
                "export_timestamp": time.time(),
                "summary": {
                    "total_compressions": self._total_compressions,
                    "total_retrievals": self._total_retrievals,
                    "total_tokens_saved": self._total_tokens_saved,
                    "tool_signatures_tracked": len(self._tool_stats),
                },
                "tool_stats": {
                    sig_hash: stats.to_dict() for sig_hash, stats in self._tool_stats.items()
                },
            }

            if self._config.include_recommendations:
                data["recommendations"] = {
                    sig_hash: {
                        "recommended_min_items": stats.recommended_min_items,
                        "skip_compression_recommended": stats.skip_compression_recommended,
                        "confidence": stats.confidence,
                    }
                    for sig_hash, stats in self._tool_stats.items()
                    if stats.sample_size >= self._config.min_samples_for_recommendation
                }

            path = Path(self._config.storage_path)
            path.parent.mkdir(parents=True, exist_ok=True)

            with open(path, "w") as f:
                json.dump(data, f, indent=2)

            self._dirty = False
            self._last_save_time = time.time()

    def _load_from_disk(self) -> None:
        """Load telemetry data from disk."""
        if not self._config.storage_path:
            return

        path = Path(self._config.storage_path)
        if not path.exists():
            return

        try:
            with open(path) as f:
                data = json.load(f)
            self.import_stats(data)
            self._dirty = False
        except (json.JSONDecodeError, OSError):
            pass  # Start fresh if file is corrupted

    def _analyze_fields(self, items: list[dict[str, Any]]) -> list[FieldDistribution]:
        """Analyze field distributions in items."""
        if not items:
            return []

        distributions: list[FieldDistribution] = []

        # Get all field names from first item
        sample = items[0] if isinstance(items[0], dict) else {}
        for field_name, _sample_value in sample.items():
            # Collect all values for this field
            values = [
                item.get(field_name)
                for item in items
                if isinstance(item, dict) and field_name in item
            ]

            if not values:
                continue

            dist = self._create_field_distribution(field_name, values)
            distributions.append(dist)

        return distributions

    def _create_field_distribution(
        self,
        field_name: str,
        values: list[Any],
    ) -> FieldDistribution:
        """Create a FieldDistribution from values."""
        field_hash = self._hash_field_name(field_name)

        # Determine type
        type_counts: dict[str, int] = {}
        for v in values:
            if isinstance(v, str):
                type_counts["string"] = type_counts.get("string", 0) + 1
            elif isinstance(v, bool):
                type_counts["boolean"] = type_counts.get("boolean", 0) + 1
            elif isinstance(v, (int, float)):
                type_counts["numeric"] = type_counts.get("numeric", 0) + 1
            elif isinstance(v, list):
                type_counts["array"] = type_counts.get("array", 0) + 1
            elif isinstance(v, dict):
                type_counts["object"] = type_counts.get("object", 0) + 1
            elif v is None:
                type_counts["null"] = type_counts.get("null", 0) + 1

        # Get dominant type
        if not type_counts:
            field_type = "null"
        elif len(type_counts) > 1:
            field_type = "mixed"
        else:
            field_type = list(type_counts.keys())[0]

        dist = FieldDistribution(
            field_name_hash=field_hash,
            field_type=field_type,  # type: ignore[arg-type]
        )

        # Type-specific analysis
        if field_type == "string":
            str_values = [v for v in values if isinstance(v, str)]
            if str_values:
                dist.avg_length = sum(len(s) for s in str_values) / len(str_values)
                unique_count = len(set(str_values))
                dist.unique_ratio = unique_count / len(str_values)
                dist.looks_like_id = dist.unique_ratio > 0.9 and dist.avg_length > 5

        elif field_type == "numeric":
            num_values = [v for v in values if isinstance(v, (int, float))]
            # Filter out infinity and NaN which can cause issues
            num_values = [
                v
                for v in num_values
                if not (
                    isinstance(v, float) and (v != v or v == float("inf") or v == float("-inf"))
                )
            ]
            if num_values:
                dist.has_negative = any(v < 0 for v in num_values)
                # Safe integer check (avoid OverflowError from int(inf))
                dist.is_integer = all(
                    isinstance(v, int) or (isinstance(v, float) and v.is_integer())
                    for v in num_values
                )

                if len(num_values) > 1:
                    mean = sum(num_values) / len(num_values)
                    variance = sum((v - mean) ** 2 for v in num_values) / len(num_values)
                    dist.has_variance = variance > 0

                    if variance == 0:
                        dist.variance_bucket = "zero"
                    elif variance < 10:
                        dist.variance_bucket = "low"
                    elif variance < 1000:
                        dist.variance_bucket = "medium"
                    else:
                        dist.variance_bucket = "high"

                    # Check for outliers
                    std = variance**0.5
                    if std > 0:
                        outliers = sum(1 for v in num_values if abs(v - mean) > 2 * std)
                        dist.has_outliers = outliers > 0

                    # Pattern detection
                    sorted_vals = sorted(num_values)
                    is_monotonic = (
                        sorted_vals == num_values or list(reversed(sorted_vals)) == num_values
                    )
                    if is_monotonic and dist.variance_bucket in ("medium", "high"):
                        dist.is_likely_score = True

        elif field_type == "array":
            arr_values = [v for v in values if isinstance(v, list)]
            if arr_values:
                dist.avg_array_length = sum(len(a) for a in arr_values) / len(arr_values)

        return dist

    def _update_tool_stats(self, signature: ToolSignature, event: CompressionEvent) -> None:
        """Update aggregated stats for a tool signature."""
        sig_hash = signature.structure_hash

        if sig_hash not in self._tool_stats:
            self._tool_stats[sig_hash] = AnonymizedToolStats(signature=signature)

        stats = self._tool_stats[sig_hash]

        # Update counts
        stats.total_compressions += 1
        stats.total_items_seen += event.original_item_count
        stats.total_items_kept += event.compressed_item_count
        stats.sample_size += 1

        # Update averages (rolling)
        n = stats.total_compressions
        stats.avg_compression_ratio = (
            stats.avg_compression_ratio * (n - 1) + event.compression_ratio
        ) / n
        stats.avg_token_reduction = (
            stats.avg_token_reduction * (n - 1) + event.token_reduction_ratio
        ) / n

        # Update strategy counts
        strategy = event.strategy
        stats.strategy_counts[strategy] = stats.strategy_counts.get(strategy, 0) + 1

        # Update confidence based on sample size
        stats.confidence = min(0.95, stats.sample_size / 100)

        # Update recommendations
        self._update_recommendations(sig_hash)

    def _update_recommendations(self, sig_hash: str) -> None:
        """Update recommendations based on current data."""
        if sig_hash not in self._tool_stats:
            return

        stats = self._tool_stats[sig_hash]

        # Not enough data yet
        if stats.sample_size < self._config.min_samples_for_recommendation:
            return

        # Check retrieval rate to determine if compression is too aggressive
        if stats.retrieval_stats:
            retrieval_rate = stats.retrieval_stats.retrieval_rate
            full_rate = stats.retrieval_stats.full_retrieval_rate

            # High retrieval rate = compression too aggressive
            if retrieval_rate > 0.5:
                if full_rate > 0.8:
                    # Almost all retrievals are full = skip compression
                    stats.skip_compression_recommended = True
                else:
                    # Increase min items
                    stats.recommended_min_items = 50
            elif retrieval_rate > 0.2:
                # Medium retrieval rate = slightly less aggressive
                stats.recommended_min_items = 30
            else:
                # Low retrieval rate = current settings work
                stats.recommended_min_items = 15

            # Track frequently queried fields
            if stats.retrieval_stats.query_field_frequency:
                top_fields = sorted(
                    stats.retrieval_stats.query_field_frequency.items(),
                    key=lambda x: x[1],
                    reverse=True,
                )[:5]
                stats.recommended_preserve_fields = [f for f, _ in top_fields]

    def _merge_tool_stats(
        self,
        existing: AnonymizedToolStats,
        imported: AnonymizedToolStats,
    ) -> None:
        """Merge imported stats into existing."""
        # Weighted average based on sample sizes
        total_samples = existing.sample_size + imported.sample_size
        if total_samples == 0:
            return

        w_existing = existing.sample_size / total_samples
        w_imported = imported.sample_size / total_samples

        existing.total_compressions += imported.total_compressions
        existing.total_items_seen += imported.total_items_seen
        existing.total_items_kept += imported.total_items_kept
        existing.avg_compression_ratio = (
            existing.avg_compression_ratio * w_existing
            + imported.avg_compression_ratio * w_imported
        )
        existing.avg_token_reduction = (
            existing.avg_token_reduction * w_existing + imported.avg_token_reduction * w_imported
        )
        existing.sample_size = total_samples

        # Merge strategy counts
        for strategy, count in imported.strategy_counts.items():
            existing.strategy_counts[strategy] = existing.strategy_counts.get(strategy, 0) + count

        # Update confidence
        existing.confidence = min(0.95, total_samples / 100)

    def _hash_field_name(self, field_name: str) -> str:
        """Hash a field name for anonymization."""
        if self._config.collect_field_names:
            return field_name
        return hashlib.sha256(field_name.encode()).hexdigest()[:8]

    def _calculate_avg_compression_ratio(self) -> float:
        """Calculate average compression ratio across all tools."""
        if not self._tool_stats:
            return 0.0
        ratios = [s.avg_compression_ratio for s in self._tool_stats.values()]
        return sum(ratios) / len(ratios)

    def _calculate_avg_token_reduction(self) -> float:
        """Calculate average token reduction across all tools."""
        if not self._tool_stats:
            return 0.0
        reductions = [s.avg_token_reduction for s in self._tool_stats.values()]
        return sum(reductions) / len(reductions)

    def _should_auto_save(self) -> bool:
        """Check if auto-save should run. Must be called with lock held."""
        if not self._config.auto_save_interval or not self._config.storage_path:
            return False

        if not self._dirty:
            return False

        elapsed = time.time() - self._last_save_time
        return elapsed >= self._config.auto_save_interval


# Global collector instance (lazy initialization)
_telemetry_collector: TelemetryCollector | None = None
_collector_lock = threading.Lock()


def get_telemetry_collector(
    config: TelemetryConfig | None = None,
) -> TelemetryCollector:
    """Get the global telemetry collector instance.

    Args:
        config: Configuration (only used on first call).

    Returns:
        Global TelemetryCollector instance.
    """
    global _telemetry_collector

    if _telemetry_collector is None:
        with _collector_lock:
            if _telemetry_collector is None:
                # Check environment for opt-out
                if os.environ.get("HEADROOM_TELEMETRY_DISABLED", "").lower() in ("1", "true"):
                    config = config or TelemetryConfig()
                    config.enabled = False

                _telemetry_collector = TelemetryCollector(config)

    return _telemetry_collector


def reset_telemetry_collector() -> None:
    """Reset the global telemetry collector. Mainly for testing."""
    global _telemetry_collector

    with _collector_lock:
        if _telemetry_collector is not None:
            _telemetry_collector.clear()
        _telemetry_collector = None
