import asyncio
from datetime import datetime, timedelta, timezone
from typing import Optional

from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
    SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
    SPEND_LOG_CLEANUP_BATCH_SIZE,
    SPEND_LOG_CLEANUP_JOB_NAME,
    SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES,
    SPEND_LOG_RUN_LOOPS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
    SpendLogsPartitionManager,
)
from litellm.proxy.utils import PrismaClient


class SpendLogCleanup:
    """
    Handles cleaning up old spend logs based on maximum retention period.

    When LiteLLM_SpendLogs is range-partitioned, expired data is reclaimed by
    dropping whole partitions (instant, frees disk immediately). Otherwise it
    falls back to deleting logs in batches.
    Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
    """

    def __init__(
        self,
        general_settings=None,
        redis_cache: Optional[RedisCache] = None,
        partition_manager: Optional[SpendLogsPartitionManager] = None,
    ):
        self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE
        self.retention_seconds: Optional[int] = None
        self.partition_manager = partition_manager or SpendLogsPartitionManager()
        from litellm.proxy.proxy_server import general_settings as default_settings

        self.general_settings = general_settings or default_settings
        from litellm.proxy.proxy_server import proxy_logging_obj

        pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
        self.pod_lock_manager = pod_lock_manager
        verbose_proxy_logger.info(
            f"SpendLogCleanup initialized with batch size: {self.batch_size}"
        )

    def _should_delete_spend_logs(self) -> bool:
        """
        Determines if logs should be deleted based on the max retention period in settings.
        """
        retention_setting = self.general_settings.get(
            "maximum_spend_logs_retention_period"
        )
        verbose_proxy_logger.info(f"Checking retention setting: {retention_setting}")

        if retention_setting is None:
            verbose_proxy_logger.info("No retention setting found")
            return False

        try:
            if isinstance(retention_setting, int):
                verbose_proxy_logger.warning(
                    f"maximum_spend_logs_retention_period is an integer ({retention_setting}); treating as days. "
                    "Use a string like '3d' to be explicit."
                )
                retention_setting = f"{retention_setting}d"
            self.retention_seconds = duration_in_seconds(retention_setting)
            verbose_proxy_logger.info(
                f"Retention period set to {self.retention_seconds} seconds"
            )
            return True
        except ValueError as e:
            verbose_proxy_logger.warning(
                f"Invalid maximum_spend_logs_retention_period value: {retention_setting}, error: {str(e)}"
            )
            return False

    async def _delete_old_logs(
        self, prisma_client: PrismaClient, cutoff_date: datetime
    ) -> int:
        """
        Helper method to delete old logs in batches.
        Returns the total number of logs deleted.
        """
        total_deleted = 0
        run_count = 0
        consecutive_failures = 0
        while True:
            if run_count > SPEND_LOG_RUN_LOOPS:
                verbose_proxy_logger.info(
                    "Max logs deleted - 1,00,000, rest of the logs will be deleted in next run"
                )
                break
            # Step 1: Find logs and delete them in one go without fetching to application
            # Delete in batches, limited by self.batch_size
            try:
                deleted_result = await prisma_client.db.execute_raw(
                    """
                    DELETE FROM "LiteLLM_SpendLogs"
                    WHERE ("request_id", "startTime") IN (
                        SELECT "request_id", "startTime" FROM "LiteLLM_SpendLogs"
                        WHERE "startTime" < $1::timestamptz
                        LIMIT $2
                    )
                    """,
                    cutoff_date,
                    self.batch_size,
                )
            except Exception as batch_exc:
                # A single batch failure (e.g. Prisma/DB timeout) must not abort
                # the whole run — subsequent batches may still succeed.
                consecutive_failures += 1
                verbose_proxy_logger.exception(
                    "Spend log cleanup batch failed "
                    "(run_count=%d, consecutive_failures=%d, batch_size=%d, "
                    "cutoff=%s, total_deleted_so_far=%d): %s: %s",
                    run_count,
                    consecutive_failures,
                    self.batch_size,
                    cutoff_date.isoformat(),
                    total_deleted,
                    type(batch_exc).__name__,
                    batch_exc,
                )
                if (
                    consecutive_failures
                    >= SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES
                ):
                    verbose_proxy_logger.error(
                        "Aborting spend log cleanup after %d consecutive batch "
                        "failures; total deleted before abort: %d",
                        consecutive_failures,
                        total_deleted,
                    )
                    break
                await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
                continue

            consecutive_failures = 0

            deleted_count = 0
            if isinstance(deleted_result, int):
                deleted_count = deleted_result
            else:
                verbose_proxy_logger.error(
                    f"Unexpected execute_raw return type for spend log cleanup: {type(deleted_result)}; "
                    "aborting cleanup to avoid infinite loop"
                )
                break

            verbose_proxy_logger.info(f"Deleted {deleted_count} logs in this batch")

            if deleted_count == 0:
                verbose_proxy_logger.info(
                    f"No more logs to delete. Total deleted: {total_deleted}"
                )
                break

            total_deleted += deleted_count
            run_count += 1

            # Add a small sleep to prevent overwhelming the database
            await asyncio.sleep(0.1)

        return total_deleted

    async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None:
        """
        Main cleanup function. Deletes old spend logs in batches.
        If pod_lock_manager is available, ensures only one pod runs cleanup.
        If no pod_lock_manager, runs cleanup without distributed locking.
        """
        lock_acquired = False
        try:
            verbose_proxy_logger.info(f"Cleanup job triggered at {datetime.now()}")

            if not self._should_delete_spend_logs():
                return

            if self.retention_seconds is None:
                verbose_proxy_logger.error(
                    "Retention seconds is None, cannot proceed with cleanup"
                )
                return

            # If we have a pod lock manager, try to acquire the lock
            if self.pod_lock_manager and self.pod_lock_manager.redis_cache:
                lock_acquired = (
                    await self.pod_lock_manager.acquire_lock(
                        cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME,
                    )
                    or False
                )
                verbose_proxy_logger.info(
                    f"Lock acquisition attempt: {'successful' if lock_acquired else 'failed'}  at {datetime.now()}"
                )

                if not lock_acquired:
                    verbose_proxy_logger.info("Another pod is already running cleanup")
                    return

            cutoff_date = datetime.now(timezone.utc) - timedelta(
                seconds=float(self.retention_seconds)
            )
            verbose_proxy_logger.info(
                f"Removing logs older than {cutoff_date.isoformat()}"
            )

            if self.general_settings.get(
                "use_spend_logs_partitioning", False
            ) and await self.partition_manager.is_partitioned(prisma_client):
                await self.partition_manager.ensure_partitions(prisma_client)
                dropped = await self.partition_manager.drop_partitions_older_than(
                    prisma_client, cutoff_date
                )
                verbose_proxy_logger.info(
                    "Dropped %d expired spend-log partitions: %s",
                    len(dropped),
                    dropped,
                )
                # DROP only reclaims whole expired partitions. Expired rows can
                # still sit in the DEFAULT partition (backfill, coverage gaps)
                # or in a partition that spans the cutoff, so retention must
                # also delete those stragglers row-wise.
                total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
                verbose_proxy_logger.info(
                    f"Deleted {total_deleted} expired logs not covered by dropped partitions"
                )
            else:
                total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
                verbose_proxy_logger.info(f"Deleted {total_deleted} logs")

        except Exception as e:
            # .exception() captures the traceback; str(e) alone on a Prisma/DB
            # timeout is often empty and gives operators no signal to diagnose.
            verbose_proxy_logger.exception(
                "Error during spend log cleanup: %s: %s",
                type(e).__name__,
                e,
            )
            return  # Return after error handling
        finally:
            # Only release the lock if it was actually acquired
            if (
                lock_acquired
                and self.pod_lock_manager
                and self.pod_lock_manager.redis_cache
            ):
                await self.pod_lock_manager.release_lock(
                    cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME
                )
                verbose_proxy_logger.info("Released cleanup lock")
