"""Cost calculation service."""

from decimal import Decimal
from typing import Optional
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, and_
from datetime import date

from src.repositories.base import BaseRepository
from src.models.aws_pricing import AwsPricing


class CostCalculator:
    """Service for calculating AWS costs."""

    def __init__(self):
        self.pricing_repo = BaseRepository(AwsPricing)

    async def get_pricing(
        self, db: AsyncSession, service: str, region: str, tier: str
    ) -> Optional[Decimal]:
        """Get active pricing for a service/region/tier."""
        result = await db.execute(
            select(AwsPricing).where(
                and_(
                    AwsPricing.service == service,
                    AwsPricing.region == region,
                    AwsPricing.tier == tier,
                    AwsPricing.is_active == True,
                )
            )
        )
        pricing = result.scalar_one_or_none()
        return pricing.price_per_unit if pricing else None

    async def calculate_sqs_cost(
        self, db: AsyncSession, blocks: int, region: str
    ) -> Decimal:
        """Calculate SQS cost."""
        price = await self.get_pricing(db, "sqs", region, "standard")
        if price is None:
            price = Decimal("0.40")

        # Formula: blocks * price_per_million / 1,000,000
        return Decimal(blocks) * price / Decimal("1000000")

    async def calculate_lambda_cost(
        self, db: AsyncSession, invocations: int, gb_seconds: float, region: str
    ) -> Decimal:
        """Calculate Lambda cost (requests + compute)."""
        request_price = await self.get_pricing(db, "lambda", region, "x86_request")
        compute_price = await self.get_pricing(db, "lambda", region, "x86_compute")

        if request_price is None:
            request_price = Decimal("0.20")
        if compute_price is None:
            compute_price = Decimal("0.0000166667")

        request_cost = Decimal(invocations) * request_price / Decimal("1000000")
        compute_cost = Decimal(str(gb_seconds)) * compute_price
        return request_cost + compute_cost

    async def calculate_bedrock_cost(
        self,
        db: AsyncSession,
        input_tokens: int,
        output_tokens: int,
        region: str,
        model: str = "claude_3_sonnet",
    ) -> Decimal:
        """Calculate Bedrock LLM cost."""
        input_price = await self.get_pricing(db, "bedrock", region, f"{model}_input")
        output_price = await self.get_pricing(db, "bedrock", region, f"{model}_output")

        if input_price is None:
            input_price = Decimal("0.003")
        if output_price is None:
            output_price = Decimal("0.015")

        input_cost = Decimal(input_tokens) * input_price / Decimal("1000")
        output_cost = Decimal(output_tokens) * output_price / Decimal("1000")
        return input_cost + output_cost


# Singleton instance
cost_calculator = CostCalculator()
