# subscription_routes.py

from flask import Blueprint, render_template, request, jsonify, session, redirect, url_for, flash, current_app
import stripe
from datetime import datetime, timedelta
from database import db
from models import User, Subscription, SubscriptionPlan, PaymentHistory, SingleAnalysisPurchase
from auth_routes import login_required, log_activity
from services.stripe_service import StripeService
import os

subscription_bp = Blueprint('subscription', __name__, url_prefix='/subscription')

@subscription_bp.route('/enterprise-features')
def enterprise_features():
    """Display Enterprise features page for Pro users"""
    user_id = session.get('user_id')
    if not user_id:
        flash('Please log in to view Enterprise features', 'warning')
        return redirect(url_for('auth.login'))

    # Get current user's subscription to verify they're on Pro
    current_subscription = Subscription.query.filter_by(
        user_id=user_id,
        status='active'
    ).first()

    # Redirect non-Pro users to regular plans page
    if not current_subscription or current_subscription.plan.name != 'pro':
        return redirect(url_for('subscription.view_plans'))

    return render_template('subscription/enterprise_features.html')

@subscription_bp.route('/plans')
def view_plans():
    """Display available subscription plans"""
    plans = SubscriptionPlan.query.filter_by(is_active=True).order_by(SubscriptionPlan.price).all()
    
    # Get current user's subscription
    user_id = session.get('user_id')
    current_subscription = None
    if user_id:
        current_subscription = Subscription.query.filter_by(
            user_id=user_id, 
            status='active'
        ).first()
    
    return render_template('subscription/plans.html', 
                          plans=plans, 
                          current_subscription=current_subscription,
                          stripe_publishable_key=current_app.config.get('STRIPE_PUBLISHABLE_KEY'))

@subscription_bp.route('/subscribe/<int:plan_id>')
@login_required
def subscribe(plan_id):
    """Start subscription process"""
    user_id = session.get('user_id')
    user = User.query.get(user_id)
    plan = SubscriptionPlan.query.get_or_404(plan_id)
    
    # Check for existing subscription
    existing_subscription = Subscription.query.filter_by(
        user_id=user_id, 
        status='active'
    ).first()
    
    if existing_subscription:
        # Allow upgrades from free plan
        if existing_subscription.plan.name == 'free' and plan.name != 'free':
            # User is upgrading from free to paid - this is allowed
            pass
        # Prevent downgrades to free if user has paid plan
        elif existing_subscription.plan.name != 'free' and plan.name == 'free':
            flash('Cannot downgrade to free plan. Please cancel your subscription first.', 'warning')
            return redirect(url_for('subscription.manage'))
        # Prevent switching between paid plans (require cancellation first)
        elif existing_subscription.plan.name != 'free' and plan.name != 'free':
            flash('Please cancel your current subscription before switching plans.', 'warning')
            return redirect(url_for('subscription.manage'))
        # Prevent subscribing to same plan
        elif existing_subscription.plan_id == plan_id:
            flash('You are already subscribed to this plan', 'info')
            return redirect(url_for('subscription.manage'))
    
    # For free plan, create subscription directly
    if plan.name == 'free':
        # If upgrading from free, update existing subscription
        if existing_subscription and existing_subscription.plan.name == 'free':
            flash('You are already on the free plan', 'info')
            return redirect(url_for('subscription.manage'))
            
        subscription = Subscription(
            user_id=user_id,
            company_id=user.company_id,
            plan_id=plan_id,
            status='active',
            current_period_start=datetime.utcnow(),
            current_period_end=datetime.utcnow() + timedelta(days=30),
            monthly_reset_date=datetime.utcnow() + timedelta(days=30)
        )
        
        db.session.add(subscription)
        db.session.commit()
        
        log_activity(user_id, 'subscription_created', f'Subscribed to {plan.display_name}')
        flash(f'Successfully subscribed to {plan.display_name}!', 'success')
        return redirect(url_for('dashboard'))
    
    # For paid plans, redirect to Stripe Checkout
    try:
        # If upgrading from free, we'll update the existing subscription later
        if existing_subscription and existing_subscription.plan.name == 'free':
            # Mark the free subscription as pending upgrade
            existing_subscription.status = 'pending_upgrade'
            db.session.commit()
        
        # Create or get Stripe customer
        if not user.stripe_customer_id:
            customer = StripeService.create_customer(user)
            if customer:
                user.stripe_customer_id = customer.id
                db.session.commit()
            else:
                flash('Payment setup failed. Please try again.', 'danger')
                return redirect(url_for('subscription.view_plans'))
        
        # Create checkout session
        success_url = url_for('subscription.success', _external=True) + '?session_id={CHECKOUT_SESSION_ID}'
        cancel_url = url_for('subscription.view_plans', _external=True)
        
        checkout_session = StripeService.create_checkout_session(
            price_id=plan.stripe_price_id,
            customer_id=user.stripe_customer_id,
            success_url=success_url,
            cancel_url=cancel_url,
            user_id=user_id
        )
        
        if checkout_session:
            return redirect(checkout_session.url)
        else:
            flash('Payment setup failed. Please try again.', 'danger')
            return redirect(url_for('subscription.view_plans'))
            
    except Exception as e:
        current_app.logger.error(f"Subscription creation error: {str(e)}")
        flash('An error occurred. Please try again.', 'danger')
        return redirect(url_for('subscription.view_plans'))
@subscription_bp.route('/success')
@login_required
def success():
    """Handle successful subscription"""
    session_id = request.args.get('session_id')
    
    if not session_id:
        flash('Invalid payment session', 'danger')
        return redirect(url_for('dashboard'))
    
    try:
        # Retrieve the checkout session
        checkout_session = stripe.checkout.Session.retrieve(session_id)
        
        # Get subscription details
        stripe_subscription = stripe.Subscription.retrieve(checkout_session.subscription)
        
        user_id = session.get('user_id')
        user = User.query.get(user_id)
        
        # Find the plan by stripe price id
        plan = SubscriptionPlan.query.filter_by(
            stripe_price_id=stripe_subscription.items.data[0].price.id
        ).first()
        
        if not plan:
            flash('Plan not found', 'danger')
            return redirect(url_for('dashboard'))
        
        # Create subscription record
        subscription = Subscription(
            user_id=user_id,
            company_id=user.company_id,
            plan_id=plan.id,
            stripe_customer_id=checkout_session.customer,
            stripe_subscription_id=stripe_subscription.id,
            status='active',
            current_period_start=datetime.fromtimestamp(stripe_subscription.current_period_start),
            current_period_end=datetime.fromtimestamp(stripe_subscription.current_period_end),
            monthly_reset_date=datetime.utcnow() + timedelta(days=30)
        )
        
        db.session.add(subscription)
        db.session.commit()
        
        log_activity(user_id, 'subscription_created', f'Subscribed to {plan.display_name}')
        flash(f'Welcome to {plan.display_name}! Your subscription is now active.', 'success')
        
    except Exception as e:
        current_app.logger.error(f"Subscription success handling error: {str(e)}")
        flash('There was an issue setting up your subscription. Please contact support.', 'danger')
    
    return redirect(url_for('dashboard'))

@subscription_bp.route('/manage')
@login_required
def manage():
    """Subscription management page"""
    user_id = session.get('user_id')
    user = User.query.get(user_id)
    
    subscription = Subscription.query.filter_by(
        user_id=user_id, 
        status='active'
    ).first()
    
    # Get payment history
    payment_history = []
    if subscription:
        payment_history = PaymentHistory.query.filter_by(
            subscription_id=subscription.id
        ).order_by(PaymentHistory.created_at.desc()).limit(10).all()
    
    # Get usage statistics
    usage_stats = None
    if subscription:
        usage_stats = {
            'analyses_used': subscription.monthly_analyses_used,
            'analyses_limit': subscription.plan.features.get('monthly_analyses', 0),
            'period_end': subscription.current_period_end,
            'can_analyze': subscription.can_analyze_document()
        }
    
    return render_template('subscription/manage.html',
                          subscription=subscription,
                          payment_history=payment_history,
                          usage_stats=usage_stats,
                          user=user)

@subscription_bp.route('/cancel', methods=['POST'])
@login_required
def cancel_subscription():
    """Cancel user's subscription"""
    user_id = session.get('user_id')
    
    subscription = Subscription.query.filter_by(
        user_id=user_id, 
        status='active'
    ).first()
    
    if not subscription:
        return jsonify({'error': 'No active subscription found'}), 404
    
    # Don't cancel free plans through Stripe
    if subscription.plan.name == 'free':
        subscription.status = 'canceled'
        db.session.commit()
        log_activity(user_id, 'subscription_canceled', 'Canceled free subscription')
        return jsonify({'success': True, 'message': 'Subscription canceled'})
    
    # Cancel Stripe subscription
    if subscription.stripe_subscription_id:
        stripe_subscription = StripeService.cancel_subscription(
            subscription.stripe_subscription_id,
            cancel_at_period_end=True
        )
        
        if stripe_subscription:
            subscription.cancel_at_period_end = True
            db.session.commit()
            log_activity(user_id, 'subscription_canceled', f'Canceled {subscription.plan.display_name}')
            return jsonify({
                'success': True, 
                'message': 'Subscription will be canceled at the end of the billing period'
            })
    
    return jsonify({'error': 'Failed to cancel subscription'}), 500

@subscription_bp.route('/billing-portal')
@login_required
def billing_portal():
    """Redirect to Stripe billing portal"""
    user_id = session.get('user_id')
    user = User.query.get(user_id)
    
    if not user.stripe_customer_id:
        flash('No billing information found', 'danger')
        return redirect(url_for('subscription.manage'))
    
    return_url = url_for('subscription.manage', _external=True)
    
    session_obj = StripeService.create_billing_portal_session(
        customer_id=user.stripe_customer_id,
        return_url=return_url
    )
    
    if session_obj:
        return redirect(session_obj.url)
    else:
        flash('Unable to access billing portal', 'danger')
        return redirect(url_for('subscription.manage'))

@subscription_bp.route('/webhook', methods=['POST'])
def stripe_webhook():
    """Handle Stripe webhooks"""
    payload = request.get_data()
    sig_header = request.headers.get('Stripe-Signature')
    webhook_secret = current_app.config.get('STRIPE_WEBHOOK_SECRET')
    
    try:
        event = stripe.Webhook.construct_event(payload, sig_header, webhook_secret)
    except ValueError:
        current_app.logger.error("Invalid payload in webhook")
        return '', 400
    except stripe.error.SignatureVerificationError:
        current_app.logger.error("Invalid signature in webhook")
        return '', 400
    
    # Handle the event
    if event['type'] == 'customer.subscription.updated':
        handle_subscription_updated(event['data']['object'])
    elif event['type'] == 'customer.subscription.deleted':
        handle_subscription_deleted(event['data']['object'])
    elif event['type'] == 'invoice.payment_succeeded':
        handle_payment_succeeded(event['data']['object'])
    elif event['type'] == 'invoice.payment_failed':
        handle_payment_failed(event['data']['object'])
    
    return '', 200

def handle_subscription_updated(subscription):
    """Handle subscription update webhook"""
    try:
        db_subscription = Subscription.query.filter_by(
            stripe_subscription_id=subscription['id']
        ).first()
        
        if db_subscription:
            db_subscription.status = subscription['status']
            db_subscription.current_period_start = datetime.fromtimestamp(subscription['current_period_start'])
            db_subscription.current_period_end = datetime.fromtimestamp(subscription['current_period_end'])
            db_subscription.cancel_at_period_end = subscription.get('cancel_at_period_end', False)
            
            db.session.commit()
            current_app.logger.info(f"Updated subscription {subscription['id']}")
    except Exception as e:
        current_app.logger.error(f"Error handling subscription update: {str(e)}")

def handle_subscription_deleted(subscription):
    """Handle subscription deletion webhook"""
    try:
        db_subscription = Subscription.query.filter_by(
            stripe_subscription_id=subscription['id']
        ).first()
        
        if db_subscription:
            db_subscription.status = 'canceled'
            db.session.commit()
            current_app.logger.info(f"Canceled subscription {subscription['id']}")
    except Exception as e:
        current_app.logger.error(f"Error handling subscription deletion: {str(e)}")

def handle_payment_succeeded(invoice):
    """Handle successful payment webhook"""
    try:
        subscription_id = invoice['subscription']
        
        db_subscription = Subscription.query.filter_by(
            stripe_subscription_id=subscription_id
        ).first()
        
        if db_subscription:
            # Create payment record
            payment = PaymentHistory(
                subscription_id=db_subscription.id,
                stripe_payment_intent_id=invoice.get('payment_intent'),
                amount=invoice['amount_paid'] / 100,  # Convert from cents
                currency=invoice['currency'],
                status='succeeded',
                description=f"Payment for {db_subscription.plan.display_name}"
            )
            
            db.session.add(payment)
            db.session.commit()
            current_app.logger.info(f"Recorded successful payment for subscription {subscription_id}")
    except Exception as e:
        current_app.logger.error(f"Error handling payment success: {str(e)}")

def handle_payment_failed(invoice):
    """Handle failed payment webhook"""
    try:
        subscription_id = invoice['subscription']
        
        db_subscription = Subscription.query.filter_by(
            stripe_subscription_id=subscription_id
        ).first()
        
        if db_subscription:
            # Create payment record
            payment = PaymentHistory(
                subscription_id=db_subscription.id,
                amount=invoice['amount_due'] / 100,  # Convert from cents
                currency=invoice['currency'],
                status='failed',
                description=f"Failed payment for {db_subscription.plan.display_name}"
            )
            
            db.session.add(payment)
            db.session.commit()
            current_app.logger.info(f"Recorded failed payment for subscription {subscription_id}")
    except Exception as e:
        current_app.logger.error(f"Error handling payment failure: {str(e)}")

# Decorators for subscription checking
def subscription_required(required_plan=None):
    """Decorator to check if user has required subscription"""
    def decorator(f):
        def decorated_function(*args, **kwargs):
            user_id = session.get('user_id')
            if not user_id:
                return redirect(url_for('auth.login'))
            
            user = User.query.get(user_id)
            subscription = user.get_subscription()
            
            if required_plan:
                if not subscription or subscription.plan.name != required_plan:
                    flash(f'This feature requires {required_plan.title()} plan', 'warning')
                    return redirect(url_for('subscription.view_plans'))
            
            return f(*args, **kwargs)
        
        decorated_function.__name__ = f.__name__
        return decorated_function
    return decorator

def check_usage_limit():
    """Middleware to check usage limits before analysis"""
    def decorator(f):
        def decorated_function(*args, **kwargs):
            user_id = session.get('user_id')
            if user_id:
                user = User.query.get(user_id)
                subscription = user.get_subscription()

                if subscription and not subscription.can_analyze_document():
                    # Return special status code to trigger upgrade/purchase modal
                    return jsonify({
                        'error': 'Monthly analysis limit reached',
                        'message': 'Please upgrade your plan or purchase a single analysis',
                        'show_options': True
                    }), 429

            return f(*args, **kwargs)

        decorated_function.__name__ = f.__name__
        return decorated_function
    return decorator

@subscription_bp.route('/purchase-single-analysis')
@login_required
def purchase_single_analysis():
    """Initiate purchase of a single analysis for $15"""
    user_id = session.get('user_id')
    user = User.query.get(user_id)

    try:
        # Create Stripe checkout session for single analysis
        success_url = url_for('subscription.single_analysis_success', _external=True) + '?session_id={CHECKOUT_SESSION_ID}'
        cancel_url = url_for('dashboard', _external=True)

        # Create or get Stripe customer
        if not user.stripe_customer_id:
            customer = StripeService.create_customer(user)
            if customer:
                user.stripe_customer_id = customer.id
                db.session.commit()

        # Create checkout session for $15 single analysis
        checkout_session = stripe.checkout.Session.create(
            payment_method_types=['card'],
            customer=user.stripe_customer_id if user.stripe_customer_id else None,
            line_items=[{
                'price_data': {
                    'currency': 'usd',
                    'product_data': {
                        'name': 'Single Document Analysis',
                        'description': 'One-time document analysis with AI insights',
                    },
                    'unit_amount': 975,  # $9.75 in cents
                },
                'quantity': 1,
            }],
            mode='payment',
            success_url=success_url,
            cancel_url=cancel_url,
            metadata={
                'user_id': user_id,
                'type': 'single_analysis'
            }
        )

        # Create pending purchase record
        purchase = SingleAnalysisPurchase(
            user_id=user_id,
            company_id=user.company_id,
            stripe_session_id=checkout_session.id,
            amount=15.00,
            status='pending'
        )
        db.session.add(purchase)
        db.session.commit()

        return redirect(checkout_session.url)

    except Exception as e:
        current_app.logger.error(f"Single analysis purchase error: {str(e)}")
        flash('Unable to process payment. Please try again.', 'danger')
        return redirect(url_for('dashboard'))

@subscription_bp.route('/single-analysis-success')
@login_required
def single_analysis_success():
    """Handle successful single analysis purchase"""
    session_id = request.args.get('session_id')

    if not session_id:
        flash('Invalid payment session', 'danger')
        return redirect(url_for('dashboard'))

    try:
        # Retrieve the checkout session
        checkout_session = stripe.checkout.Session.retrieve(session_id)

        # Update purchase record
        purchase = SingleAnalysisPurchase.query.filter_by(
            stripe_session_id=session_id
        ).first()

        if purchase:
            purchase.status = 'completed'
            purchase.stripe_payment_intent_id = checkout_session.payment_intent
            db.session.commit()

            log_activity(session.get('user_id'), 'single_analysis_purchased', 'Purchased single document analysis')
            flash('Payment successful! You can now analyze one additional document.', 'success')
        else:
            flash('Purchase record not found. Please contact support.', 'danger')

    except Exception as e:
        current_app.logger.error(f"Single analysis success handling error: {str(e)}")
        flash('There was an issue confirming your purchase. Please contact support.', 'danger')

    return redirect(url_for('dashboard'))