# -*- coding: utf-8 -*-
"""Metrics Cleanup Service for automatic deletion of old metrics.

This service provides automatic and manual cleanup of old metrics data to prevent
unbounded table growth and maintain query performance.

Features:
- Batched deletion to prevent long locks
- Configurable retention period
- Background task for periodic cleanup
- Manual cleanup trigger via admin API
- Per-table cleanup with statistics

Copyright 2025
SPDX-License-Identifier: Apache-2.0
"""

# Standard
import asyncio
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
import logging
import time
from typing import Dict, Optional

# Third-Party
from sqlalchemy import delete, func, select
from sqlalchemy.orm import Session

# First-Party
from mcpgateway.config import settings
from mcpgateway.db import (
    A2AAgentMetric,
    A2AAgentMetricsHourly,
    fresh_db_session,
    PromptMetric,
    PromptMetricsHourly,
    ResourceMetric,
    ResourceMetricsHourly,
    ServerMetric,
    ServerMetricsHourly,
    ToolMetric,
    ToolMetricsHourly,
)
from mcpgateway.services.metrics_rollup_service import get_metrics_rollup_service_if_initialized

logger = logging.getLogger(__name__)


def delete_metrics_in_batches(db: Session, model_class, filter_column, entity_id, batch_size: Optional[int] = None) -> int:
    """Delete metrics rows for a specific entity in batches within the current transaction.

    Args:
        db: Database session.
        model_class: SQLAlchemy model to delete from.
        filter_column: Column used to filter by entity_id.
        entity_id: Entity identifier to delete metrics for.
        batch_size: Optional batch size override.

    Returns:
        int: Total rows deleted.
    """
    effective_batch_size = batch_size or getattr(settings, "metrics_cleanup_batch_size", 10000)
    total_deleted = 0

    while True:
        subq = select(model_class.id).where(filter_column == entity_id).limit(effective_batch_size)
        delete_stmt = delete(model_class).where(model_class.id.in_(subq))
        result = db.execute(delete_stmt)

        rowcount = result.rowcount
        batch_deleted = rowcount if isinstance(rowcount, int) else 0
        total_deleted += batch_deleted

        if batch_deleted <= 0 or batch_deleted < effective_batch_size:
            break

    return total_deleted


@contextmanager
def pause_rollup_during_purge(reason: str = "purge_metrics"):
    """Pause rollup task while purging metrics to reduce race conditions.

    Args:
        reason: Reason for pausing the rollup task.

    Yields:
        None
    """
    rollup_service = get_metrics_rollup_service_if_initialized()
    if rollup_service:
        rollup_service.pause(reason)
    try:
        yield
    finally:
        if rollup_service:
            rollup_service.resume()


@dataclass
class CleanupResult:
    """Result of a cleanup operation."""

    table_name: str
    deleted_count: int
    remaining_count: int
    cutoff_date: datetime
    duration_seconds: float
    error: Optional[str] = None


@dataclass
class CleanupSummary:
    """Summary of all cleanup operations."""

    total_deleted: int
    tables: Dict[str, CleanupResult]
    duration_seconds: float
    started_at: datetime
    completed_at: datetime


class MetricsCleanupService:
    """Service for cleaning up old metrics data.

    This service provides:
    - Batched deletion of old metrics (prevents long locks)
    - Configurable retention period per table type
    - Background task for periodic cleanup
    - Manual cleanup trigger
    - Cleanup statistics and reporting

    Configuration (via environment variables):
    - METRICS_CLEANUP_ENABLED: Enable automatic cleanup (default: True)
    - METRICS_RETENTION_DAYS: Days to retain raw metrics (default: 7)
    - METRICS_CLEANUP_INTERVAL_HOURS: Hours between cleanup runs (default: 1)
    - METRICS_CLEANUP_BATCH_SIZE: Batch size for deletion (default: 10000)
    - METRICS_DELETE_RAW_AFTER_ROLLUP: Delete raw after rollup exists (default: True)
    - METRICS_DELETE_RAW_AFTER_ROLLUP_HOURS: Hours after which to delete if rollup exists (default: 1)
    """

    # Map of raw metric tables to their hourly rollup counterparts
    METRIC_TABLES = [
        ("tool_metrics", ToolMetric, ToolMetricsHourly, "tool_id"),
        ("resource_metrics", ResourceMetric, ResourceMetricsHourly, "resource_id"),
        ("prompt_metrics", PromptMetric, PromptMetricsHourly, "prompt_id"),
        ("server_metrics", ServerMetric, ServerMetricsHourly, "server_id"),
        ("a2a_agent_metrics", A2AAgentMetric, A2AAgentMetricsHourly, "a2a_agent_id"),
    ]

    def __init__(
        self,
        retention_days: Optional[int] = None,
        batch_size: Optional[int] = None,
        cleanup_interval_hours: Optional[int] = None,
        enabled: Optional[bool] = None,
    ):
        """Initialize the metrics cleanup service.

        Args:
            retention_days: Days to retain raw metrics (default: from settings or 7)
            batch_size: Batch size for deletion (default: from settings or 10000)
            cleanup_interval_hours: Hours between cleanup runs (default: from settings or 1)
            enabled: Whether cleanup is enabled (default: from settings or True)
        """
        self.retention_days = retention_days or getattr(settings, "metrics_retention_days", 7)
        self.batch_size = batch_size or getattr(settings, "metrics_cleanup_batch_size", 10000)
        self.cleanup_interval_hours = cleanup_interval_hours or getattr(settings, "metrics_cleanup_interval_hours", 1)
        self.enabled = enabled if enabled is not None else getattr(settings, "metrics_cleanup_enabled", True)

        # Raw deletion after rollup is handled by MetricsRollupService; stored for stats/reporting.
        self.delete_raw_after_rollup = getattr(settings, "metrics_delete_raw_after_rollup", True)
        self.delete_raw_after_rollup_hours = getattr(settings, "metrics_delete_raw_after_rollup_hours", 1)

        # Rollup retention
        self.rollup_retention_days = getattr(settings, "metrics_rollup_retention_days", 365)

        # Background task
        self._cleanup_task: Optional[asyncio.Task] = None
        self._shutdown_event = asyncio.Event()

        # Stats
        self._total_cleaned = 0
        self._cleanup_runs = 0

        logger.info(
            f"MetricsCleanupService initialized: enabled={self.enabled}, " f"retention_days={self.retention_days}, batch_size={self.batch_size}, " f"interval_hours={self.cleanup_interval_hours}"
        )

    async def start(self) -> None:
        """Start the background cleanup task."""
        if not self.enabled:
            logger.info("MetricsCleanupService disabled, skipping start")
            return

        if self._cleanup_task is None or self._cleanup_task.done():
            self._shutdown_event.clear()
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())
            logger.info("MetricsCleanupService background task started")

    async def shutdown(self) -> None:
        """Shutdown the cleanup service."""
        logger.info("MetricsCleanupService shutting down...")

        # Signal shutdown
        self._shutdown_event.set()

        # Cancel the cleanup task
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass

        logger.info(f"MetricsCleanupService shutdown complete: " f"total_cleaned={self._total_cleaned}, cleanup_runs={self._cleanup_runs}")

    async def _cleanup_loop(self) -> None:
        """Background task that periodically cleans up old metrics.

        Raises:
            asyncio.CancelledError: When the task is cancelled during shutdown.
        """
        logger.info(f"Metrics cleanup loop started (interval={self.cleanup_interval_hours}h)")

        # Calculate interval in seconds
        interval_seconds = self.cleanup_interval_hours * 3600

        while not self._shutdown_event.is_set():
            try:
                # Wait for interval or shutdown
                try:
                    await asyncio.wait_for(
                        self._shutdown_event.wait(),
                        timeout=interval_seconds,
                    )
                    # Shutdown signaled
                    break
                except asyncio.TimeoutError:
                    # Normal timeout, proceed to cleanup
                    pass

                # Run cleanup
                summary = await self.cleanup_all()
                self._cleanup_runs += 1
                self._total_cleaned += summary.total_deleted

                if summary.total_deleted > 0:
                    logger.info(f"Metrics cleanup #{self._cleanup_runs}: deleted {summary.total_deleted} records " f"in {summary.duration_seconds:.2f}s")

            except asyncio.CancelledError:
                logger.debug("Cleanup loop cancelled")
                raise
            except Exception as e:
                logger.error(f"Error in metrics cleanup loop: {e}", exc_info=True)
                # Continue the loop despite errors
                await asyncio.sleep(60)

    async def cleanup_all(
        self,
        retention_days: Optional[int] = None,
        include_rollup: bool = True,
    ) -> CleanupSummary:
        """Clean up all old metrics across all tables.

        Args:
            retention_days: Override retention period (optional). Use 0 to delete all.
            include_rollup: Also clean up old rollup tables

        Returns:
            CleanupSummary: Summary of cleanup operations
        """
        started_at = datetime.now(timezone.utc)
        start_time = time.monotonic()

        # Use provided retention_days if set (including 0), otherwise use default
        retention = retention_days if retention_days is not None else self.retention_days
        # For retention=0 (delete all), use current time without hour-alignment to delete everything
        # Otherwise, hour-align cutoff to match query service's get_retention_cutoff() and prevent gaps
        if retention == 0:
            cutoff = datetime.now(timezone.utc) + timedelta(hours=1)  # Future time ensures all data is deleted
        else:
            cutoff = (datetime.now(timezone.utc) - timedelta(days=retention)).replace(minute=0, second=0, microsecond=0)

        results: Dict[str, CleanupResult] = {}
        total_deleted = 0

        # Clean up raw metrics tables
        for table_name, model_class, _, _ in self.METRIC_TABLES:
            result = await asyncio.to_thread(
                self._cleanup_table,
                model_class,
                table_name,
                cutoff,
            )
            results[table_name] = result
            total_deleted += result.deleted_count

        # Clean up rollup tables if enabled
        if include_rollup:
            # If retention_days is explicitly 0 (delete all), also delete all rollups
            # Otherwise use the rollup retention period
            if retention_days is not None and retention_days == 0:
                rollup_cutoff = datetime.now(timezone.utc)  # Delete all rollups too
            else:
                rollup_cutoff = datetime.now(timezone.utc) - timedelta(days=self.rollup_retention_days)
            for table_name, _, hourly_model_class, _ in self.METRIC_TABLES:
                hourly_table_name = f"{table_name}_hourly"
                result = await asyncio.to_thread(
                    self._cleanup_table,
                    hourly_model_class,
                    hourly_table_name,
                    rollup_cutoff,
                    timestamp_column="hour_start",
                )
                results[hourly_table_name] = result
                total_deleted += result.deleted_count

        duration = time.monotonic() - start_time
        completed_at = datetime.now(timezone.utc)

        return CleanupSummary(
            total_deleted=total_deleted,
            tables=results,
            duration_seconds=duration,
            started_at=started_at,
            completed_at=completed_at,
        )

    def _cleanup_table(
        self,
        model_class,
        table_name: str,
        cutoff: datetime,
        timestamp_column: str = "timestamp",
    ) -> CleanupResult:
        """Clean up old records from a single table using batched deletion.

        Args:
            model_class: SQLAlchemy model class
            table_name: Table name for logging
            cutoff: Delete records older than this timestamp
            timestamp_column: Name of the timestamp column

        Returns:
            CleanupResult: Result of the cleanup operation
        """
        start_time = time.monotonic()
        total_deleted = 0
        error_msg = None

        try:
            with fresh_db_session() as db:
                timestamp_col = getattr(model_class, timestamp_column)

                # Delete in batches to prevent long locks
                while True:
                    # Get batch of IDs to delete
                    stmt = select(model_class.id).where(timestamp_col < cutoff).limit(self.batch_size)
                    ids_to_delete = [row[0] for row in db.execute(stmt).fetchall()]

                    if not ids_to_delete:
                        break

                    # Delete batch
                    delete_stmt = delete(model_class).where(model_class.id.in_(ids_to_delete))
                    result = db.execute(delete_stmt)
                    db.commit()

                    batch_deleted = result.rowcount
                    total_deleted += batch_deleted

                    logger.debug(f"Cleaned {batch_deleted} records from {table_name}")

                    # If we deleted less than batch size, we're done
                    if batch_deleted < self.batch_size:
                        break

                # Get remaining count
                remaining_count = db.execute(select(func.count()).select_from(model_class)).scalar() or 0  # pylint: disable=not-callable

        except Exception as e:
            logger.error(f"Error cleaning up {table_name}: {e}")
            error_msg = str(e)
            remaining_count = -1

        duration = time.monotonic() - start_time

        if total_deleted > 0:
            logger.info(f"Cleaned {total_deleted} records from {table_name} (cutoff: {cutoff})")

        return CleanupResult(
            table_name=table_name,
            deleted_count=total_deleted,
            remaining_count=remaining_count,
            cutoff_date=cutoff,
            duration_seconds=duration,
            error=error_msg,
        )

    async def cleanup_table(
        self,
        table_type: str,
        retention_days: Optional[int] = None,
    ) -> CleanupResult:
        """Clean up old records from a specific table.

        Args:
            table_type: One of 'tool', 'resource', 'prompt', 'server', 'a2a_agent'
            retention_days: Override retention period (optional). Use 0 to delete all.

        Returns:
            CleanupResult: Result of the cleanup operation

        Raises:
            ValueError: If table_type is not recognized
        """
        table_map = {
            "tool": ("tool_metrics", ToolMetric),
            "resource": ("resource_metrics", ResourceMetric),
            "prompt": ("prompt_metrics", PromptMetric),
            "server": ("server_metrics", ServerMetric),
            "a2a_agent": ("a2a_agent_metrics", A2AAgentMetric),
        }

        if table_type not in table_map:
            raise ValueError(f"Unknown table type: {table_type}. Must be one of: {list(table_map.keys())}")

        table_name, model_class = table_map[table_type]
        # Use provided retention_days if set (including 0), otherwise use default
        retention = retention_days if retention_days is not None else self.retention_days
        # For retention=0 (delete all), use current time without hour-alignment to delete everything
        # Otherwise, hour-align cutoff to match query service's get_retention_cutoff() and prevent gaps
        if retention == 0:
            cutoff = datetime.now(timezone.utc) + timedelta(hours=1)  # Future time ensures all data is deleted
        else:
            cutoff = (datetime.now(timezone.utc) - timedelta(days=retention)).replace(minute=0, second=0, microsecond=0)

        return await asyncio.to_thread(
            self._cleanup_table,
            model_class,
            table_name,
            cutoff,
        )

    def get_stats(self) -> dict:
        """Get cleanup service statistics.

        Returns:
            dict: Cleanup statistics
        """
        return {
            "enabled": self.enabled,
            "retention_days": self.retention_days,
            "rollup_retention_days": self.rollup_retention_days,
            "batch_size": self.batch_size,
            "cleanup_interval_hours": self.cleanup_interval_hours,
            "total_cleaned": self._total_cleaned,
            "cleanup_runs": self._cleanup_runs,
            "delete_raw_after_rollup": self.delete_raw_after_rollup,
            "delete_raw_after_rollup_hours": self.delete_raw_after_rollup_hours,
        }

    async def get_table_sizes(self) -> Dict[str, int]:
        """Get the current size of all metrics tables.

        Returns:
            Dict[str, int]: Table name to row count mapping
        """
        sizes = {}

        def _get_sizes() -> Dict[str, int]:
            with fresh_db_session() as db:
                for table_name, model_class, hourly_class, _ in self.METRIC_TABLES:
                    # pylint: disable=not-callable
                    sizes[table_name] = db.execute(select(func.count()).select_from(model_class)).scalar() or 0
                    hourly_name = f"{table_name}_hourly"
                    sizes[hourly_name] = db.execute(select(func.count()).select_from(hourly_class)).scalar() or 0
            return sizes

        return await asyncio.to_thread(_get_sizes)


# Singleton instance
_metrics_cleanup_service: Optional[MetricsCleanupService] = None


def get_metrics_cleanup_service() -> MetricsCleanupService:
    """Get or create the singleton MetricsCleanupService instance.

    Returns:
        MetricsCleanupService: The singleton cleanup service instance
    """
    global _metrics_cleanup_service  # pylint: disable=global-statement
    if _metrics_cleanup_service is None:
        _metrics_cleanup_service = MetricsCleanupService()
    return _metrics_cleanup_service
