import os
import time
import threading
import logging
from datetime import datetime, timezone
from typing import Dict, List, Any, Optional

from matrice_inference.server.stream.worker_metrics import WorkerMetrics, MetricSnapshot
from matrice_inference.server.stream.metric_publisher import MetricPublisher, KafkaMetricPublisher, NoOpMetricPublisher 

logger = logging.getLogger(__name__)

class InferenceMetricLogger:
    """
    Background aggregator for worker metrics with periodic publishing.
    
    This class:
    - Runs on a dedicated background thread using threading.Timer
    - Periodically collects metrics from all workers via StreamingPipeline
    - Aggregates by worker_type (merges multiple instances)
    - Produces InferenceMetricLog matching the required schema
    - Publishes via configurable MetricPublisher
    - Handles graceful shutdown with timeout
    
    Thread Safety:
        - Timer-based execution ensures single aggregator thread
        - Worker metrics use internal locks for snapshot operations
        - No shared mutable state between collection cycles
    
    Lifecycle:
        logger = InferenceMetricLogger(pipeline, ...)
        logger.start()
        # ... runs in background ...
        logger.stop(timeout=10)
    """
    
    def __init__(
        self,
        streaming_pipeline: Any,  # StreamingPipeline reference
        interval_seconds: float = 60.0,  # 2 minutes default
        publisher: Optional[MetricPublisher] = None,
        deployment_id: Optional[str] = None,
        deployment_instance_id: Optional[str] = None,
        app_deploy_id: Optional[str] = None,
        action_id: Optional[str] = None,
        app_id: Optional[str] = None,
        log_metrics_every_n: int = 5  # Log detailed metrics every N collections
    ):
        """
        Initialize metric logger.
        
        Args:
            streaming_pipeline: Reference to StreamingPipeline instance
            interval_seconds: Reporting interval (INFERENCE_METRIC_LOGGING_INTERVAL)
            publisher: MetricPublisher implementation (defaults to Kafka)
            deployment_id: Deployment identifier for metric log
            deployment_instance_id: Deployment instance identifier for metric log
            app_deploy_id: App deployment identifier
            action_id: Action identifier
            app_id: Application identifier
            log_metrics_every_n: Log detailed metrics every N collections (0 = never)
        """
        self.pipeline = streaming_pipeline
        self.interval_seconds = interval_seconds
        self.publisher = publisher
        self.log_metrics_every_n = log_metrics_every_n
        
        # NOTE : Opt 1: Use `action_details` but API call overhead
        # NOTE : Opt 2: Extract these params from `server.py` via streaming_pipeline
        
        # Metric log identifiers
        self.deployment_id = deployment_id 
        self.deployment_instance_id = deployment_instance_id
        self.app_deploy_id = app_deploy_id
        self.action_id = action_id
        self.app_id = app_id

        # State
        self._running = False
        self._timer: Optional[threading.Timer] = None
        self._lock = threading.Lock()
        self._last_collection_time = time.time()
        
        # Statistics
        self._total_collections = 0
        self._failed_collections = 0
        self._failed_publishes = 0
        self._skipped_publishes = 0  # Count of skipped publishes due to no data
        
        logger.info(
            f"Initialized InferenceMetricLogger: "
            f"interval={interval_seconds}s, "
            f"log_metrics_every_n={log_metrics_every_n}, "
            f"deployment_id={self.deployment_id}, "
            f"deployment_instance_id={self.deployment_instance_id}"
        )
    
    def start(self) -> None:
        """
        Start the background metric collection loop.
        
        Spawns a timer-based thread that wakes every interval_seconds
        to collect and publish metrics.
        
        Thread Safety:
            Uses lock to prevent multiple start calls from creating
            duplicate timer threads.
        """
        with self._lock:
            if self._running:
                logger.warning("Metric logger already running")
                return
            
            self._running = True
            self._last_collection_time = time.time()
            
            # Initialize default publisher if none provided
            if self.publisher is None:
                try:
                    self.publisher = KafkaMetricPublisher()
                except Exception as e:
                    logger.warning(
                        f"Failed to initialize Kafka publisher: {e}. "
                        f"Using NoOp publisher."
                    )
                    self.publisher = NoOpMetricPublisher()
            
            # Start timer loop
            self._schedule_next_collection()
            
            logger.info("Metric logger started")
    
    def stop(self, timeout: float = 10.0) -> None:
        """
        Stop the background metric collection loop.
        
        Args:
            timeout: Maximum time to wait for final collection (seconds)
        
        Note:
            Performs one final collection before stopping to avoid
            losing metrics from the last interval.
        """
        with self._lock:
            if not self._running:
                return
            
            self._running = False
            
            # Cancel pending timer
            if self._timer:
                self._timer.cancel()
                self._timer = None
        
        # Final collection outside lock to avoid deadlock
        try:
            logger.info("Performing final metric collection...")
            self._collect_and_publish()
        except Exception as e:
            logger.error(f"Error during final collection: {e}")
        
        # Close publisher
        if self.publisher:
            try:
                self.publisher.close()
            except Exception as e:
                logger.error(f"Error closing publisher: {e}")
        
        logger.info(
            f"Metric logger stopped. "
            f"Collections: {self._total_collections}, "
            f"Skipped (no data): {self._skipped_publishes}, "
            f"Failed: {self._failed_collections}, "
            f"Publish failures: {self._failed_publishes}"
        )
    
    def wait(self, timeout: Optional[float] = None) -> None:
        """
        Wait for the metric logger to stop.
        
        Args:
            timeout: Maximum time to wait (None = wait indefinitely)
        
        Note:
            This is a passive wait - use stop() to actually stop the logger.
        """
        start_time = time.time()
        
        while self._running:
            if timeout and (time.time() - start_time) >= timeout:
                logger.warning(f"Wait timeout after {timeout}s")
                break
            time.sleep(0.1)
    
    def _schedule_next_collection(self) -> None:
        """Schedule the next metric collection using threading.Timer."""
        if not self._running:
            return
        
        self._timer = threading.Timer(
            self.interval_seconds,
            self._timer_callback
        )
        self._timer.daemon = False  # Non-daemon for graceful shutdown
        self._timer.start()
    
    def _timer_callback(self) -> None:
        """Timer callback that performs collection and reschedules."""
        try:
            self._collect_and_publish()
        except Exception as e:
            logger.error(f"Error in metric collection: {e}", exc_info=True)
            self._failed_collections += 1
        finally:
            # Schedule next collection
            self._schedule_next_collection()
    
    def _collect_and_publish(self) -> None:
        """
        Collect metrics from all workers, aggregate, and publish.
        
        This method:
        1. Determines interval boundaries
        2. Collects snapshots from all active workers
        3. Aggregates by worker_type
        4. Builds InferenceMetricLog
        5. Publishes via configured publisher
        6. Resets worker metrics for next interval
        """
        interval_end = time.time()
        interval_start = self._last_collection_time
        self._last_collection_time = interval_end
        
        try:
            # Collect snapshots from all workers
            snapshots = self._collect_worker_snapshots(interval_start, interval_end)

            # Aggregate by worker type (empty snapshots will result in all inactive metrics)
            aggregated = self._aggregate_by_type(snapshots)

            # Check if there's any actual data to publish
            has_data = self._has_metric_data(aggregated)

            # ALWAYS publish metrics, even when all workers are inactive (zero metrics)
            # This ensures backend visibility into deployment instance health
            if not has_data:
                self._skipped_publishes += 1
                logger.debug(
                    f"Publishing zero metrics (all workers inactive) for interval={self.interval_seconds}s "
                    f"(total inactive intervals: {self._skipped_publishes})"
                )
                # Continue to build and publish metric_log (don't return early)

            # Build metric log
            metric_log = self._build_metric_log(aggregated, interval_end)
            
            # Log detailed metrics info every N collections
            if self.log_metrics_every_n > 0 and (self._total_collections + 1) % self.log_metrics_every_n == 0:
                self._log_detailed_metrics(aggregated, metric_log)
            
            # Publish
            success = self.publisher.publish(metric_log)
            
            if success:
                self._total_collections += 1
                logger.info(
                    f"Published metrics: interval={self.interval_seconds}s, "
                    f"workers={len(snapshots)}, "
                    f"timestamp={metric_log['timestamp']}"
                )
            else:
                self._failed_publishes += 1
                logger.error("Failed to publish metric log")
                
        except Exception as e:
            logger.error(f"Error collecting/publishing metrics: {e}", exc_info=True)
            self._failed_collections += 1
    
    def _collect_worker_snapshots(self, interval_start: float, interval_end: float) -> List[MetricSnapshot]:
        """Collect snapshots from shared metrics instances (optimized for refactored architecture)."""
        snapshots = []

        # Collect once per worker type (not per worker instance)
        # All workers of the same type share a single WorkerMetrics instance
        worker_types_collected = set()

        try:
            # Consumer metrics (AsyncConsumerManager)
            if self.pipeline.consumer_manager and "consumer" not in worker_types_collected:
                try:
                    snapshot = self.pipeline.consumer_manager.metrics.snapshot_and_reset(
                        interval_start, interval_end
                    )
                    snapshots.append(snapshot)
                    worker_types_collected.add("consumer")
                except Exception as e:
                    logger.warning(f"Error collecting consumer metrics: {e}")

            # Inference workers (multiprocessing pool with shared metrics)
            if self.pipeline.inference_pool and "inference" not in worker_types_collected:
                try:
                    # Access shared metrics directly from WorkerMetrics class
                    metrics = WorkerMetrics.get_shared("inference")
                    snapshot = metrics.snapshot_and_reset(interval_start, interval_end)
                    snapshots.append(snapshot)
                    worker_types_collected.add("inference")
                except Exception as e:
                    logger.warning(f"Error collecting inference metrics: {e}")

            # Post-processing workers (multiprocessing pool with shared metrics)
            if self.pipeline.postproc_pool and "post_processing" not in worker_types_collected:
                try:
                    # Access shared metrics directly from WorkerMetrics class
                    metrics = WorkerMetrics.get_shared("post_processing")
                    snapshot = metrics.snapshot_and_reset(interval_start, interval_end)
                    snapshots.append(snapshot)
                    worker_types_collected.add("post_processing")
                except Exception as e:
                    logger.warning(f"Error collecting post-processing metrics: {e}")

            # Producer workers (list of threading-based workers with shared metrics)
            if self.pipeline.producer_workers and "producer" not in worker_types_collected:
                try:
                    worker = self.pipeline.producer_workers[0]  # Any worker, they share metrics
                    snapshot = worker.metrics.snapshot_and_reset(interval_start, interval_end)
                    snapshots.append(snapshot)
                    worker_types_collected.add("producer")
                except Exception as e:
                    logger.warning(f"Error collecting producer metrics: {e}")

        except Exception as e:
            logger.error(f"Error accessing pipeline workers: {e}", exc_info=True)

        return snapshots
    
    def _aggregate_by_type(
        self,
        snapshots: List[MetricSnapshot]
    ) -> Dict[str, Dict[str, Any]]:
        """
        Aggregate snapshots by worker_type.
        
        Combines metrics from multiple workers of the same type
        (e.g., multiple consumer instances) into single aggregated stats.
        
        Args:
            snapshots: List of MetricSnapshot objects
        
        Returns:
            Dictionary mapping worker_type to aggregated statistics
        """
        # Group snapshots by type
        by_type: Dict[str, List[MetricSnapshot]] = {
            "consumer": [],
            "inference": [],
            "post_processing": [],
            "producer": []
        }
        
        for snapshot in snapshots:
            if snapshot.worker_type in by_type:
                by_type[snapshot.worker_type].append(snapshot)
        
        # Aggregate each type
        aggregated = {}
        
        for worker_type, type_snapshots in by_type.items():
            if type_snapshots:
                aggregated[worker_type] = self._aggregate_snapshots(type_snapshots)
            else:
                # No active workers of this type - use inactive metrics
                aggregated[worker_type] = self._inactive_metrics()
        
        return aggregated
    
    def _aggregate_snapshots(self, snapshots: List[MetricSnapshot]) -> Dict[str, Any]:
        """
        Aggregate multiple snapshots into combined latency and throughput statistics.
        """
        if not snapshots:
            return self._inactive_metrics()

        all_latency_samples = []
        throughput_rates = []
        any_active = False

        # Default fallback for zero/negative intervals
        FALLBACK_INTERVAL = 1.0

        total_throughput = 0

        for snapshot in snapshots:
            # Compute per-snapshot interval with validation
            interval_seconds = snapshot.interval_end_ts - snapshot.interval_start_ts
            if interval_seconds <= 0:
                logger.warning(
                    f"Invalid interval for worker {snapshot.worker_id}: "
                    f"start={snapshot.interval_start_ts}, end={snapshot.interval_end_ts}, "
                    f"using fallback {FALLBACK_INTERVAL}s"
                )
                interval_seconds = FALLBACK_INTERVAL

            # Compute and record per-worker rate
            worker_rate = snapshot.throughput_count / interval_seconds
            throughput_rates.append(worker_rate)

            # Aggregate latency and throughput data
            all_latency_samples.extend(snapshot.latency_samples)
            total_throughput += snapshot.throughput_count
            any_active = any_active or snapshot.was_active

        # Use modular stat calculators
        latency_stats = self._compute_latency_stats(all_latency_samples)
        throughput_stats = self._compute_throughput_stats(throughput_rates)

        result = {
            "active": any_active,
            "total_throughput": total_throughput
        }
        
        # Only include stats if we have data
        if latency_stats:
            result["latency"] = latency_stats
        if throughput_stats:
            result["throughput"] = throughput_stats
        
        return result

    
    def _compute_latency_stats(self, samples: List[float]) -> Dict[str, Any]:
        """
        Compute latency statistics from samples.
        
        Args:
            samples: List of latency measurements in milliseconds
        
        Returns:
            Dictionary with min, max, avg, p0, p50, p100, unit
            Returns empty dict if no samples
        """
        if not samples:
            return {}
        
        sorted_samples = sorted(samples)
        n = len(sorted_samples)
        
        return {
            "min": sorted_samples[0],
            "max": sorted_samples[-1],
            "avg": sum(samples) / n,
            "p0": sorted_samples[0],
            "p50": WorkerMetrics._percentile(sorted_samples, 50),
            "p100": sorted_samples[-1],
            "unit": "ms"
        }
    
    def _compute_throughput_stats(self, rates: List[float]) -> Dict[str, Any]:
        """
        Compute throughput statistics from per-worker rates.

        Args:
            rates: List of per-worker throughput rates (msg/sec)

        Returns:
            Dictionary with min, max, avg, p0, p50, p100, unit
            Returns empty dict if no rates
        """
        if not rates:
            return {}

        sorted_rates = sorted(rates)
        n = len(sorted_rates)

        return {
            "min": sorted_rates[0],
            "max": sorted_rates[-1],
            "avg": sum(sorted_rates) / n,
            "p0": sorted_rates[0],
            "p50": WorkerMetrics._percentile(sorted_rates, 50),
            "p100": sorted_rates[-1],
            "unit": "msg/sec"
        }

    
    def _inactive_metrics(self) -> Dict[str, Any]:
        """
        Generate metrics structure for inactive worker type.
        
        Returns:
            Dictionary with only active=False, no metric data
        """
        return {
            "active": False
        }
    
    def _has_metric_data(self, aggregated: Dict[str, Dict[str, Any]]) -> bool:
        """
        Check if aggregated metrics contain any actual data.
        
        Args:
            aggregated: Aggregated metrics by worker_type
        
        Returns:
            True if any worker has latency or throughput data, False otherwise
        """
        for worker_type, metrics in aggregated.items():
            # Check if worker has any actual metric data (not just active flag)
            if "latency" in metrics or "throughput" in metrics:
                return True
        return False
    
    def _log_detailed_metrics(self, aggregated: Dict[str, Dict[str, Any]], metric_log: Dict[str, Any]) -> None:
        """
        Log detailed metrics information for debugging and monitoring.
        
        Args:
            aggregated: Aggregated metrics by worker_type
            metric_log: The complete metric log being published
        """
        logger.info("=" * 80)
        logger.info(f"DETAILED METRICS REPORT (Collection #{self._total_collections + 1})")
        logger.info(f"Timestamp: {metric_log['timestamp']}")
        logger.info(f"Deployment: {self.deployment_id} / Instance: {self.deployment_instance_id}")
        logger.info("-" * 80)
        
        for worker_type in ["consumer", "inference", "post_processing", "producer"]:
            metrics = aggregated.get(worker_type, {})
            
            if not metrics or not metrics.get("active"):
                logger.info(f"{worker_type.upper()}: INACTIVE (no data)")
                continue
            
            logger.info(f"{worker_type.upper()}:")
            
            # Latency metrics
            if "latency" in metrics:
                lat = metrics["latency"]
                logger.info(
                    f"  Latency: "
                    f"min={lat.get('min', 0):.2f}{lat.get('unit', 'ms')}, "
                    f"avg={lat.get('avg', 0):.2f}{lat.get('unit', 'ms')}, "
                    f"p50={lat.get('p50', 0):.2f}{lat.get('unit', 'ms')}, "
                    f"max={lat.get('max', 0):.2f}{lat.get('unit', 'ms')}"
                )
            else:
                logger.info("  Latency: No data")
            
            # Throughput metrics
            if "throughput" in metrics:
                tput = metrics["throughput"]
                logger.info(
                    f"  Throughput: "
                    f"min={tput.get('min', 0):.2f}{tput.get('unit', '')}, "
                    f"avg={tput.get('avg', 0):.2f}{tput.get('unit', '')}, "
                    f"p50={tput.get('p50', 0):.2f}{tput.get('unit', '')}, "
                    f"max={tput.get('max', 0):.2f}{tput.get('unit', '')}"
                )
            else:
                logger.info("  Throughput: No data")
            
            # Total throughput if available
            if "total_throughput" in metrics:
                logger.info(f"  Total Items: {metrics['total_throughput']}")
        
        logger.info("-" * 80)
        logger.info(
            f"Stats: Total={self._total_collections + 1}, "
            f"Skipped={self._skipped_publishes}, "
            f"Failed={self._failed_collections}, "
            f"Publish Failures={self._failed_publishes}"
        )
        logger.info("=" * 80)
    
    def _build_metric_log(
        self,
        aggregated: Dict[str, Dict[str, Any]],
        timestamp: float
    ) -> Dict[str, Any]:
        """
        Build InferenceMetricLog matching required schema.
        
        Args:
            aggregated: Aggregated metrics by worker_type
            timestamp: Collection timestamp (Unix epoch)
        
        Returns:
            InferenceMetricLog dictionary ready for publishing
            Only includes worker types that have data
        """
        # Convert timestamp to ISO8601 UTC
        dt = datetime.fromtimestamp(timestamp, tz=timezone.utc)
        iso_timestamp = dt.strftime("%Y-%m-%dT%H:%M:%SZ")
        
        # Build metrics dict, only including workers with data
        metrics = {}
        for worker_type in ["consumer", "inference", "post_processing", "producer"]:
            worker_metrics = aggregated.get(worker_type, self._inactive_metrics())
            # Only include if active or has actual metric data
            if worker_metrics.get("active") or "latency" in worker_metrics or "throughput" in worker_metrics:
                metrics[worker_type] = worker_metrics
        
        return {
            "deployment_id": self.deployment_id,
            "deployment_instance_id": self.deployment_instance_id,
            "app_deploy_id": self.app_deploy_id,
            "action_id": self.action_id,
            "app_id": self.app_id,
            "timestamp": iso_timestamp,
            "metrics": metrics
        }
    
    def get_stats(self) -> Dict[str, Any]:
        """
        Get logger statistics.
        
        Returns:
            Dictionary with collection statistics
        """
        return {
            "running": self._running,
            "interval_seconds": self.interval_seconds,
            "log_metrics_every_n": self.log_metrics_every_n,
            "total_collections": self._total_collections,
            "skipped_publishes": self._skipped_publishes,
            "failed_collections": self._failed_collections,
            "failed_publishes": self._failed_publishes,
            "last_collection_time": self._last_collection_time
        }