"""
Payment Processing Worker Tasks
Handles payment retries, refunds, and subscription management
"""
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 app.workers.celery_app import celery
from app.database import get_sync_session
from app.models.billing import Invoice, Payment, InvoiceStatus, PaymentStatus, PaymentMethod
from app.models.subscription import Subscription, SubscriptionStatus
from app.models.client import Client
from app.config import settings

logger = logging.getLogger(__name__)


# Retry schedule in seconds: 1 hour, 1 day, 3 days, 7 days
RETRY_SCHEDULE = [3600, 86400, 259200, 604800]
MAX_RETRY_ATTEMPTS = 4


@celery.task(
    bind=True,
    name="app.workers.tasks.payment_tasks.retry_payment",
    max_retries=0,  # Manual retry scheduling
    queue="invoices"
)
def retry_payment(self, invoice_id: int) -> Dict[str, Any]:
    """
    Retry a failed payment for an invoice

    Args:
        invoice_id: Invoice ID to retry payment for
    """
    try:
        logger.info(f"Retrying payment for invoice {invoice_id}")

        with get_sync_session() as session:
            invoice = session.execute(
                select(Invoice).where(Invoice.id == invoice_id)
            ).scalar_one_or_none()

            if not invoice:
                return {"success": False, "error": "Invoice not found"}

            if invoice.status == InvoiceStatus.PAID:
                return {"success": False, "error": "Invoice already paid"}

            # Count previous attempts
            attempt_count = session.execute(
                select(func.count(Payment.id)).where(
                    and_(
                        Payment.invoice_id == invoice_id,
                        Payment.status == PaymentStatus.FAILED
                    )
                )
            ).scalar() or 0

            if attempt_count >= MAX_RETRY_ATTEMPTS:
                logger.warning(f"Max retry attempts reached for invoice {invoice_id}")
                _handle_payment_failure(session, invoice)
                return {
                    "success": False,
                    "error": "Max retry attempts reached",
                    "action": "subscription_suspended"
                }

            # Get client subscription
            client = session.execute(
                select(Client).where(Client.id == invoice.client_id)
            ).scalar_one_or_none()

            if not client:
                return {"success": False, "error": "Client not found"}

            subscription = session.execute(
                select(Subscription).where(
                    and_(
                        Subscription.client_id == client.id,
                        Subscription.status.in_([
                            SubscriptionStatus.ACTIVE,
                            SubscriptionStatus.PAUSED
                        ])
                    )
                )
            ).scalar_one_or_none()

            if not subscription:
                return {"success": False, "error": "No active subscription"}

            # Attempt payment with available gateway
            result = None
            if subscription.razorpay_subscription_id:
                result = _retry_razorpay_payment(session, invoice, subscription)
            elif subscription.stripe_subscription_id:
                result = _retry_stripe_payment(session, invoice, subscription)
            else:
                # No stored payment method, send reminder
                _send_payment_reminder(invoice, attempt_count + 1)
                return {
                    "success": False,
                    "error": "No stored payment method",
                    "action": "reminder_sent"
                }

            if result and result.get("success"):
                return result
            else:
                # Schedule next retry
                if attempt_count + 1 < MAX_RETRY_ATTEMPTS:
                    next_delay = RETRY_SCHEDULE[attempt_count]
                    retry_payment.apply_async(
                        args=[invoice_id],
                        countdown=next_delay
                    )
                    logger.info(f"Scheduled retry {attempt_count + 2} for invoice {invoice_id} in {next_delay}s")

                return {
                    "success": False,
                    "error": result.get("error") if result else "Payment failed",
                    "next_retry_in": RETRY_SCHEDULE[attempt_count] if attempt_count + 1 < MAX_RETRY_ATTEMPTS else None
                }

    except Exception as e:
        logger.error(f"Error retrying payment for invoice {invoice_id}: {str(e)}")
        return {"success": False, "error": str(e)}


@celery.task(
    bind=True,
    name="app.workers.tasks.payment_tasks.process_failed_payments",
    max_retries=2,
    queue="invoices"
)
def process_failed_payments(self) -> Dict[str, Any]:
    """
    Process all invoices with failed payments that need retry
    """
    try:
        logger.info("Processing failed payments")

        with get_sync_session() as session:
            # Find overdue invoices that haven't exceeded max retries
            overdue_invoices = session.execute(
                select(Invoice).where(
                    and_(
                        Invoice.status.in_([InvoiceStatus.PENDING, InvoiceStatus.OVERDUE]),
                        Invoice.due_date < datetime.utcnow()
                    )
                )
            ).scalars().all()

            results = {
                "invoices_processed": 0,
                "retries_scheduled": 0,
                "failures_handled": 0,
                "errors": []
            }

            for invoice in overdue_invoices:
                try:
                    # Count failed payment attempts
                    attempt_count = session.execute(
                        select(func.count(Payment.id)).where(
                            and_(
                                Payment.invoice_id == invoice.id,
                                Payment.status == PaymentStatus.FAILED
                            )
                        )
                    ).scalar() or 0

                    # Get last failed payment
                    last_payment = session.execute(
                        select(Payment).where(
                            and_(
                                Payment.invoice_id == invoice.id,
                                Payment.status == PaymentStatus.FAILED
                            )
                        ).order_by(Payment.created_at.desc())
                    ).scalar_one_or_none()

                    # Check if enough time has passed for retry
                    if last_payment and attempt_count < MAX_RETRY_ATTEMPTS:
                        retry_delay = RETRY_SCHEDULE[min(attempt_count, len(RETRY_SCHEDULE) - 1)]
                        next_retry_time = last_payment.created_at + timedelta(seconds=retry_delay)

                        if datetime.utcnow() >= next_retry_time:
                            retry_payment.delay(invoice.id)
                            results["retries_scheduled"] += 1
                    elif attempt_count >= MAX_RETRY_ATTEMPTS:
                        _handle_payment_failure(session, invoice)
                        results["failures_handled"] += 1
                    elif not last_payment:
                        # No payment attempt yet, schedule first retry
                        retry_payment.delay(invoice.id)
                        results["retries_scheduled"] += 1

                    results["invoices_processed"] += 1

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

            logger.info(f"Failed payments processed: {results}")
            return results

    except Exception as e:
        logger.error(f"Error processing failed payments: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.payment_tasks.process_refund",
    max_retries=3,
    queue="invoices"
)
def process_refund(
    self,
    payment_id: int,
    amount: Optional[float] = None,
    reason: str = "Customer request"
) -> Dict[str, Any]:
    """
    Process a refund for a payment

    Args:
        payment_id: Payment ID to refund
        amount: Amount to refund (default: full amount)
        reason: Refund reason
    """
    try:
        logger.info(f"Processing refund for payment {payment_id}")

        with get_sync_session() as session:
            payment = session.execute(
                select(Payment).where(Payment.id == payment_id)
            ).scalar_one_or_none()

            if not payment:
                return {"success": False, "error": "Payment not found"}

            if payment.status != PaymentStatus.SUCCESS:
                return {"success": False, "error": "Payment is not successful"}

            refund_amount = Decimal(str(amount)) if amount else payment.amount

            if refund_amount > payment.amount - (payment.refund_amount or Decimal("0")):
                return {"success": False, "error": "Refund amount exceeds available amount"}

            # Process refund with gateway
            result = None
            if payment.method == PaymentMethod.RAZORPAY:
                result = _process_razorpay_refund(payment, refund_amount, reason)
            elif payment.method == PaymentMethod.STRIPE:
                result = _process_stripe_refund(payment, refund_amount, reason)
            else:
                # Manual refund
                result = {"success": True, "refund_id": f"manual_{payment_id}"}

            if result and result.get("success"):
                # Update payment record
                payment.refund_amount = (payment.refund_amount or Decimal("0")) + refund_amount
                payment.refund_reason = reason
                payment.refunded_at = datetime.utcnow()

                if payment.refund_amount >= payment.amount:
                    payment.status = PaymentStatus.REFUNDED

                # Update invoice if fully refunded
                if payment.status == PaymentStatus.REFUNDED:
                    invoice = payment.invoice
                    invoice.status = InvoiceStatus.REFUNDED

                session.commit()

                logger.info(f"Refund processed for payment {payment_id}: ₹{refund_amount}")

                # Send refund notification
                _send_refund_notification(payment, refund_amount)

                return {
                    "success": True,
                    "refund_id": result.get("refund_id"),
                    "amount": float(refund_amount)
                }
            else:
                return {
                    "success": False,
                    "error": result.get("error") if result else "Refund failed"
                }

    except Exception as e:
        logger.error(f"Error processing refund for payment {payment_id}: {str(e)}")
        self.retry(exc=e)


@celery.task(
    name="app.workers.tasks.payment_tasks.suspend_subscription",
    queue="invoices"
)
def suspend_subscription(client_id: int, reason: str = "Payment failure") -> Dict[str, Any]:
    """
    Suspend a client's subscription due to payment failure

    Args:
        client_id: Client ID
        reason: Suspension reason
    """
    try:
        logger.info(f"Suspending subscription for client {client_id}")

        with get_sync_session() as session:
            subscription = session.execute(
                select(Subscription).where(
                    and_(
                        Subscription.client_id == client_id,
                        Subscription.status == SubscriptionStatus.ACTIVE
                    )
                )
            ).scalar_one_or_none()

            if not subscription:
                return {"success": False, "error": "No active subscription"}

            subscription.status = SubscriptionStatus.PAUSED
            subscription.paused_at = datetime.utcnow()
            session.commit()

            # Send suspension notification
            from app.workers.tasks.email_tasks import send_templated_email
            client = session.execute(
                select(Client).where(Client.id == client_id)
            ).scalar_one()

            send_templated_email.delay(
                template_name="subscription_suspended",
                recipient_email=client.email,
                recipient_name=client.company_name,
                context={
                    "client_name": client.company_name,
                    "reason": reason,
                    "reactivation_url": "https://portal.ai360.com/billing"
                }
            )

            logger.info(f"Subscription suspended for client {client_id}")

            return {"success": True, "status": "paused"}

    except Exception as e:
        logger.error(f"Error suspending subscription: {str(e)}")
        return {"success": False, "error": str(e)}


@celery.task(
    name="app.workers.tasks.payment_tasks.reactivate_subscription",
    queue="invoices"
)
def reactivate_subscription(client_id: int) -> Dict[str, Any]:
    """
    Reactivate a paused subscription after payment

    Args:
        client_id: Client ID
    """
    try:
        logger.info(f"Reactivating subscription for client {client_id}")

        with get_sync_session() as session:
            subscription = session.execute(
                select(Subscription).where(
                    and_(
                        Subscription.client_id == client_id,
                        Subscription.status == SubscriptionStatus.PAUSED
                    )
                )
            ).scalar_one_or_none()

            if not subscription:
                return {"success": False, "error": "No paused subscription"}

            subscription.status = SubscriptionStatus.ACTIVE
            subscription.paused_at = None
            session.commit()

            # Send reactivation notification
            from app.workers.tasks.email_tasks import send_templated_email
            client = session.execute(
                select(Client).where(Client.id == client_id)
            ).scalar_one()

            send_templated_email.delay(
                template_name="subscription_reactivated",
                recipient_email=client.email,
                recipient_name=client.company_name,
                context={
                    "client_name": client.company_name,
                    "dashboard_url": "https://portal.ai360.com/dashboard"
                }
            )

            logger.info(f"Subscription reactivated for client {client_id}")

            return {"success": True, "status": "active"}

    except Exception as e:
        logger.error(f"Error reactivating subscription: {str(e)}")
        return {"success": False, "error": str(e)}


# ===========================================
# HELPER FUNCTIONS
# ===========================================

def _retry_razorpay_payment(
    session,
    invoice: Invoice,
    subscription: Subscription
) -> Dict[str, Any]:
    """Retry payment via Razorpay"""
    try:
        import razorpay

        client = razorpay.Client(auth=(
            settings.RAZORPAY_KEY_ID,
            settings.RAZORPAY_KEY_SECRET
        ))

        # Create payment link
        payment_link = client.payment_link.create({
            "amount": int(invoice.total_amount * 100),
            "currency": "INR",
            "accept_partial": False,
            "description": f"Invoice {invoice.invoice_number}",
            "customer": {
                "email": invoice.client.email,
                "name": invoice.client.company_name,
            },
            "notify": {
                "sms": True,
                "email": True
            },
            "reminder_enable": True,
            "notes": {
                "invoice_id": str(invoice.id)
            },
            "callback_url": f"https://api.ai360.com/webhooks/razorpay?invoice_id={invoice.id}",
            "callback_method": "get"
        })

        # Create payment record
        payment = Payment(
            invoice_id=invoice.id,
            amount=invoice.total_amount,
            currency=invoice.currency,
            status=PaymentStatus.PENDING,
            method=PaymentMethod.RAZORPAY,
            gateway_order_id=payment_link["id"],
            gateway_response=payment_link,
            initiated_at=datetime.utcnow()
        )
        session.add(payment)
        session.commit()

        return {
            "success": True,
            "payment_link": payment_link["short_url"],
            "payment_id": payment.id
        }

    except Exception as e:
        logger.error(f"Razorpay retry failed: {str(e)}")
        return {"success": False, "error": str(e)}


def _retry_stripe_payment(
    session,
    invoice: Invoice,
    subscription: Subscription
) -> Dict[str, Any]:
    """Retry payment via Stripe"""
    try:
        import stripe
        stripe.api_key = settings.STRIPE_SECRET_KEY

        # Create payment intent
        intent = stripe.PaymentIntent.create(
            amount=int(invoice.total_amount * 100),
            currency="inr",
            customer=subscription.stripe_subscription_id,
            description=f"Invoice {invoice.invoice_number}",
            metadata={"invoice_id": str(invoice.id)},
            automatic_payment_methods={"enabled": True},
        )

        # Create payment record
        payment = Payment(
            invoice_id=invoice.id,
            amount=invoice.total_amount,
            currency=invoice.currency,
            status=PaymentStatus.PENDING,
            method=PaymentMethod.STRIPE,
            gateway_payment_id=intent.id,
            gateway_response=dict(intent),
            initiated_at=datetime.utcnow()
        )
        session.add(payment)
        session.commit()

        return {
            "success": True,
            "client_secret": intent.client_secret,
            "payment_id": payment.id
        }

    except Exception as e:
        logger.error(f"Stripe retry failed: {str(e)}")
        return {"success": False, "error": str(e)}


def _process_razorpay_refund(
    payment: Payment,
    amount: Decimal,
    reason: str
) -> Dict[str, Any]:
    """Process refund via Razorpay"""
    try:
        import razorpay

        client = razorpay.Client(auth=(
            settings.RAZORPAY_KEY_ID,
            settings.RAZORPAY_KEY_SECRET
        ))

        refund = client.payment.refund(payment.gateway_payment_id, {
            "amount": int(amount * 100),
            "notes": {"reason": reason}
        })

        return {
            "success": True,
            "refund_id": refund["id"]
        }

    except Exception as e:
        return {"success": False, "error": str(e)}


def _process_stripe_refund(
    payment: Payment,
    amount: Decimal,
    reason: str
) -> Dict[str, Any]:
    """Process refund via Stripe"""
    try:
        import stripe
        stripe.api_key = settings.STRIPE_SECRET_KEY

        refund = stripe.Refund.create(
            payment_intent=payment.gateway_payment_id,
            amount=int(amount * 100),
            reason="requested_by_customer"
        )

        return {
            "success": True,
            "refund_id": refund.id
        }

    except Exception as e:
        return {"success": False, "error": str(e)}


def _handle_payment_failure(session, invoice: Invoice) -> None:
    """Handle final payment failure - suspend subscription"""
    logger.warning(f"Handling payment failure for invoice {invoice.id}")

    # Mark invoice as overdue if not already
    if invoice.status != InvoiceStatus.OVERDUE:
        invoice.status = InvoiceStatus.OVERDUE
        session.commit()

    # Suspend subscription
    suspend_subscription.delay(invoice.client_id, "Payment failure - max retries exceeded")


def _send_payment_reminder(invoice: Invoice, attempt_number: int) -> None:
    """Send payment reminder email"""
    from app.workers.tasks.email_tasks import send_templated_email

    send_templated_email.delay(
        template_name="payment_retry_reminder",
        recipient_email=invoice.client.email,
        recipient_name=invoice.client.company_name,
        context={
            "client_name": invoice.client.company_name,
            "invoice_number": invoice.invoice_number,
            "amount": str(invoice.total_amount),
            "currency": invoice.currency,
            "attempt_number": attempt_number,
            "payment_url": f"https://portal.ai360.com/invoices/{invoice.id}/pay"
        }
    )


def _send_refund_notification(payment: Payment, amount: Decimal) -> None:
    """Send refund notification email"""
    from app.workers.tasks.email_tasks import send_templated_email

    invoice = payment.invoice
    send_templated_email.delay(
        template_name="refund_processed",
        recipient_email=invoice.client.email,
        recipient_name=invoice.client.company_name,
        context={
            "client_name": invoice.client.company_name,
            "invoice_number": invoice.invoice_number,
            "refund_amount": str(amount),
            "currency": invoice.currency,
            "original_amount": str(payment.amount),
            "refund_date": datetime.utcnow().strftime("%B %d, %Y")
        }
    )