# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/services/metrics_cleanup_service.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

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
"""

# 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
