"""Background AWS pricing update tasks."""

import asyncio
from datetime import datetime

from celery import shared_task

from src.core.celery_app import celery_app
from src.core.database import AsyncSessionLocal
from src.core.cache import cache_manager
from src.core.logging_config import get_logger, set_correlation_id
from src.core.monitoring import metrics


logger = get_logger(__name__)


@celery_app.task(
    bind=True,
    max_retries=3,
    default_retry_delay=3600,  # 1 hour
    time_limit=1800,  # 30 minutes
)
def update_aws_pricing(self):
    """Update AWS pricing data from AWS Pricing API.

    This task fetches the latest AWS pricing information and updates
    the local cache and database.
    """
    correlation_id = set_correlation_id()
    start_time = datetime.utcnow()

    logger.info("Starting AWS pricing update", extra={"correlation_id": correlation_id})

    try:
        # Run update
        updated_count = asyncio.run(_update_pricing_async())

        duration = (datetime.utcnow() - start_time).total_seconds()

        metrics.increment_counter("aws_pricing_updates_total")

        logger.info(
            "AWS pricing update completed",
            extra={
                "updated_count": updated_count,
                "duration_seconds": duration,
            },
        )

        return {
            "status": "completed",
            "updated_count": updated_count,
            "duration_seconds": duration,
        }

    except Exception as exc:
        logger.exception("AWS pricing update failed")
        raise self.retry(exc=exc, countdown=3600)


async def _update_pricing_async() -> int:
    """Async helper for pricing update."""
    async with AsyncSessionLocal() as db:
        from src.services.cost_calculator import cost_calculator

        try:
            # Initialize cache
            await cache_manager.initialize()

            # Update pricing for different services
            updated_count = 0

            # This would typically fetch from AWS Pricing API
            # For now, we'll just clear the pricing cache to force refresh
            cleared = await cache_manager.delete_pattern("l3:pricing:*")

            logger.info(f"Cleared {cleared} cached pricing entries")

            # Pre-warm cache with common queries
            services = ["sqs", "lambda", "bedrock"]
            regions = ["us-east-1", "us-west-2", "eu-west-1", "ap-southeast-1"]

            for service in services:
                for region in regions:
                    try:
                        # Warm cache for each service/region combination
                        if service == "sqs":
                            await cost_calculator.calculate_sqs_cost(db, 1, region)
                        elif service == "lambda":
                            await cost_calculator.calculate_lambda_cost(
                                db, 1, 1.0, region
                            )
                        elif service == "bedrock":
                            await cost_calculator.calculate_bedrock_cost(
                                db, 1000, 0, region
                            )

                        updated_count += 1
                    except Exception as e:
                        logger.warning(
                            f"Failed to warm cache for {service}/{region}",
                            extra={"error": str(e)},
                        )

            await db.commit()

            return updated_count

        except Exception as e:
            await db.rollback()
            raise


@celery_app.task(
    bind=True,
    time_limit=300,
)
def warm_pricing_cache(self, services: list[str] = None, regions: list[str] = None):
    """Pre-warm the pricing cache for common services and regions.

    Args:
        services: List of services to warm (default: all)
        regions: List of regions to warm (default: common ones)
    """
    correlation_id = set_correlation_id()

    services = services or ["sqs", "lambda", "bedrock"]
    regions = regions or ["us-east-1", "us-west-2", "eu-west-1"]

    logger.info(
        "Warming pricing cache",
        extra={
            "services": services,
            "regions": regions,
            "correlation_id": correlation_id,
        },
    )

    try:
        warmed_count = asyncio.run(_warm_cache_async(services, regions))

        logger.info(
            "Pricing cache warming completed",
            extra={"warmed_count": warmed_count},
        )

        return {"status": "completed", "warmed_count": warmed_count}

    except Exception as exc:
        logger.exception("Cache warming failed")
        raise self.retry(exc=exc, countdown=300)


async def _warm_cache_async(services: list[str], regions: list[str]) -> int:
    """Async helper for cache warming."""
    async with AsyncSessionLocal() as db:
        from src.services.cost_calculator import cost_calculator

        await cache_manager.initialize()

        warmed_count = 0

        for service in services:
            for region in regions:
                try:
                    cache_key = f"{service}:{region}"

                    # Calculate pricing (this will cache the result)
                    if service == "sqs":
                        await cost_calculator.calculate_sqs_cost(db, 1, region)
                    elif service == "lambda":
                        await cost_calculator.calculate_lambda_cost(db, 1, 1.0, region)
                    elif service == "bedrock":
                        await cost_calculator.calculate_bedrock_cost(
                            db, 1000, 0, region
                        )

                    warmed_count += 1

                except Exception as e:
                    logger.warning(
                        f"Failed to warm cache for {service}/{region}",
                        extra={"error": str(e)},
                    )

        return warmed_count
