# Replace your existing billing.py with this Paystack-integrated version

import os
import requests
import json
from datetime import datetime, timedelta
from flask import Blueprint, render_template, request, redirect, url_for, flash, jsonify
from flask_login import login_required, current_user
from functools import wraps
from models import db, Organization, Subscription, Invoice

billing_bp = Blueprint('billing', __name__)

# Paystack Configuration
PAYSTACK_SECRET_KEY = os.getenv('PAYSTACK_SECRET_KEY')
PAYSTACK_PUBLIC_KEY = os.getenv('PAYSTACK_PUBLIC_KEY')
PAYSTACK_BASE_URL = os.getenv('PAYSTACK_BASE_URL', 'https://api.paystack.co')

# Subscription Plans (Kenya-optimized pricing in KES)
SUBSCRIPTION_PLANS = {
    'free': {
        'name': 'Free Plan',
        'price': 0,
        'currency': 'KES',
        'limits': {
            'assessments_per_month': 5,
            'projects': 3,
            'users': 1,
            'reports': 'basic'
        },
        'features': ['Basic assessments', 'Limited reports', 'Community support']
    },
    'professional': {
        'name': 'Professional',
        'price': 6500,  # KES ~$50
        'currency': 'KES',
        'plan_code': os.getenv('PAYSTACK_PRO_PLAN_CODE'),
        'limits': {
            'assessments_per_month': 100,
            'projects': 'unlimited',
            'users': 10,
            'reports': 'advanced'
        },
        'features': ['Unlimited assessments', 'Advanced analytics', 'Priority support', 'Custom reports']
    },
    'enterprise': {
        'name': 'Enterprise',
        'price': 26000,  # KES ~$200
        'currency': 'KES', 
        'plan_code': os.getenv('PAYSTACK_ENTERPRISE_PLAN_CODE'),
        'limits': {
            'assessments_per_month': 'unlimited',
            'projects': 'unlimited',
            'users': 'unlimited',
            'reports': 'premium'
        },
        'features': ['Everything in Pro', 'SSO integration', 'Compliance reports', 'Dedicated support', 'API access']
    }
}

def paystack_request(method, endpoint, data=None):
    """Make requests to Paystack API"""
    headers = {
        'Authorization': f'Bearer {PAYSTACK_SECRET_KEY}',
        'Content-Type': 'application/json'
    }
    
    url = f"{PAYSTACK_BASE_URL}/{endpoint}"
    
    if method == 'GET':
        response = requests.get(url, headers=headers, params=data)
    elif method == 'POST':
        response = requests.post(url, headers=headers, json=data)
    elif method == 'PUT':
        response = requests.put(url, headers=headers, json=data)
    else:
        raise ValueError(f"Unsupported method: {method}")
    
    return response.json()

@billing_bp.route('/billing')
@login_required
def billing_dashboard():
    org = current_user.organization
    subscription = Subscription.query.filter_by(organization_id=org.id).first()
    recent_invoices = Invoice.query.filter_by(organization_id=org.id).order_by(Invoice.created_at.desc()).limit(5).all()

    return render_template('billing/dashboard.html',
                         organization=org,
                         subscription=subscription,
                         recent_invoices=recent_invoices,
                         plans=SUBSCRIPTION_PLANS,
                         paystack_public_key=PAYSTACK_PUBLIC_KEY)

@billing_bp.route('/upgrade/<plan>')
@login_required
def upgrade_plan(plan):
    if plan not in SUBSCRIPTION_PLANS:
        flash('Invalid subscription plan', 'error')
        return redirect(url_for('billing.billing_dashboard'))

    if plan == 'free':
        flash('You are already on the free plan', 'info')
        return redirect(url_for('billing.billing_dashboard'))

    org = current_user.organization

    # Check if already on this plan
    if org.subscription_plan == plan:
        flash(f'You are already on the {SUBSCRIPTION_PLANS[plan]["name"]} plan', 'info')
        return redirect(url_for('billing.billing_dashboard'))

    try:
        # Create or get Paystack customer
        if not org.paystack_customer_code:
            customer_data = {
                'email': current_user.email,
                'first_name': current_user.first_name,
                'last_name': current_user.last_name,
                'phone': getattr(current_user, 'phone', ''),
                'metadata': {
                    'organization_id': org.id,
                    'organization_name': org.name
                }
            }
            
            response = paystack_request('POST', 'customer', customer_data)
            
            if response['status']:
                org.paystack_customer_code = response['data']['customer_code']
                db.session.commit()
            else:
                flash('Error creating customer profile', 'error')
                return redirect(url_for('billing.billing_dashboard'))

        # Initialize transaction for subscription
        transaction_data = {
            'email': current_user.email,
            'amount': SUBSCRIPTION_PLANS[plan]['price'] * 100,  # Convert to kobo
            'currency': 'KES',
            'plan': SUBSCRIPTION_PLANS[plan]['plan_code'],
            'customer': org.paystack_customer_code,
            'callback_url': url_for('billing.upgrade_success', _external=True),
            'metadata': {
                'organization_id': org.id,
                'plan': plan,
                'cancel_action': url_for('billing.billing_dashboard', _external=True)
            }
        }

        response = paystack_request('POST', 'transaction/initialize', transaction_data)
        
        if response['status']:
            return redirect(response['data']['authorization_url'])
        else:
            flash('Error initializing payment', 'error')
            return redirect(url_for('billing.billing_dashboard'))

    except Exception as e:
        flash(f'Error processing upgrade: {str(e)}', 'error')
        return redirect(url_for('billing.billing_dashboard'))

@billing_bp.route('/upgrade/success')
@login_required
def upgrade_success():
    reference = request.args.get('reference')
    
    if not reference:
        flash('No payment reference found', 'error')
        return redirect(url_for('billing.billing_dashboard'))
    
    try:
        # Verify the transaction
        response = paystack_request('GET', f'transaction/verify/{reference}')
        
        if response['status'] and response['data']['status'] == 'success':
            # Update organization subscription
            org = current_user.organization
            metadata = response['data']['metadata']
            plan = metadata.get('plan')
            
            if plan in SUBSCRIPTION_PLANS:
                org.subscription_plan = plan
                org.subscription_status = 'active'
                org.paystack_subscription_code = response['data'].get('subscription', {}).get('subscription_code', '')
                
                # Update limits based on plan
                limits = SUBSCRIPTION_PLANS[plan]['limits']
                if limits['assessments_per_month'] != 'unlimited':
                    org.monthly_assessments_limit = limits['assessments_per_month']
                else:
                    org.monthly_assessments_limit = 999999  # Large number for unlimited
                
                org.subscription_ends_at = datetime.now() + timedelta(days=30)  # Monthly subscription
                
                db.session.commit()
                
                flash(f'Successfully upgraded to {SUBSCRIPTION_PLANS[plan]["name"]} plan!', 'success')
            else:
                flash('Invalid plan in payment metadata', 'error')
        else:
            flash('Payment verification failed', 'error')
            
    except Exception as e:
        flash(f'Error verifying payment: {str(e)}', 'error')

    return redirect(url_for('billing.billing_dashboard'))

@billing_bp.route('/cancel-subscription', methods=['POST'])
@login_required
def cancel_subscription():
    if not current_user.can_manage_org():
        flash('You do not have permission to cancel the subscription', 'error')
        return redirect(url_for('billing.billing_dashboard'))

    org = current_user.organization

    if not org.paystack_subscription_code:
        flash('No active subscription to cancel', 'error')
        return redirect(url_for('billing.billing_dashboard'))

    try:
        # Disable the subscription
        response = paystack_request('POST', f'subscription/disable', {
            'code': org.paystack_subscription_code,
            'token': current_user.email
        })
        
        if response['status']:
            org.subscription_status = 'canceled'
            db.session.commit()
            flash('Your subscription has been canceled', 'info')
        else:
            flash('Error canceling subscription', 'error')

    except Exception as e:
        flash(f'Error canceling subscription: {str(e)}', 'error')

    return redirect(url_for('billing.billing_dashboard'))

@billing_bp.route('/webhooks/paystack', methods=['POST'])
def paystack_webhook():
    """Handle Paystack webhooks"""
    payload = request.get_data()
    sig_header = request.headers.get('X-Paystack-Signature')
    
    # Verify webhook signature
    webhook_secret = os.getenv('WEBHOOK_SECRET', '')
    
    try:
        import hmac
        import hashlib
        
        expected_signature = hmac.new(
            webhook_secret.encode('utf-8'),
            payload,
            hashlib.sha512
        ).hexdigest()
        
        if not hmac.compare_digest(expected_signature, sig_header or ''):
            return 'Invalid signature', 400
            
    except Exception:
        return 'Invalid payload', 400

    try:
        event = json.loads(payload.decode('utf-8'))
    except json.JSONDecodeError:
        return 'Invalid JSON', 400

    # Handle different event types
    event_type = event.get('event')
    
    if event_type == 'subscription.create':
        handle_subscription_created(event['data'])
    elif event_type == 'subscription.disable':
        handle_subscription_disabled(event['data'])
    elif event_type == 'invoice.create':
        handle_invoice_created(event['data'])
    elif event_type == 'invoice.update':
        handle_invoice_updated(event['data'])
    elif event_type == 'invoice.payment_failed':
        handle_payment_failed(event['data'])
    elif event_type == 'charge.success':
        handle_charge_success(event['data'])

    return 'Success', 200

def handle_subscription_created(data):
    """Handle subscription creation webhook"""
    customer_code = data.get('customer', {}).get('customer_code')
    
    org = Organization.query.filter_by(paystack_customer_code=customer_code).first()
    if not org:
        return

    org.paystack_subscription_code = data['subscription_code']
    org.subscription_status = data['status']
    org.subscription_ends_at = datetime.fromisoformat(data['next_payment_date'].replace('Z', '+00:00'))
    
    db.session.commit()

def handle_subscription_disabled(data):
    """Handle subscription cancellation webhook"""
    subscription_code = data['subscription_code']
    
    org = Organization.query.filter_by(paystack_subscription_code=subscription_code).first()
    if not org:
        return

    org.subscription_plan = 'free'
    org.subscription_status = 'canceled'
    org.monthly_assessments_limit = 5
    
    db.session.commit()

def handle_invoice_created(data):
    """Handle invoice creation webhook"""
    customer_code = data.get('customer', {}).get('customer_code')
    
    org = Organization.query.filter_by(paystack_customer_code=customer_code).first()
    if not org:
        return
    
    # Create invoice record
    invoice = Invoice(
        organization_id=org.id,
        paystack_invoice_id=data['id'],
        amount=data['amount'] / 100,  # Convert from kobo
        currency=data['currency'],
        status=data['status'],
        created_at=datetime.fromisoformat(data['created_at'].replace('Z', '+00:00'))
    )
    
    db.session.add(invoice)
    db.session.commit()

def handle_invoice_updated(data):
    """Handle invoice update webhook"""
    invoice = Invoice.query.filter_by(paystack_invoice_id=data['id']).first()
    if invoice:
        invoice.status = data['status']
        if data['status'] == 'success':
            invoice.paid_at = datetime.now()
        db.session.commit()

def handle_payment_failed(data):
    """Handle payment failure webhook"""
    customer_code = data.get('customer', {}).get('customer_code')
    
    org = Organization.query.filter_by(paystack_customer_code=customer_code).first()
    if not org:
        return

    # Update subscription status
    if org.subscription_status != 'past_due':
        org.subscription_status = 'past_due'
        db.session.commit()

def handle_charge_success(data):
    """Handle successful charge webhook"""
    customer_code = data.get('customer', {}).get('customer_code')
    
    org = Organization.query.filter_by(paystack_customer_code=customer_code).first()
    if not org:
        return
    
    # Update subscription status if needed
    if org.subscription_status == 'past_due':
        org.subscription_status = 'active'
        db.session.commit()

# Usage tracking decorators (keep your existing ones)
def track_assessment_usage(f):
    """Decorator to track assessment usage"""
    def decorated_function(*args, **kwargs):
        if current_user.is_authenticated:
            org = current_user.organization
            
            # Check if within limits
            if hasattr(org, 'monthly_assessments_used') and hasattr(org, 'monthly_assessments_limit'):
                if org.monthly_assessments_used >= org.monthly_assessments_limit:
                    flash('You have reached your monthly assessment limit. Please upgrade your plan.', 'warning')
                    return redirect(url_for('billing.billing_dashboard'))
                
                # Increment usage
                org.monthly_assessments_used = (org.monthly_assessments_used or 0) + 1
                db.session.commit()
        
        return f(*args, **kwargs)
    return decorated_function