"""
Usage history aggregation service.

This service creates monthly usage snapshots by aggregating data from:
- EnrichmentUsage: Enrichment operations and credits consumed
- UsageEvent: Generic billing events (matches, dedupes, API calls)
- UsageQuotaModel: Current usage state
- Job: Completed jobs and rows processed

Snapshots are stored in the usage_history table for trend analysis
and customer intelligence reporting.
"""

from __future__ import annotations

import logging
from datetime import datetime
from decimal import Decimal
from typing import Any, Dict, List, Optional

from sqlalchemy import cast, func, select, Text
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import (
    Account,
    EnrichmentUsage,
    Job,
    UsageEvent,
    UsageEventType,
    UsageHistory,
    UsageQuotaModel,
)

logger = logging.getLogger(__name__)


class UsageHistoryService:
    """Service for aggregating and managing usage history snapshots."""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def create_snapshot(
        self,
        account_id: str,
        period: str,
        force_refresh: bool = False,
    ) -> UsageHistory:
        """
        Create or update a usage snapshot for the given account and period.

        Args:
            account_id: Account to snapshot
            period: Billing period in YYYY-MM format
            force_refresh: If True, recalculate even if snapshot exists

        Returns:
            UsageHistory snapshot record
        """
        # Check if snapshot already exists
        stmt = select(UsageHistory).where(
            UsageHistory.account_id == account_id,
            UsageHistory.period == period,
        )
        result = await self.db.execute(stmt)
        existing = result.scalar_one_or_none()

        if existing and not force_refresh:
            logger.info(f"Snapshot already exists for {account_id} {period}")
            return existing

        # Calculate period boundaries
        year, month = period.split("-")
        period_start = datetime(int(year), int(month), 1)
        if int(month) == 12:
            period_end = datetime(int(year) + 1, 1, 1)
        else:
            period_end = datetime(int(year), int(month) + 1, 1)

        # Aggregate enrichment usage
        enrichment_stats = await self._aggregate_enrichment_usage(
            account_id, period_start, period_end
        )

        # Aggregate usage events (matches, dedupes, jobs)
        event_stats = await self._aggregate_usage_events(
            account_id, period_start, period_end
        )

        # Get account metadata
        account_metadata = await self._get_account_metadata(account_id)

        # Build operation breakdown
        operation_breakdown = {
            "enrichment": {
                "scan": enrichment_stats.get("scan_count", 0),
                "preview": enrichment_stats.get("preview_count", 0),
                "apply": enrichment_stats.get("apply_count", 0),
            },
            "matching": {
                "matches": event_stats.get("match_count", 0),
                "dedupes": event_stats.get("dedupe_count", 0),
            },
            "jobs": {
                "completed": event_stats.get("job_count", 0),
                "failed": event_stats.get("job_failed_count", 0),
            },
        }

        # Create or update snapshot
        snapshot_data = {
            "account_id": account_id,
            "period": period,
            "rows_processed": event_stats.get("total_rows_processed", 0),
            "matches_found": event_stats.get("total_matches_found", 0),
            "dedupes_run": event_stats.get("dedupe_count", 0),
            "enrichments_performed": enrichment_stats.get("total_enrichments", 0),
            "enrichment_credits_consumed": enrichment_stats.get(
                "total_credits_consumed", 0
            ),
            "records_enriched": enrichment_stats.get("total_records_enriched", 0),
            "api_calls": enrichment_stats.get("total_api_calls", 0),
            "jobs_completed": event_stats.get("job_count", 0),
            "cost_usd": enrichment_stats.get("total_cost_usd"),
            "price_usd": enrichment_stats.get("total_price_usd"),
            "margin_usd": enrichment_stats.get("total_margin_usd"),
            "operation_breakdown": operation_breakdown,
            "tier_at_snapshot": account_metadata.get("tier"),
            "subscription_status": account_metadata.get("subscription_status"),
        }

        if existing:
            # Update existing snapshot
            for key, value in snapshot_data.items():
                if key not in ("account_id", "period"):  # Don't update keys
                    setattr(existing, key, value)
            existing.updated_at = datetime.utcnow()
            self.db.add(existing)
            await self.db.commit()
            await self.db.refresh(existing)
            logger.info(f"Updated snapshot for {account_id} {period}")
            return existing
        else:
            # Create new snapshot
            snapshot = UsageHistory(**snapshot_data)
            self.db.add(snapshot)
            await self.db.commit()
            await self.db.refresh(snapshot)
            logger.info(f"Created snapshot for {account_id} {period}")
            return snapshot

    async def get_history(
        self,
        account_id: str,
        months: int = 6,
        ensure_snapshots: bool = False,
    ) -> List[UsageHistory]:
        """
        Get usage history for the last N months.

        Args:
            account_id: Account to retrieve history for
            months: Number of months to retrieve (default 6)
            ensure_snapshots: If True, create missing snapshots for recent months

        Returns:
            List of UsageHistory records, ordered by period descending
        """
        # Calculate period range using proper month arithmetic
        now = datetime.utcnow()
        periods = []
        for i in range(months):
            # Proper month subtraction
            year = now.year
            month = now.month - i
            while month <= 0:
                month += 12
                year -= 1
            period = f"{year:04d}-{month:02d}"
            periods.append(period)

        # Ensure snapshots exist for recent months if requested
        if ensure_snapshots:
            for period in periods[:3]:  # Only create for last 3 months
                try:
                    await self.create_snapshot(account_id, period, force_refresh=False)
                except Exception as e:
                    logger.warning(
                        f"Failed to create snapshot for {account_id} {period}: {e}"
                    )

        # Retrieve snapshots
        stmt = (
            select(UsageHistory)
            .where(
                UsageHistory.account_id == account_id,
                UsageHistory.period.in_(periods),
            )
            .order_by(UsageHistory.period.desc())
        )
        result = await self.db.execute(stmt)
        return list(result.scalars().all())

    async def _aggregate_enrichment_usage(
        self,
        account_id: str,
        period_start: datetime,
        period_end: datetime,
    ) -> Dict[str, Any]:
        """Aggregate enrichment usage for the period."""
        stmt = select(
            func.count(EnrichmentUsage.id).label("total_enrichments"),
            func.sum(EnrichmentUsage.credits_debited).label("total_credits_consumed"),
            func.sum(EnrichmentUsage.records_written).label("total_records_enriched"),
            func.sum(EnrichmentUsage.pdl_requests).label("total_api_calls"),
            func.sum(EnrichmentUsage.cost_usd).label("total_cost_usd"),
            func.sum(EnrichmentUsage.price_usd).label("total_price_usd"),
            func.sum(EnrichmentUsage.margin_usd).label("total_margin_usd"),
            func.count(EnrichmentUsage.id)
            .filter(EnrichmentUsage.operation == "scan")
            .label("scan_count"),
            func.count(EnrichmentUsage.id)
            .filter(EnrichmentUsage.operation == "preview")
            .label("preview_count"),
            func.count(EnrichmentUsage.id)
            .filter(EnrichmentUsage.operation == "apply")
            .label("apply_count"),
        ).where(
            EnrichmentUsage.account_id == account_id,
            EnrichmentUsage.created_at >= period_start,
            EnrichmentUsage.created_at < period_end,
        )

        result = await self.db.execute(stmt)
        row = result.first()

        if not row:
            return {}

        return {
            "total_enrichments": row.total_enrichments or 0,
            "total_credits_consumed": row.total_credits_consumed or 0,
            "total_records_enriched": row.total_records_enriched or 0,
            "total_api_calls": row.total_api_calls or 0,
            "total_cost_usd": row.total_cost_usd or Decimal("0"),
            "total_price_usd": row.total_price_usd or Decimal("0"),
            "total_margin_usd": row.total_margin_usd or Decimal("0"),
            "scan_count": row.scan_count or 0,
            "preview_count": row.preview_count or 0,
            "apply_count": row.apply_count or 0,
        }

    async def _aggregate_usage_events(
        self,
        account_id: str,
        period_start: datetime,
        period_end: datetime,
    ) -> Dict[str, Any]:
        """Aggregate usage events (matches, dedupes, jobs) for the period."""
        # Get event counts and quantities
        stmt = select(
            UsageEvent.event_type,
            func.count(UsageEvent.id).label("event_count"),
            func.sum(UsageEvent.quantity).label("total_quantity"),
        ).where(
            UsageEvent.account_id == account_id,
            UsageEvent.created_at >= period_start,
            UsageEvent.created_at < period_end,
        ).group_by(UsageEvent.event_type)

        result = await self.db.execute(stmt)
        event_data = {}
        total_rows = 0
        total_matches = 0

        for row in result:
            event_type = row.event_type
            count = row.event_count or 0
            quantity = row.total_quantity or 0

            if event_type == UsageEventType.DEDUPE_REQUEST:
                event_data["dedupe_count"] = count
            elif event_type == UsageEventType.MATCH_REQUEST:
                event_data["match_count"] = count
                total_matches += quantity
            elif event_type == UsageEventType.ROWS_PROCESSED:
                total_rows += quantity
            elif event_type == UsageEventType.JOB_COMPLETED:
                event_data["job_count"] = count

        # Get job statistics
        # Note: Job.status is TEXT in DB, not enum, so cast to TEXT for comparison
        job_stmt = select(
            func.count(Job.id).label("total_jobs"),
            func.sum(Job.matches_found).label("total_results"),
        ).where(
            Job.account_id == account_id,
            Job.updated_at >= period_start,
            Job.updated_at < period_end,
            cast(Job.status, Text) == "COMPLETED",  # Cast to TEXT to match DB type
        )

        job_result = await self.db.execute(job_stmt)
        job_row = job_result.first()

        if job_row:
            if "job_count" not in event_data:
                event_data["job_count"] = job_row.total_jobs or 0
            total_matches += job_row.total_results or 0

        return {
            **event_data,
            "total_rows_processed": total_rows,
            "total_matches_found": total_matches,
            "match_count": event_data.get("match_count", 0),
            "dedupe_count": event_data.get("dedupe_count", 0),
            "job_count": event_data.get("job_count", 0),
            "job_failed_count": 0,  # TODO: Add failed job tracking if needed
        }

    async def _get_account_metadata(self, account_id: str) -> Dict[str, Any]:
        """Get account tier and subscription status."""
        # Get current quota for tier information
        quota_stmt = select(UsageQuotaModel).where(
            UsageQuotaModel.account_id == account_id
        )
        quota_result = await self.db.execute(quota_stmt)
        quota = quota_result.scalar_one_or_none()

        # Get account for subscription status
        account_stmt = select(Account).where(Account.id == account_id)
        account_result = await self.db.execute(account_stmt)
        account = account_result.scalar_one_or_none()

        return {
            "tier": quota.tier if quota else None,
            "subscription_status": getattr(account, "subscription_status", None) if account else None,
        }

    async def backfill_snapshots(
        self,
        account_id: Optional[str] = None,
        months: int = 12,
        force_refresh: bool = False,
    ) -> int:
        """
        Backfill usage snapshots for all or specific accounts.

        Args:
            account_id: If provided, only backfill for this account
            months: Number of historical months to backfill
            force_refresh: If True, recalculate existing snapshots

        Returns:
            Number of snapshots created/updated
        """
        # Get accounts to backfill
        if account_id:
            account_ids = [account_id]
        else:
            stmt = select(Account.id)
            result = await self.db.execute(stmt)
            account_ids = [row[0] for row in result.all()]

        # Generate period list using proper month arithmetic
        now = datetime.utcnow()
        periods = []
        for i in range(months):
            # Proper month subtraction
            year = now.year
            month = now.month - i
            while month <= 0:
                month += 12
                year -= 1
            period = f"{year:04d}-{month:02d}"
            periods.append(period)

        # Create snapshots
        snapshot_count = 0
        for acc_id in account_ids:
            for period in periods:
                try:
                    await self.create_snapshot(acc_id, period, force_refresh=force_refresh)
                    snapshot_count += 1
                except Exception as e:
                    logger.error(f"Failed to backfill {acc_id} {period}: {e}")

        logger.info(f"Backfilled {snapshot_count} usage snapshots")
        return snapshot_count
