"""
Usage Aggregation Worker Tasks
Handles usage tracking, aggregation, and limit checking
"""
import logging
from datetime import datetime, timedelta
from decimal import Decimal
from typing import Optional, Dict, Any, List

from celery import shared_task
from sqlalchemy import select, and_, func
from sqlalchemy.orm import selectinload

from app.workers.celery_app import celery
from app.database import get_sync_session
from app.models.billing import UsageLog, UsageSummary, UsageType
from app.models.subscription import Subscription, Plan
from app.models.client import Client
from app.models.conversation import Conversation, Message

logger = logging.getLogger(__name__)


# Token costs per model (cost per 1000 tokens in INR)
MODEL_COSTS = {
    "gpt-4o": {"input": 0.375, "output": 1.125},
    "gpt-4o-mini": {"input": 0.0113, "output": 0.045},
    "gpt-4-turbo": {"input": 0.75, "output": 2.25},
    "gpt-3.5-turbo": {"input": 0.0375, "output": 0.1125},
    "claude-3-opus": {"input": 1.125, "output": 5.625},
    "claude-3-sonnet": {"input": 0.225, "output": 1.125},
    "claude-3-haiku": {"input": 0.01875, "output": 0.09375},
}


@celery.task(
    bind=True,
    name="app.workers.tasks.usage_tasks.log_usage",
    max_retries=3,
    default_retry_delay=60,
    queue="analytics"
)
def log_usage(
    self,
    client_id: int,
    usage_type: str,
    quantity: int = 1,
    model: Optional[str] = None,
    input_tokens: int = 0,
    output_tokens: int = 0,
    conversation_id: Optional[int] = None,
    api_key_id: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Log usage event for a client

    Args:
        client_id: Client ID
        usage_type: Type of usage (chat_message, chat_token, etc.)
        quantity: Quantity of usage
        model: AI model used
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        conversation_id: Related conversation ID
        api_key_id: API key used
        metadata: Additional metadata
    """
    try:
        logger.debug(f"Logging usage for client {client_id}: {usage_type}")

        with get_sync_session() as session:
            # Calculate costs
            total_tokens = input_tokens + output_tokens
            unit_cost = Decimal("0")
            total_cost = Decimal("0")

            if model and model in MODEL_COSTS:
                costs = MODEL_COSTS[model]
                input_cost = Decimal(str(input_tokens)) * Decimal(str(costs["input"])) / 1000
                output_cost = Decimal(str(output_tokens)) * Decimal(str(costs["output"])) / 1000
                total_cost = input_cost + output_cost
                unit_cost = total_cost / max(total_tokens, 1)

            # Get billing period
            billing_period = datetime.utcnow().strftime("%Y-%m")

            # Create usage log
            usage_log = UsageLog(
                client_id=client_id,
                usage_type=UsageType(usage_type),
                quantity=quantity,
                conversation_id=conversation_id,
                api_key_id=api_key_id,
                model=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                unit_cost=unit_cost,
                total_cost=total_cost,
                metadata=metadata,
                billing_period=billing_period
            )

            session.add(usage_log)
            session.commit()

            logger.debug(f"Usage logged: {usage_log.id}")

            return {
                "success": True,
                "usage_log_id": usage_log.id,
                "total_cost": float(total_cost)
            }

    except Exception as e:
        logger.error(f"Error logging usage for client {client_id}: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.usage_tasks.aggregate_daily_usage",
    max_retries=2,
    queue="analytics"
)
def aggregate_daily_usage(self, date: Optional[str] = None) -> Dict[str, Any]:
    """
    Aggregate usage data for a specific day

    Args:
        date: Date to aggregate (YYYY-MM-DD), defaults to yesterday
    """
    try:
        # Default to yesterday
        if date:
            target_date = datetime.strptime(date, "%Y-%m-%d").date()
        else:
            target_date = (datetime.utcnow() - timedelta(days=1)).date()

        date_str = target_date.strftime("%Y-%m-%d")
        logger.info(f"Aggregating daily usage for {date_str}")

        with get_sync_session() as session:
            # Get all clients with usage on this day
            start_time = datetime.combine(target_date, datetime.min.time())
            end_time = datetime.combine(target_date, datetime.max.time())

            # Aggregate by client
            usage_stats = session.execute(
                select(
                    UsageLog.client_id,
                    func.count(UsageLog.id).label("total_logs"),
                    func.sum(UsageLog.quantity).label("total_quantity"),
                    func.sum(UsageLog.total_tokens).label("total_tokens"),
                    func.sum(UsageLog.total_cost).label("total_cost")
                )
                .where(
                    and_(
                        UsageLog.created_at >= start_time,
                        UsageLog.created_at <= end_time
                    )
                )
                .group_by(UsageLog.client_id)
            ).all()

            results = {
                "date": date_str,
                "clients_processed": 0,
                "summaries_created": 0,
                "errors": []
            }

            for stat in usage_stats:
                try:
                    client_id = stat.client_id

                    # Count messages and conversations
                    message_count = session.execute(
                        select(func.count(UsageLog.id))
                        .where(
                            and_(
                                UsageLog.client_id == client_id,
                                UsageLog.usage_type == UsageType.CHAT_MESSAGE,
                                UsageLog.created_at >= start_time,
                                UsageLog.created_at <= end_time
                            )
                        )
                    ).scalar() or 0

                    conversation_count = session.execute(
                        select(func.count(func.distinct(UsageLog.conversation_id)))
                        .where(
                            and_(
                                UsageLog.client_id == client_id,
                                UsageLog.conversation_id.isnot(None),
                                UsageLog.created_at >= start_time,
                                UsageLog.created_at <= end_time
                            )
                        )
                    ).scalar() or 0

                    vision_count = session.execute(
                        select(func.count(UsageLog.id))
                        .where(
                            and_(
                                UsageLog.client_id == client_id,
                                UsageLog.usage_type == UsageType.VISION_REQUEST,
                                UsageLog.created_at >= start_time,
                                UsageLog.created_at <= end_time
                            )
                        )
                    ).scalar() or 0

                    api_call_count = session.execute(
                        select(func.count(UsageLog.id))
                        .where(
                            and_(
                                UsageLog.client_id == client_id,
                                UsageLog.usage_type == UsageType.API_CALL,
                                UsageLog.created_at >= start_time,
                                UsageLog.created_at <= end_time
                            )
                        )
                    ).scalar() or 0

                    # Create or update daily summary
                    existing = session.execute(
                        select(UsageSummary)
                        .where(
                            and_(
                                UsageSummary.client_id == client_id,
                                UsageSummary.period_type == "daily",
                                UsageSummary.period_date == date_str
                            )
                        )
                    ).scalar_one_or_none()

                    if existing:
                        existing.total_messages = message_count
                        existing.total_conversations = conversation_count
                        existing.total_tokens = stat.total_tokens or 0
                        existing.total_vision_requests = vision_count
                        existing.total_api_calls = api_call_count
                        existing.total_cost = stat.total_cost or 0
                    else:
                        summary = UsageSummary(
                            client_id=client_id,
                            period_type="daily",
                            period_date=date_str,
                            total_messages=message_count,
                            total_conversations=conversation_count,
                            total_tokens=stat.total_tokens or 0,
                            total_vision_requests=vision_count,
                            total_api_calls=api_call_count,
                            total_cost=stat.total_cost or 0
                        )
                        session.add(summary)
                        results["summaries_created"] += 1

                    session.commit()
                    results["clients_processed"] += 1

                except Exception as e:
                    results["errors"].append({
                        "client_id": stat.client_id,
                        "error": str(e)
                    })

            logger.info(f"Daily usage aggregation complete: {results}")
            return results

    except Exception as e:
        logger.error(f"Error aggregating daily usage: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.usage_tasks.aggregate_monthly_usage",
    max_retries=2,
    queue="analytics"
)
def aggregate_monthly_usage(self, month: Optional[str] = None) -> Dict[str, Any]:
    """
    Aggregate usage data for a specific month

    Args:
        month: Month to aggregate (YYYY-MM), defaults to previous month
    """
    try:
        # Default to previous month
        if month:
            year, mon = map(int, month.split("-"))
        else:
            today = datetime.utcnow()
            if today.month == 1:
                year, mon = today.year - 1, 12
            else:
                year, mon = today.year, today.month - 1

        month_str = f"{year:04d}-{mon:02d}"
        logger.info(f"Aggregating monthly usage for {month_str}")

        with get_sync_session() as session:
            # Calculate date range
            start_date = datetime(year, mon, 1)
            if mon == 12:
                end_date = datetime(year + 1, 1, 1) - timedelta(seconds=1)
            else:
                end_date = datetime(year, mon + 1, 1) - timedelta(seconds=1)

            # Get all daily summaries for this month
            daily_summaries = session.execute(
                select(
                    UsageSummary.client_id,
                    func.sum(UsageSummary.total_messages).label("total_messages"),
                    func.sum(UsageSummary.total_conversations).label("total_conversations"),
                    func.sum(UsageSummary.total_tokens).label("total_tokens"),
                    func.sum(UsageSummary.total_vision_requests).label("total_vision_requests"),
                    func.sum(UsageSummary.total_api_calls).label("total_api_calls"),
                    func.sum(UsageSummary.total_cost).label("total_cost")
                )
                .where(
                    and_(
                        UsageSummary.period_type == "daily",
                        UsageSummary.period_date >= start_date.strftime("%Y-%m-%d"),
                        UsageSummary.period_date <= end_date.strftime("%Y-%m-%d")
                    )
                )
                .group_by(UsageSummary.client_id)
            ).all()

            results = {
                "month": month_str,
                "clients_processed": 0,
                "summaries_created": 0,
                "errors": []
            }

            for stat in daily_summaries:
                try:
                    client_id = stat.client_id

                    # Create or update monthly summary
                    existing = session.execute(
                        select(UsageSummary)
                        .where(
                            and_(
                                UsageSummary.client_id == client_id,
                                UsageSummary.period_type == "monthly",
                                UsageSummary.period_date == month_str
                            )
                        )
                    ).scalar_one_or_none()

                    if existing:
                        existing.total_messages = stat.total_messages or 0
                        existing.total_conversations = stat.total_conversations or 0
                        existing.total_tokens = stat.total_tokens or 0
                        existing.total_vision_requests = stat.total_vision_requests or 0
                        existing.total_api_calls = stat.total_api_calls or 0
                        existing.total_cost = stat.total_cost or 0
                    else:
                        summary = UsageSummary(
                            client_id=client_id,
                            period_type="monthly",
                            period_date=month_str,
                            total_messages=stat.total_messages or 0,
                            total_conversations=stat.total_conversations or 0,
                            total_tokens=stat.total_tokens or 0,
                            total_vision_requests=stat.total_vision_requests or 0,
                            total_api_calls=stat.total_api_calls or 0,
                            total_cost=stat.total_cost or 0
                        )
                        session.add(summary)
                        results["summaries_created"] += 1

                    session.commit()
                    results["clients_processed"] += 1

                except Exception as e:
                    results["errors"].append({
                        "client_id": stat.client_id,
                        "error": str(e)
                    })

            logger.info(f"Monthly usage aggregation complete: {results}")
            return results

    except Exception as e:
        logger.error(f"Error aggregating monthly usage: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.usage_tasks.check_usage_limits",
    max_retries=2,
    queue="analytics"
)
def check_usage_limits(self) -> Dict[str, Any]:
    """
    Check usage limits for all clients and send notifications
    """
    try:
        logger.info("Checking usage limits")

        with get_sync_session() as session:
            # Get current month
            current_month = datetime.utcnow().strftime("%Y-%m")

            # Get all active clients with subscriptions
            clients = session.execute(
                select(Client)
                .options(selectinload(Client.subscription))
                .where(Client.is_active == True)
            ).scalars().all()

            results = {
                "clients_checked": 0,
                "warnings_sent": 0,
                "limits_exceeded": 0,
                "errors": []
            }

            for client in clients:
                try:
                    if not client.subscription:
                        continue

                    # Get plan limits
                    plan = session.execute(
                        select(Plan)
                        .where(Plan.id == client.subscription.plan_id)
                    ).scalar_one_or_none()

                    if not plan:
                        continue

                    # Get current month usage
                    usage = session.execute(
                        select(UsageSummary)
                        .where(
                            and_(
                                UsageSummary.client_id == client.id,
                                UsageSummary.period_type == "monthly",
                                UsageSummary.period_date == current_month
                            )
                        )
                    ).scalar_one_or_none()

                    if not usage:
                        results["clients_checked"] += 1
                        continue

                    # Check message limits
                    if plan.messages_per_month:
                        usage_percentage = (usage.total_messages / plan.messages_per_month) * 100

                        # 80% warning
                        if 80 <= usage_percentage < 100:
                            _send_usage_warning(
                                client,
                                "messages",
                                usage.total_messages,
                                plan.messages_per_month,
                                usage_percentage
                            )
                            results["warnings_sent"] += 1

                        # 100% exceeded
                        elif usage_percentage >= 100:
                            _send_limit_exceeded_notification(
                                client,
                                "messages",
                                usage.total_messages,
                                plan.messages_per_month
                            )
                            results["limits_exceeded"] += 1

                    results["clients_checked"] += 1

                except Exception as e:
                    results["errors"].append({
                        "client_id": client.id,
                        "error": str(e)
                    })

            logger.info(f"Usage limits check complete: {results}")
            return results

    except Exception as e:
        logger.error(f"Error checking usage limits: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.usage_tasks.generate_daily_reports",
    max_retries=2,
    queue="analytics"
)
def generate_daily_reports(self) -> Dict[str, Any]:
    """
    Generate daily usage reports for clients who have enabled them
    """
    try:
        logger.info("Generating daily reports")

        yesterday = (datetime.utcnow() - timedelta(days=1)).date()
        date_str = yesterday.strftime("%Y-%m-%d")

        with get_sync_session() as session:
            # Get daily summaries for yesterday
            summaries = session.execute(
                select(UsageSummary)
                .options(selectinload(UsageSummary.client))
                .where(
                    and_(
                        UsageSummary.period_type == "daily",
                        UsageSummary.period_date == date_str
                    )
                )
            ).scalars().all()

            results = {
                "date": date_str,
                "reports_generated": 0,
                "emails_sent": 0,
                "errors": []
            }

            for summary in summaries:
                try:
                    client = session.execute(
                        select(Client).where(Client.id == summary.client_id)
                    ).scalar_one_or_none()

                    if not client:
                        continue

                    # Send daily report email
                    from app.workers.tasks.email_tasks import send_templated_email

                    send_templated_email.delay(
                        template_name="daily_usage_report",
                        recipient_email=client.email,
                        recipient_name=client.company_name,
                        context={
                            "date": yesterday.strftime("%B %d, %Y"),
                            "total_messages": summary.total_messages,
                            "total_conversations": summary.total_conversations,
                            "total_tokens": summary.total_tokens,
                            "total_cost": f"₹{summary.total_cost:.2f}",
                            "client_name": client.company_name
                        }
                    )

                    results["reports_generated"] += 1
                    results["emails_sent"] += 1

                except Exception as e:
                    results["errors"].append({
                        "client_id": summary.client_id,
                        "error": str(e)
                    })

            logger.info(f"Daily reports generated: {results}")
            return results

    except Exception as e:
        logger.error(f"Error generating daily reports: {str(e)}")
        self.retry(exc=e)


@celery.task(
    name="app.workers.tasks.usage_tasks.get_usage_stats",
    queue="analytics"
)
def get_usage_stats(
    client_id: int,
    start_date: str,
    end_date: str
) -> Dict[str, Any]:
    """
    Get detailed usage statistics for a client

    Args:
        client_id: Client ID
        start_date: Start date (YYYY-MM-DD)
        end_date: End date (YYYY-MM-DD)
    """
    try:
        with get_sync_session() as session:
            start = datetime.strptime(start_date, "%Y-%m-%d")
            end = datetime.strptime(end_date, "%Y-%m-%d")

            # Get usage logs
            logs = session.execute(
                select(UsageLog)
                .where(
                    and_(
                        UsageLog.client_id == client_id,
                        UsageLog.created_at >= start,
                        UsageLog.created_at <= end
                    )
                )
            ).scalars().all()

            # Calculate statistics
            stats = {
                "period": {"start": start_date, "end": end_date},
                "total_messages": 0,
                "total_tokens": 0,
                "total_cost": Decimal("0"),
                "by_type": {},
                "by_model": {},
                "daily_breakdown": {}
            }

            for log in logs:
                # Totals
                stats["total_tokens"] += log.total_tokens
                stats["total_cost"] += log.total_cost

                # By type
                usage_type = log.usage_type.value
                if usage_type not in stats["by_type"]:
                    stats["by_type"][usage_type] = {"count": 0, "tokens": 0, "cost": Decimal("0")}
                stats["by_type"][usage_type]["count"] += log.quantity
                stats["by_type"][usage_type]["tokens"] += log.total_tokens
                stats["by_type"][usage_type]["cost"] += log.total_cost

                if usage_type == "chat_message":
                    stats["total_messages"] += log.quantity

                # By model
                if log.model:
                    if log.model not in stats["by_model"]:
                        stats["by_model"][log.model] = {"count": 0, "tokens": 0, "cost": Decimal("0")}
                    stats["by_model"][log.model]["count"] += 1
                    stats["by_model"][log.model]["tokens"] += log.total_tokens
                    stats["by_model"][log.model]["cost"] += log.total_cost

                # Daily breakdown
                day = log.created_at.strftime("%Y-%m-%d")
                if day not in stats["daily_breakdown"]:
                    stats["daily_breakdown"][day] = {"messages": 0, "tokens": 0, "cost": Decimal("0")}
                if usage_type == "chat_message":
                    stats["daily_breakdown"][day]["messages"] += log.quantity
                stats["daily_breakdown"][day]["tokens"] += log.total_tokens
                stats["daily_breakdown"][day]["cost"] += log.total_cost

            # Convert Decimal to float for JSON serialization
            stats["total_cost"] = float(stats["total_cost"])
            for key in stats["by_type"]:
                stats["by_type"][key]["cost"] = float(stats["by_type"][key]["cost"])
            for key in stats["by_model"]:
                stats["by_model"][key]["cost"] = float(stats["by_model"][key]["cost"])
            for key in stats["daily_breakdown"]:
                stats["daily_breakdown"][key]["cost"] = float(stats["daily_breakdown"][key]["cost"])

            return stats

    except Exception as e:
        logger.error(f"Error getting usage stats: {str(e)}")
        return {"error": str(e)}


def _send_usage_warning(
    client: Client,
    resource_type: str,
    current_usage: int,
    limit: int,
    percentage: float
):
    """Send usage warning notification"""
    from app.workers.tasks.email_tasks import send_templated_email

    send_templated_email.delay(
        template_name="usage_warning",
        recipient_email=client.email,
        recipient_name=client.company_name,
        context={
            "resource_type": resource_type,
            "current_usage": current_usage,
            "limit": limit,
            "percentage": f"{percentage:.1f}%",
            "client_name": client.company_name,
            "upgrade_url": "https://portal.ai360.com/subscription/upgrade"
        }
    )


def _send_limit_exceeded_notification(
    client: Client,
    resource_type: str,
    current_usage: int,
    limit: int
):
    """Send limit exceeded notification"""
    from app.workers.tasks.email_tasks import send_templated_email

    send_templated_email.delay(
        template_name="limit_exceeded",
        recipient_email=client.email,
        recipient_name=client.company_name,
        context={
            "resource_type": resource_type,
            "current_usage": current_usage,
            "limit": limit,
            "client_name": client.company_name,
            "upgrade_url": "https://portal.ai360.com/subscription/upgrade"
        }
    )