"""
Subscription Management System for Toxo
Handles user subscriptions, payment processing, and feature restrictions.
"""

import razorpay
import json
import uuid
from datetime import datetime, timedelta
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum

from ..utils.logger import get_logger
from ..utils.exceptions import SecurityError
from .database import get_database_manager


class SubscriptionPlan(Enum):
    """Subscription plan types."""
    FREE = "free"
    PROFESSIONAL = "professional" 
    ENTERPRISE = "enterprise"


class PaymentStatus(Enum):
    """Payment status types."""
    PENDING = "pending"
    SUCCESS = "success"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class PlanLimits:
    """Plan limits and features."""
    max_layers: Optional[int]  # None means unlimited
    max_training_sessions_per_month: Optional[int]
    advanced_training: bool
    priority_support: bool
    private_layers: bool
    collaborative_workspaces: bool
    custom_integrations: bool
    dedicated_infrastructure: bool
    sla_guarantees: bool
    
    # Professional features
    analytics_dashboard: bool = False
    advanced_api_access: bool = False
    custom_model_integrations: bool = False
    
    # Enterprise features
    white_label: bool = False
    dedicated_account_manager: bool = False
    on_premise: bool = False


# Plan configurations
PLAN_CONFIGS = {
    SubscriptionPlan.FREE: PlanLimits(
        max_layers=1,
        max_training_sessions_per_month=10,
        advanced_training=False,
        priority_support=False,
        private_layers=False,
        collaborative_workspaces=False,
        custom_integrations=False,
        dedicated_infrastructure=False,
        sla_guarantees=False,
        analytics_dashboard=False,
        advanced_api_access=False,
        custom_model_integrations=False
    ),
    SubscriptionPlan.PROFESSIONAL: PlanLimits(
        max_layers=None,  # Unlimited
        max_training_sessions_per_month=None,  # Unlimited
        advanced_training=True,
        priority_support=True,
        private_layers=True,
        collaborative_workspaces=True,
        custom_integrations=True,
        dedicated_infrastructure=False,
        sla_guarantees=False,
        analytics_dashboard=True,
        advanced_api_access=True,
        custom_model_integrations=True
    ),
    SubscriptionPlan.ENTERPRISE: PlanLimits(
        max_layers=None,  # Unlimited
        max_training_sessions_per_month=None,  # Unlimited
        advanced_training=True,
        priority_support=True,
        private_layers=True,
        collaborative_workspaces=True,
        custom_integrations=True,
        dedicated_infrastructure=True,
        sla_guarantees=True,
        analytics_dashboard=True,
        advanced_api_access=True,
        custom_model_integrations=True,
        white_label=True,
        dedicated_account_manager=True,
        on_premise=True
    )
}

# Plan pricing (in INR for Razorpay)
PLAN_PRICING = {
    SubscriptionPlan.FREE: 0,
    SubscriptionPlan.PROFESSIONAL: 2900,  # $29 = ~₹2900
    SubscriptionPlan.ENTERPRISE: None  # Custom pricing
}


class SubscriptionManager:
    """
    Manages user subscriptions, payment processing, and feature access.
    """
    
    def __init__(self, razorpay_key_id: str = None, razorpay_key_secret: str = None):
        self.logger = get_logger(__name__)
        self.db = get_database_manager()
        
        # Initialize Razorpay client
        if razorpay_key_id and razorpay_key_secret:
            self.razorpay_client = razorpay.Client(auth=(razorpay_key_id, razorpay_key_secret))
        else:
            self.razorpay_client = None
            self.logger.warning("Razorpay credentials not provided - payment processing disabled")
        
        # Initialize subscription tables
        self._init_subscription_tables()
        
        self.logger.info("SubscriptionManager initialized")
    
    def _init_subscription_tables(self):
        """Initialize subscription-related database tables."""
        try:
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            cursor = conn.cursor()
            
            # Create subscriptions table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS user_subscriptions (
                    id TEXT PRIMARY KEY,
                    user_id TEXT NOT NULL,
                    plan TEXT NOT NULL,
                    status TEXT DEFAULT 'active',
                    current_period_start TEXT NOT NULL,
                    current_period_end TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    razorpay_subscription_id TEXT,
                    razorpay_customer_id TEXT,
                    trial_end TEXT,
                    cancelled_at TEXT,
                    FOREIGN KEY (user_id) REFERENCES users (id)
                )
            ''')
            
            # Create payments table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS subscription_payments (
                    id TEXT PRIMARY KEY,
                    user_id TEXT NOT NULL,
                    subscription_id TEXT NOT NULL,
                    razorpay_payment_id TEXT NOT NULL,
                    razorpay_order_id TEXT,
                    amount INTEGER NOT NULL,
                    currency TEXT DEFAULT 'INR',
                    status TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    paid_at TEXT,
                    failure_reason TEXT,
                    FOREIGN KEY (user_id) REFERENCES users (id),
                    FOREIGN KEY (subscription_id) REFERENCES user_subscriptions (id)
                )
            ''')
            
            # Create usage tracking table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS usage_tracking (
                    id TEXT PRIMARY KEY,
                    user_id TEXT NOT NULL,
                    resource_type TEXT NOT NULL,
                    resource_count INTEGER DEFAULT 1,
                    month_year TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    FOREIGN KEY (user_id) REFERENCES users (id)
                )
            ''')
            
            # Create indexes
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_subscriptions_user_id ON user_subscriptions(user_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_subscriptions_status ON user_subscriptions(status)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_payments_user_id ON subscription_payments(user_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_usage_user_month ON usage_tracking(user_id, month_year)')
            
            conn.commit()
            conn.close()
            
            self.logger.info("Subscription tables initialized")
            
        except Exception as e:
            self.logger.error(f"Failed to initialize subscription tables: {e}")
            raise
    
    def get_user_subscription(self, user_id: str) -> Dict[str, Any]:
        """Get user's current subscription details."""
        try:
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM user_subscriptions 
                WHERE user_id = ? AND status = 'active'
                ORDER BY current_period_end DESC
                LIMIT 1
            ''', (user_id,))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                subscription = dict(row)
                subscription['plan_config'] = PLAN_CONFIGS[SubscriptionPlan(subscription['plan'])]
                return subscription
            else:
                # Return default free plan
                return {
                    'user_id': user_id,
                    'plan': SubscriptionPlan.FREE.value,
                    'status': 'active',
                    'plan_config': PLAN_CONFIGS[SubscriptionPlan.FREE]
                }
                
        except Exception as e:
            self.logger.error(f"Failed to get user subscription: {e}")
            return {
                'user_id': user_id,
                'plan': SubscriptionPlan.FREE.value,
                'status': 'active',
                'plan_config': PLAN_CONFIGS[SubscriptionPlan.FREE]
            }
    
    def create_payment_order(self, user_id: str, plan: SubscriptionPlan) -> Dict[str, Any]:
        """Create a Razorpay payment order for subscription."""
        if not self.razorpay_client:
            raise SecurityError("Payment processing not configured")
        
        if plan == SubscriptionPlan.ENTERPRISE:
            raise SecurityError("Enterprise plans require custom pricing - contact sales")
        
        try:
            amount = PLAN_PRICING[plan]
            if amount == 0:
                raise SecurityError("Free plans don't require payment")
            
            # Create Razorpay order
            # Generate short receipt (max 40 chars) - use first 8 chars of user_id + timestamp
            user_id_short = user_id[:8] if len(user_id) > 8 else user_id
            timestamp = datetime.now().strftime("%m%d%H%M%S")  # Short timestamp
            receipt = f"sub_{user_id_short}_{timestamp}"
            
            # Ensure receipt doesn't exceed 40 characters (Razorpay limit)
            if len(receipt) > 40:
                receipt = receipt[:40]
            
            order_data = {
                'amount': amount * 100,  # Amount in paise
                'currency': 'INR',
                'receipt': receipt,  # Ensures max 40 chars: sub_12345678_0626142530 = 22 chars
                'notes': {
                    'user_id': user_id,
                    'plan': plan.value,
                    'type': 'subscription'
                }
            }
            
            razorpay_order = self.razorpay_client.order.create(data=order_data)
            
            # Store payment record
            payment_id = str(uuid.uuid4())
            now = datetime.now().isoformat()
            
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO subscription_payments 
                (id, user_id, subscription_id, razorpay_payment_id, razorpay_order_id, amount, status, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', (payment_id, user_id, '', razorpay_order['id'], razorpay_order['id'], amount, PaymentStatus.PENDING.value, now))
            
            conn.commit()
            conn.close()
            
            return {
                'order_id': razorpay_order['id'],
                'amount': amount,
                'currency': 'INR',
                'payment_id': payment_id,
                'key': self.razorpay_client.auth[0]  # Razorpay key ID for frontend
            }
            
        except Exception as e:
            self.logger.error(f"Failed to create payment order: {e}")
            raise SecurityError(f"Payment order creation failed: {e}")
    
    def verify_payment_and_activate_subscription(self, payment_id: str, razorpay_payment_id: str, razorpay_signature: str) -> Dict[str, Any]:
        """Verify payment and activate subscription."""
        if not self.razorpay_client:
            raise SecurityError("Payment processing not configured")
        
        try:
            # Get payment record
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            cursor.execute('SELECT * FROM subscription_payments WHERE id = ?', (payment_id,))
            payment_record = cursor.fetchone()
            
            if not payment_record:
                raise SecurityError("Payment record not found")
            
            # Verify payment with Razorpay
            try:
                payment_details = self.razorpay_client.payment.fetch(razorpay_payment_id)
                
                # Verify signature
                params_dict = {
                    'razorpay_order_id': payment_record['razorpay_order_id'],
                    'razorpay_payment_id': razorpay_payment_id,
                    'razorpay_signature': razorpay_signature
                }
                
                self.razorpay_client.utility.verify_payment_signature(params_dict)
                
                if payment_details['status'] in ['captured', 'authorized']:
                    # Payment successful - activate subscription (captured = live mode, authorized = test mode)
                    now = datetime.now()
                    period_end = now + timedelta(days=30)  # Monthly subscription
                    
                    # Determine plan from payment notes
                    order_details = self.razorpay_client.order.fetch(payment_record['razorpay_order_id'])
                    plan = order_details['notes']['plan']
                    user_id = order_details['notes']['user_id']
                    
                    # Create subscription record
                    subscription_id = str(uuid.uuid4())
                    cursor.execute('''
                        INSERT INTO user_subscriptions 
                        (id, user_id, plan, status, current_period_start, current_period_end, created_at, updated_at, razorpay_subscription_id)
                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ''', (subscription_id, user_id, plan, 'active', now.isoformat(), period_end.isoformat(), now.isoformat(), now.isoformat(), razorpay_payment_id))
                    
                    # Update payment record
                    cursor.execute('''
                        UPDATE subscription_payments 
                        SET status = ?, paid_at = ?, subscription_id = ?, razorpay_payment_id = ?
                        WHERE id = ?
                    ''', (PaymentStatus.SUCCESS.value, now.isoformat(), subscription_id, razorpay_payment_id, payment_id))
                    
                    # Update user subscription field
                    cursor.execute('UPDATE users SET subscription = ? WHERE id = ?', (plan, user_id))
                    
                    conn.commit()
                    conn.close()
                    
                    self.logger.info(f"Subscription activated for user {user_id} - plan: {plan}")
                    
                    return {
                        'success': True,
                        'subscription_id': subscription_id,
                        'plan': plan,
                        'period_end': period_end.isoformat()
                    }
                else:
                    raise SecurityError(f"Payment not successful: {payment_details['status']}. Expected 'captured' or 'authorized'.")
                    
            except Exception as e:
                # Update payment record as failed
                cursor.execute('''
                    UPDATE subscription_payments 
                    SET status = ?, failure_reason = ?
                    WHERE id = ?
                ''', (PaymentStatus.FAILED.value, str(e), payment_id))
                
                conn.commit()
                conn.close()
                
                raise SecurityError(f"Payment verification failed: {e}")
                
        except Exception as e:
            self.logger.error(f"Payment verification failed: {e}")
            raise SecurityError(f"Payment verification failed: {e}")
    
    def check_feature_access(self, user_id: str, feature: str) -> bool:
        """Check if user has access to a specific feature."""
        subscription = self.get_user_subscription(user_id)
        plan_config = subscription['plan_config']
        
        # Map features to plan attributes
        feature_mapping = {
            'advanced_training': 'advanced_training',
            'priority_support': 'priority_support',
            'private_layers': 'private_layers',
            'collaborative_workspaces': 'collaborative_workspaces',
            'custom_integrations': 'custom_integrations',
            'analytics_dashboard': 'analytics_dashboard',
            'advanced_api_access': 'advanced_api_access',
            'custom_model_integrations': 'custom_model_integrations',
            'dedicated_infrastructure': 'dedicated_infrastructure',
            'white_label': 'white_label',
            'dedicated_account_manager': 'dedicated_account_manager'
        }
        
        if feature in feature_mapping:
            return getattr(plan_config, feature_mapping[feature], False)
        
        return False
    
    def check_usage_limits(self, user_id: str, resource_type: str) -> Dict[str, Any]:
        """Check usage against plan limits."""
        subscription = self.get_user_subscription(user_id)
        plan_config = subscription['plan_config']
        
        # Get current month usage
        current_month = datetime.now().strftime("%Y-%m")
        
        try:
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            if resource_type == 'layers':
                # Count user's layers
                cursor.execute('SELECT COUNT(*) as count FROM user_layers WHERE user_id = ?', (user_id,))
                current_usage = cursor.fetchone()['count']
                limit = plan_config.max_layers
                
            elif resource_type == 'training_sessions':
                # Count training sessions this month
                cursor.execute('''
                    SELECT COALESCE(SUM(resource_count), 0) as count 
                    FROM usage_tracking 
                    WHERE user_id = ? AND resource_type = ? AND month_year = ?
                ''', (user_id, resource_type, current_month))
                current_usage = cursor.fetchone()['count']
                limit = plan_config.max_training_sessions_per_month
                
            else:
                current_usage = 0
                limit = None
            
            conn.close()
            
            return {
                'current_usage': current_usage,
                'limit': limit,
                'has_access': limit is None or current_usage < limit,
                'usage_percentage': 0 if limit is None else min(100, (current_usage / limit) * 100)
            }
            
        except Exception as e:
            self.logger.error(f"Failed to check usage limits: {e}")
            return {'current_usage': 0, 'limit': None, 'has_access': True, 'usage_percentage': 0}
    
    def track_usage(self, user_id: str, resource_type: str, count: int = 1):
        """Track resource usage for billing/limiting."""
        try:
            current_month = datetime.now().strftime("%Y-%m")
            now = datetime.now().isoformat()
            
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            cursor = conn.cursor()
            
            # Check if record exists for this month
            cursor.execute('''
                SELECT id, resource_count FROM usage_tracking 
                WHERE user_id = ? AND resource_type = ? AND month_year = ?
            ''', (user_id, resource_type, current_month))
            
            existing = cursor.fetchone()
            
            if existing:
                # Update existing record
                new_count = existing[1] + count
                cursor.execute('''
                    UPDATE usage_tracking 
                    SET resource_count = ?
                    WHERE id = ?
                ''', (new_count, existing[0]))
            else:
                # Create new record
                cursor.execute('''
                    INSERT INTO usage_tracking (id, user_id, resource_type, resource_count, month_year, created_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (str(uuid.uuid4()), user_id, resource_type, count, current_month, now))
            
            conn.commit()
            conn.close()
            
        except Exception as e:
            self.logger.error(f"Failed to track usage: {e}")
    
    def cancel_subscription(self, user_id: str) -> Dict[str, Any]:
        """Cancel user's subscription."""
        try:
            import sqlite3
            conn = sqlite3.connect(self.db.db_path)
            cursor = conn.cursor()
            
            now = datetime.now().isoformat()
            
            # Mark subscription as cancelled
            cursor.execute('''
                UPDATE user_subscriptions 
                SET status = 'cancelled', cancelled_at = ?, updated_at = ?
                WHERE user_id = ? AND status = 'active'
            ''', ('cancelled', now, now, user_id))
            
            # Update user subscription to free
            cursor.execute('UPDATE users SET subscription = ? WHERE id = ?', (SubscriptionPlan.FREE.value, user_id))
            
            conn.commit()
            conn.close()
            
            self.logger.info(f"Subscription cancelled for user {user_id}")
            
            return {'success': True, 'message': 'Subscription cancelled successfully'}
            
        except Exception as e:
            self.logger.error(f"Failed to cancel subscription: {e}")
            raise SecurityError(f"Subscription cancellation failed: {e}")
    
    def get_subscription_analytics(self, user_id: str) -> Dict[str, Any]:
        """Get subscription and usage analytics for user."""
        subscription = self.get_user_subscription(user_id)
        
        # Get usage stats
        layers_usage = self.check_usage_limits(user_id, 'layers')
        training_usage = self.check_usage_limits(user_id, 'training_sessions')
        
        return {
            'subscription': subscription,
            'usage': {
                'layers': layers_usage,
                'training_sessions': training_usage
            },
            'features': {
                'advanced_training': self.check_feature_access(user_id, 'advanced_training'),
                'priority_support': self.check_feature_access(user_id, 'priority_support'),
                'private_layers': self.check_feature_access(user_id, 'private_layers'),
                'analytics_dashboard': self.check_feature_access(user_id, 'analytics_dashboard'),
                'custom_integrations': self.check_feature_access(user_id, 'custom_integrations')
            }
        }


# Global subscription manager instance
_subscription_manager = None

def get_subscription_manager() -> SubscriptionManager:
    """Get the global subscription manager instance."""
    global _subscription_manager
    if _subscription_manager is None:
        import os
        from dotenv import load_dotenv
        
        # Load environment variables from .env file
        load_dotenv()
        
        # Load from environment variables
        razorpay_key_id = os.getenv('RAZORPAY_KEY_ID')
        razorpay_key_secret = os.getenv('RAZORPAY_KEY_SECRET')
        
        _subscription_manager = SubscriptionManager(
            razorpay_key_id=razorpay_key_id,
            razorpay_key_secret=razorpay_key_secret
        )
    return _subscription_manager 