bakery-ia/services/tenant/app/api/subscription.py

"""
Tenant Service API Endpoints for Subscription and Registration

NEW ARCHITECTURE (SetupIntent-first):
1. /registration-payment-setup - Creates customer + SetupIntent only (NO subscription)
2. Frontend confirms SetupIntent (handles 3DS if needed)
3. /verify-and-complete-registration - Creates subscription AFTER verification

This eliminates duplicate subscriptions by only creating the subscription
after payment verification is complete.
"""

import logging
from typing import Dict, Any
from fastapi import APIRouter, Depends, HTTPException, status, Request, Query
from sqlalchemy.ext.asyncio import AsyncSession
from app.services.subscription_orchestration_service import SubscriptionOrchestrationService
from app.services.coupon_service import CouponService
from app.core.database import get_db
from app.services.registration_state_service import (
    registration_state_service,
    RegistrationStateService,
    RegistrationState
)
from shared.exceptions.payment_exceptions import (
    PaymentServiceError,
    SetupIntentError,
    SubscriptionCreationFailed,
)
from shared.exceptions.registration_exceptions import (
    RegistrationStateError,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/v1/tenants", tags=["tenant"])


async def get_subscription_orchestration_service(
    db: AsyncSession = Depends(get_db)
) -> SubscriptionOrchestrationService:
    """Dependency injection for subscription orchestration service"""
    return SubscriptionOrchestrationService(db)


async def get_registration_state_service() -> RegistrationStateService:
    """Dependency injection for registration state service"""
    return registration_state_service


@router.post("/registration-payment-setup",
            response_model=Dict[str, Any],
            summary="Start registration payment setup")
async def create_registration_payment_setup(
    user_data: Dict[str, Any],
    request: Request,
    orchestration_service: SubscriptionOrchestrationService = Depends(get_subscription_orchestration_service),
    state_service: RegistrationStateService = Depends(get_registration_state_service)
) -> Dict[str, Any]:
    """
    Start registration payment setup (SetupIntent-first architecture).

    NEW ARCHITECTURE: Only creates customer + SetupIntent here.
    NO subscription is created - subscription is created in verify-and-complete-registration.

    Flow:
    1. Create Stripe customer
    2. Create SetupIntent for payment verification
    3. Return SetupIntent to frontend for 3DS handling
    4. Frontend confirms SetupIntent
    5. (Next endpoint) Creates subscription after verification

    Args:
        user_data: User registration data with payment info

    Returns:
        SetupIntent data for frontend confirmation
    """
    state_id = None
    try:
        logger.info("Registration payment setup started",
                   extra={"email": user_data.get('email'), "plan_id": user_data.get('plan_id')})

        # Validate required fields
        if not user_data.get('email'):
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Email is required")

        if not user_data.get('payment_method_id'):
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Payment method ID is required")

        if not user_data.get('plan_id'):
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Plan ID is required")

        # Create registration state
        state_id = await state_service.create_registration_state(
            email=user_data['email'],
            user_data=user_data
        )

        # Create customer + SetupIntent (NO subscription yet!)
        result = await orchestration_service.create_registration_payment_setup(
            user_data=user_data,
            plan_id=user_data.get('plan_id', 'professional'),
            payment_method_id=user_data.get('payment_method_id'),
            billing_interval=user_data.get('billing_cycle', 'monthly'),
            coupon_code=user_data.get('coupon_code')
        )

        # Update state with setup results
        await state_service.update_state_context(state_id, {
            'setup_intent_id': result.get('setup_intent_id'),
            'customer_id': result.get('customer_id'),
            'payment_method_id': result.get('payment_method_id'),
            'plan_id': result.get('plan_id'),
            'billing_interval': result.get('billing_interval'),
            'trial_period_days': result.get('trial_period_days'),
            'coupon_code': result.get('coupon_code')
        })

        await state_service.transition_state(state_id, RegistrationState.PAYMENT_VERIFICATION_PENDING)

        logger.info("Registration payment setup completed",
                   extra={
                       "email": user_data.get('email'),
                       "setup_intent_id": result.get('setup_intent_id'),
                       "requires_action": result.get('requires_action')
                   })

        return {
            "success": True,
            "requires_action": result.get('requires_action', True),
            "action_type": result.get('action_type', 'use_stripe_sdk'),
            "client_secret": result.get('client_secret'),
            "setup_intent_id": result.get('setup_intent_id'),
            "customer_id": result.get('customer_id'),
            "payment_customer_id": result.get('customer_id'),
            "plan_id": result.get('plan_id'),
            "payment_method_id": result.get('payment_method_id'),
            "trial_period_days": result.get('trial_period_days', 0),
            "billing_cycle": result.get('billing_interval'),
            "email": result.get('email'),
            "state_id": state_id,
            "message": result.get('message', 'Payment verification required')
        }

    except PaymentServiceError as e:
        logger.error(f"Payment setup failed: {str(e)}", extra={"email": user_data.get('email')}, exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Payment setup failed: {str(e)}") from e
    except RegistrationStateError as e:
        logger.error(f"Registration state error: {str(e)}", extra={"email": user_data.get('email')}, exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Registration state error: {str(e)}") from e
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}", extra={"email": user_data.get('email')}, exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Registration failed: {str(e)}") from e


@router.post("/verify-and-complete-registration",
            response_model=Dict[str, Any],
            summary="Complete registration after 3DS verification")
async def verify_and_complete_registration(
    verification_data: Dict[str, Any],
    request: Request,
    db: AsyncSession = Depends(get_db),
    orchestration_service: SubscriptionOrchestrationService = Depends(get_subscription_orchestration_service),
    state_service: RegistrationStateService = Depends(get_registration_state_service)
) -> Dict[str, Any]:
    """
    Complete registration after frontend confirms SetupIntent.

    NEW ARCHITECTURE: Creates subscription HERE (not in payment-setup).
    This is the ONLY place subscriptions are created during registration.

    Flow:
    1. Verify SetupIntent status is 'succeeded'
    2. Create subscription with verified payment method
    3. Update registration state

    Args:
        verification_data: SetupIntent verification data with user_data

    Returns:
        Subscription creation result
    """
    setup_intent_id = None
    user_data = {}
    state_id = None

    try:
        # Validate required fields
        if not verification_data.get('setup_intent_id'):
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="SetupIntent ID is required")

        if not verification_data.get('user_data'):
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="User data is required")

        setup_intent_id = verification_data['setup_intent_id']
        user_data = verification_data['user_data']
        state_id = verification_data.get('state_id')

        logger.info("Completing registration after verification",
                   extra={"email": user_data.get('email'), "setup_intent_id": setup_intent_id})

        # Calculate trial period from coupon if provided in the completion call
        trial_period_days = 0
        coupon_code = user_data.get('coupon_code')
        
        if coupon_code:
            logger.info("Validating coupon in completion call",
                       extra={"coupon_code": coupon_code, "email": user_data.get('email')})
            
            # Create coupon service to validate coupon
            coupon_service = CouponService(db)
            success, discount_applied, error = await coupon_service.redeem_coupon(
                coupon_code,
                None,  # No tenant_id yet
                base_trial_days=0
            )
            
            if success and discount_applied:
                trial_period_days = discount_applied.get("total_trial_days", 0)
                logger.info("Coupon validated in completion call",
                           extra={"coupon_code": coupon_code, "trial_period_days": trial_period_days})
            else:
                logger.warning("Failed to validate coupon in completion call",
                             extra={"coupon_code": coupon_code, "error": error})
        elif 'trial_period_days' in user_data:
            # Fallback: use trial_period_days if explicitly provided
            trial_period_days = int(user_data.get('trial_period_days', 0))
            logger.info("Using explicitly provided trial period",
                       extra={"trial_period_days": trial_period_days})
        
        # Create subscription AFTER verification (the core fix!)
        result = await orchestration_service.complete_registration_subscription(
            setup_intent_id=setup_intent_id,
            customer_id=user_data.get('customer_id', ''),
            plan_id=user_data.get('plan_id') or user_data.get('subscription_plan', 'professional'),
            payment_method_id=user_data.get('payment_method_id', ''),
            billing_interval=user_data.get('billing_cycle') or user_data.get('billing_interval', 'monthly'),
            trial_period_days=trial_period_days,
            user_id=user_data.get('user_id')
        )

        # Update registration state
        if state_id:
            try:
                await state_service.update_state_context(state_id, {
                    'subscription_id': result['subscription_id'],
                    'status': result['status']
                })
                await state_service.transition_state(state_id, RegistrationState.SUBSCRIPTION_CREATED)
            except Exception as e:
                logger.warning(f"Failed to update registration state: {e}", extra={"state_id": state_id})

        logger.info("Registration subscription created successfully",
                   extra={
                       "email": user_data.get('email'),
                       "subscription_id": result['subscription_id'],
                       "status": result['status']
                   })

        return {
            "success": True,
            "subscription_id": result['subscription_id'],
            "customer_id": result['customer_id'],
            "payment_customer_id": result.get('payment_customer_id', result['customer_id']),
            "status": result['status'],
            "plan_id": result.get('plan_id'),
            "payment_method_id": result.get('payment_method_id'),
            "trial_period_days": result.get('trial_period_days', 0),
            "current_period_end": result.get('current_period_end'),
            "state_id": state_id,
            "message": "Subscription created successfully"
        }

    except SetupIntentError as e:
        logger.error(f"SetupIntent verification failed: {e}",
                    extra={"setup_intent_id": setup_intent_id, "email": user_data.get('email')},
                    exc_info=True)
        if state_id:
            try:
                await state_service.mark_registration_failed(state_id, f"Verification failed: {e}")
            except Exception:
                pass
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Payment verification failed: {e}") from e

    except SubscriptionCreationFailed as e:
        logger.error(f"Subscription creation failed: {e}",
                    extra={"setup_intent_id": setup_intent_id, "email": user_data.get('email')},
                    exc_info=True)
        if state_id:
            try:
                await state_service.mark_registration_failed(state_id, f"Subscription failed: {e}")
            except Exception:
                pass
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Subscription creation failed: {e}") from e

    except HTTPException:
        raise

    except Exception as e:
        logger.error(f"Unexpected error: {e}",
                    extra={"setup_intent_id": setup_intent_id, "email": user_data.get('email')},
                    exc_info=True)
        if state_id:
            try:
                await state_service.mark_registration_failed(state_id, f"Registration failed: {e}")
            except Exception:
                pass
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Registration failed: {e}") from e


@router.get("/registration-state/{state_id}",
           response_model=Dict[str, Any],
           summary="Get registration state")
async def get_registration_state(
    state_id: str,
    request: Request,
    state_service: RegistrationStateService = Depends(get_registration_state_service)
) -> Dict[str, Any]:
    """
    Get registration state by ID
    
    Args:
        state_id: Registration state ID
        
    Returns:
        Registration state data
        
    Raises:
        HTTPException: 404 if state not found, 500 for server errors
    """
    try:
        logger.info("Getting registration state",
                   extra={"state_id": state_id})
        
        state_data = await state_service.get_registration_state(state_id)
        
        logger.info("Registration state retrieved",
                   extra={
                       "state_id": state_id,
                       "current_state": state_data['current_state']
                   })
        
        return {
            "state_id": state_data['state_id'],
            "email": state_data['email'],
            "current_state": state_data['current_state'],
            "created_at": state_data['created_at'],
            "updated_at": state_data['updated_at'],
            "setup_intent_id": state_data.get('setup_intent_id'),
            "customer_id": state_data.get('customer_id'),
            "subscription_id": state_data.get('subscription_id'),
            "error": state_data.get('error'),
            "user_data": state_data.get('user_data')
        }
        
    except RegistrationStateError as e:
        logger.error("Registration state not found",
                    extra={
                        "error": str(e),
                        "state_id": state_id
                    })
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Registration state not found: {str(e)}"
        ) from e
    except Exception as e:
        logger.error("Unexpected error getting registration state",
                    extra={
                        "error": str(e),
                        "state_id": state_id
                    },
                    exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to get registration state: {str(e)}"
        ) from e


# ============================================================================
# SUBSCRIPTION MANAGEMENT ENDPOINTS
# These endpoints handle ongoing subscription management (upgrade, cancel, etc.)
# ============================================================================

@router.post("/subscriptions/cancel",
            response_model=Dict[str, Any],
            summary="Cancel subscription - Downgrade to read-only mode")
async def cancel_subscription(
    request: Request,
    tenant_id: str = Query(..., description="Tenant ID"),
    reason: str = Query("", description="Cancellation reason"),
    orchestration_service: SubscriptionOrchestrationService = Depends(get_subscription_orchestration_service)
) -> Dict[str, Any]:
    """
    Cancel a subscription and set to read-only mode
    
    This endpoint allows users to cancel their subscription, which will:
    - Mark subscription as pending cancellation
    - Set read-only mode effective date to end of current billing period
    - Allow read-only access until the end of paid period
    
    Args:
        tenant_id: Tenant ID to cancel subscription for
        reason: Optional cancellation reason
        
    Returns:
        Dictionary with cancellation details including effective date
        
    Raises:
        HTTPException: 404 if subscription not found, 400 for validation errors
    """
    try:
        result = await orchestration_service.orchestrate_subscription_cancellation(
            tenant_id,
            reason
        )
        return {
            "success": True,
            "message": result.get("message", "Subscription cancellation initiated"),
            "status": result.get("status", "pending_cancellation"),
            "cancellation_effective_date": result.get("cancellation_effective_date"),
            "days_remaining": result.get("days_remaining"),
            "read_only_mode_starts": result.get("read_only_mode_starts")
        }
    except Exception as e:
        logger.error("Failed to cancel subscription",
                    extra={
                        "error": str(e),
                        "tenant_id": tenant_id
                    },
                    exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to cancel subscription: {str(e)}"
        ) from e

@router.post("/subscriptions/{tenant_id}/upgrade",
            response_model=Dict[str, Any],
            summary="Upgrade subscription plan")
async def upgrade_subscription_plan(
    tenant_id: str,
    request: Request,
    new_plan: str = Query(..., description="New plan name"),
    orchestration_service: SubscriptionOrchestrationService = Depends(get_subscription_orchestration_service)
) -> Dict[str, Any]:
    """
    Upgrade a tenant's subscription plan
    
    Validates the upgrade and updates the subscription if eligible.
    
    Args:
        tenant_id: Tenant ID
        new_plan: New plan name (starter, professional, enterprise)
        
    Returns:
        Dictionary with upgrade result
        
    Raises:
        HTTPException: 400 if upgrade not allowed, 404 if tenant not found
    """
    try:
        # Perform the upgrade (validation is built into the orchestration)
        result = await orchestration_service.orchestrate_plan_upgrade(tenant_id, new_plan)
        
        return {
            "success": True,
            "message": result.get("message", f"Plan successfully upgraded to {new_plan}"),
            "new_plan": new_plan,
            "old_plan": result.get("old_plan"),
            "proration_details": result.get("proration_details"),
            "next_billing_date": result.get("next_billing_date"),
            "billing_cycle": result.get("billing_cycle")
        }
    except Exception as e:
        logger.error("Failed to upgrade subscription plan",
                    extra={
                        "error": str(e),
                        "tenant_id": tenant_id,
                        "new_plan": new_plan
                    },
                    exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to upgrade subscription plan: {str(e)}"
        ) from e

@router.get("/subscriptions/{tenant_id}/validate-upgrade/{new_plan}",
           response_model=Dict[str, Any],
           summary="Validate plan upgrade eligibility")
async def validate_plan_upgrade(
    tenant_id: str,
    new_plan: str,
    orchestration_service: SubscriptionOrchestrationService = Depends(get_subscription_orchestration_service)
) -> Dict[str, Any]:
    """
    Validate if a tenant can upgrade to a new plan
    
    Checks plan hierarchy and subscription status before allowing upgrade.
    
    Args:
        tenant_id: Tenant ID
        new_plan: Plan to validate upgrade to
        
    Returns:
        Dictionary with validation result and reason if not allowed
    """
    try:
        validation = await orchestration_service.validate_plan_upgrade(tenant_id, new_plan)
        return validation
    except Exception as e:
        logger.error("Failed to validate plan upgrade",
                    extra={
                        "error": str(e),
                        "tenant_id": tenant_id,
                        "new_plan": new_plan
                    },
                    exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to validate plan upgrade: {str(e)}"
        ) from e

@router.get("/{tenant_id}/subscriptions/status",
           response_model=Dict[str, Any],
           summary="Get subscription status")
async def get_subscription_status(
    tenant_id: str,
    db: AsyncSession = Depends(get_db)
) -> Dict[str, Any]:
    """
    Get subscription status for read-only mode enforcement
    """
    try:
        from sqlalchemy import select
        from app.models.tenants import Subscription
        
        # Find active subscription for tenant
        result = await db.execute(
            select(Subscription).where(
                Subscription.tenant_id == tenant_id,
                Subscription.status == "active"
            )
        )
        subscription = result.scalars().first()
        
        if subscription:
            return {
                "status": subscription.status,
                "plan": subscription.plan,
                "is_read_only": False,
                "cancellation_effective_date": subscription.cancellation_effective_date.isoformat() if subscription.cancellation_effective_date else None
            }
        else:
            # No active subscription found
            return {
                "status": "inactive",
                "plan": None,
                "is_read_only": True,
                "cancellation_effective_date": None
            }
    except Exception as e:
        logger.error(f"Failed to get subscription status: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to get subscription status: {str(e)}"
        ) from e

@router.get("/setup-intents/{setup_intent_id}/verify",
           response_model=Dict[str, Any],
           summary="Verify SetupIntent status for registration")
async def verify_setup_intent(
    setup_intent_id: str,
    request: Request,
    orchestration_service: SubscriptionOrchestrationService = Depends(get_subscription_orchestration_service)
) -> Dict[str, Any]:
    """
    Verify SetupIntent status for registration flow

    This endpoint verifies the status of a SetupIntent that was created
    during the registration payment setup flow.

    Args:
        setup_intent_id: The SetupIntent ID to verify

    Returns:
        Dictionary with verification result including status and any required actions

    Raises:
        HTTPException: 404 if SetupIntent not found, 500 for server errors
    """
    try:
        logger.info("Verifying SetupIntent for registration",
                   extra={"setup_intent_id": setup_intent_id})

        # Verify the SetupIntent status via orchestration service
        result = await orchestration_service.verify_setup_intent_for_registration(setup_intent_id)

        logger.info("SetupIntent verification completed",
                   extra={
                       "setup_intent_id": setup_intent_id,
                       "status": result.get('status')
                   })

        return {
            "success": True,
            "setup_intent_id": setup_intent_id,
            "status": result.get('status'),
            "payment_method_id": result.get('payment_method_id'),
            "customer_id": result.get('customer_id'),
            "requires_action": result.get('requires_action', False),
            "action_type": result.get('action_type'),
            "client_secret": result.get('client_secret'),
            "message": result.get('message', 'SetupIntent verification completed successfully')
        }

    except Exception as e:
        logger.error("SetupIntent verification failed",
                    extra={
                        "error": str(e),
                        "setup_intent_id": setup_intent_id
                    },
                    exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"SetupIntent verification failed: {str(e)}"
        ) from e