token_system.py

#!/usr/bin/env python3
"""
VYNL Token & User System
- Demo mode: 3 free tokens, 5-min track limit
- Licensed mode: 300 tokens/month, full access
- Admin token distribution via simple text file
"""

import json
import hashlib
import os
from pathlib import Path
from datetime import datetime, timedelta
from typing import Optional, Tuple, Dict

# ============================================================================
# CONFIGURATION
# ============================================================================

DATA_DIR = Path(os.environ.get('VYNL_DATA_DIR', Path.home() / '.vynl_data'))
DATA_DIR.mkdir(parents=True, exist_ok=True)

USERS_FILE = DATA_DIR / 'users.json'
TOKENS_FILE = DATA_DIR / 'token_grants.txt'
SESSIONS_FILE = DATA_DIR / 'sessions.json'

# Token costs
TOKEN_COSTS = {
    'song_analysis': 1,      # Full stem + chord + DAW
    'stem_only': 1,          # Just stems
    'chord_only': 1,         # Just chords
    'ai_generate': 2,        # GROOVES generation
    'bulk_song': 1,          # Per song in bulk
}

# Limits
DEMO_TOKENS = 3
DEMO_MAX_DURATION = 300      # 5 minutes in seconds
LICENSED_MONTHLY_TOKENS = 300
LICENSED_MAX_DURATION = None  # Unlimited

# ============================================================================
# VALID LICENSES (from license_system)
# ============================================================================

VALID_LICENSES = {
    # Creator licenses - unlimited
    "VYNL-IY2M-KV47-AT7J-C74V": {"name": "R.T. Lackey", "email": "rlackey.seattle@gmail.com", "type": "CREATOR", "unlimited": True},
    "VYNL-INZW-JNZY-Y4O2-WOEB": {"name": "R.T. Lackey", "email": "rlackey.seattle@gmail.com", "type": "CREATOR", "unlimited": True},

    # Universal Licenses - 300 tokens/month, no duration limit
    # These can be distributed to users
    "VYNL-UNIV-2026-ALPHA-001A": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-ALPHA-002B": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-ALPHA-003C": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-ALPHA-004D": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-ALPHA-005E": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-BETA-001F": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-BETA-002G": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-BETA-003H": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-BETA-004J": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-UNIV-2026-BETA-005K": {"name": "Universal License", "type": "PROFESSIONAL", "tokens": 300},

    # Premium Licenses - 1000 tokens/month
    "VYNL-PREM-2026-GOLD-001A": {"name": "Premium License", "type": "PREMIUM", "tokens": 1000},
    "VYNL-PREM-2026-GOLD-002B": {"name": "Premium License", "type": "PREMIUM", "tokens": 1000},
    "VYNL-PREM-2026-GOLD-003C": {"name": "Premium License", "type": "PREMIUM", "tokens": 1000},
    "VYNL-PREM-2026-GOLD-004D": {"name": "Premium License", "type": "PREMIUM", "tokens": 1000},
    "VYNL-PREM-2026-GOLD-005E": {"name": "Premium License", "type": "PREMIUM", "tokens": 1000},

    # Demo licenses (original 15)
    "VYNL-WAFV-HBGQ-UMAY-UKRD": {"name": "Demo User 01", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-5M73-VSUB-CP5L-PABM": {"name": "Demo User 02", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-VURV-P5NN-N2IK-EV44": {"name": "Demo User 03", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-7TH6-NWHM-LNC2-KMG7": {"name": "Demo User 04", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-4W2G-NYRK-LDW7-554E": {"name": "Demo User 05", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-GGAD-AMOO-TLVQ-5O6M": {"name": "Demo User 06", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-PJM4-PRRG-AID3-VFEA": {"name": "Demo User 07", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-G45E-OBGJ-7LB6-3BKZ": {"name": "Demo User 08", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-WT7Y-ICDE-WN43-SU4B": {"name": "Demo User 09", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-J3DM-Y2KY-GLTN-PNM4": {"name": "Demo User 10", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-3FVE-RTMT-LAOJ-NH3P": {"name": "Demo User 11", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-YOS6-LESJ-WGIB-AOVM": {"name": "Demo User 12", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-ST6S-4GUY-WXVL-JWM6": {"name": "Demo User 13", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-RFRG-YUXL-7AX4-7FPY": {"name": "Demo User 14", "type": "PROFESSIONAL", "tokens": 300},
    "VYNL-54HA-343P-V5AT-6RJL": {"name": "Demo User 15", "type": "PROFESSIONAL", "tokens": 300},
}

# ============================================================================
# USER DATABASE
# ============================================================================

def load_users() -> Dict:
    """Load user database"""
    if USERS_FILE.exists():
        return json.loads(USERS_FILE.read_text())
    return {}

def save_users(users: Dict):
    """Save user database"""
    USERS_FILE.write_text(json.dumps(users, indent=2))

def hash_password(password: str) -> str:
    """Hash password for storage"""
    return hashlib.sha256(password.encode()).hexdigest()

def get_month_key() -> str:
    """Get current month key for token tracking"""
    return datetime.now().strftime('%Y-%m')

# ============================================================================
# TOKEN GRANTS FILE
# ============================================================================

def load_token_grants() -> Dict[str, int]:
    """
    Load token grants from simple text file
    Format: email,tokens (one per line)
    Example:
        john@example.com,100
        jane@example.com,50
    """
    grants = {}
    if TOKENS_FILE.exists():
        for line in TOKENS_FILE.read_text().strip().split('\n'):
            line = line.strip()
            if line and ',' in line and not line.startswith('#'):
                parts = line.split(',')
                if len(parts) >= 2:
                    email = parts[0].strip().lower()
                    try:
                        tokens = int(parts[1].strip())
                        grants[email] = grants.get(email, 0) + tokens
                    except ValueError:
                        pass
    return grants

# ============================================================================
# USER MANAGEMENT
# ============================================================================

class UserManager:
    def __init__(self):
        self.users = load_users()
        self.token_grants = load_token_grants()

    def reload_grants(self):
        """Reload token grants from file"""
        self.token_grants = load_token_grants()

    def create_account(self, email: str, password: str, name: str = "") -> Tuple[bool, str]:
        """Create new user account"""
        email = email.strip().lower()

        if not email or '@' not in email:
            return False, "Invalid email address"

        if not password or len(password) < 6:
            return False, "Password must be at least 6 characters"

        if email in self.users:
            return False, "Account already exists"

        self.users[email] = {
            'email': email,
            'name': name or email.split('@')[0],
            'password_hash': hash_password(password),
            'created': datetime.now().isoformat(),
            'license_key': None,
            'license_type': 'DEMO',
            'tokens_used': {},  # {month_key: count}
            'bonus_tokens': 0,
            'total_songs_processed': 0,
        }

        save_users(self.users)
        return True, "Account created successfully"

    def login(self, email: str, password: str) -> Tuple[bool, Optional[Dict]]:
        """Login user"""
        email = email.strip().lower()

        if email not in self.users:
            return False, None

        user = self.users[email]
        if user['password_hash'] != hash_password(password):
            return False, None

        return True, user

    def activate_license(self, email: str, license_key: str) -> Tuple[bool, str]:
        """Activate license for user"""
        email = email.strip().lower()
        license_key = license_key.strip().upper()

        if email not in self.users:
            return False, "User not found"

        if license_key not in VALID_LICENSES:
            return False, "Invalid license key"

        license_info = VALID_LICENSES[license_key]

        self.users[email]['license_key'] = license_key
        self.users[email]['license_type'] = license_info['type']
        self.users[email]['license_activated'] = datetime.now().isoformat()

        save_users(self.users)
        return True, f"License activated: {license_info['type']}"

    def get_user_status(self, email: str) -> Dict:
        """Get complete user status including tokens"""
        email = email.strip().lower()

        if email not in self.users:
            # Demo user (not registered)
            return {
                'registered': False,
                'license_type': 'DEMO',
                'tokens_remaining': DEMO_TOKENS,
                'tokens_used': 0,
                'max_duration': DEMO_MAX_DURATION,
                'unlimited': False,
            }

        user = self.users[email]
        month_key = get_month_key()
        tokens_used_this_month = user['tokens_used'].get(month_key, 0)

        # Check for bonus tokens from grants file
        self.reload_grants()
        bonus_from_grants = self.token_grants.get(email, 0)

        # Calculate based on license type
        if user['license_type'] == 'CREATOR':
            return {
                'registered': True,
                'email': email,
                'name': user['name'],
                'license_type': 'CREATOR',
                'tokens_remaining': 999999,
                'tokens_used': tokens_used_this_month,
                'max_duration': None,
                'unlimited': True,
            }
        elif user['license_key']:
            # Licensed user
            base_tokens = LICENSED_MONTHLY_TOKENS
            total_available = base_tokens + user.get('bonus_tokens', 0) + bonus_from_grants
            tokens_remaining = max(0, total_available - tokens_used_this_month)

            return {
                'registered': True,
                'email': email,
                'name': user['name'],
                'license_type': user['license_type'],
                'tokens_remaining': tokens_remaining,
                'tokens_used': tokens_used_this_month,
                'monthly_limit': base_tokens,
                'bonus_tokens': user.get('bonus_tokens', 0) + bonus_from_grants,
                'max_duration': LICENSED_MAX_DURATION,
                'unlimited': False,
            }
        else:
            # Registered but no license (demo)
            return {
                'registered': True,
                'email': email,
                'name': user['name'],
                'license_type': 'DEMO',
                'tokens_remaining': max(0, DEMO_TOKENS - tokens_used_this_month),
                'tokens_used': tokens_used_this_month,
                'max_duration': DEMO_MAX_DURATION,
                'unlimited': False,
            }

    def use_tokens(self, email: str, amount: int, action: str = 'song_analysis') -> Tuple[bool, str]:
        """Deduct tokens for an action"""
        email = email.strip().lower()
        status = self.get_user_status(email)

        if status['unlimited']:
            return True, "Unlimited access"

        if status['tokens_remaining'] < amount:
            return False, f"Insufficient tokens. Need {amount}, have {status['tokens_remaining']}"

        # Deduct tokens
        if email in self.users:
            month_key = get_month_key()
            if month_key not in self.users[email]['tokens_used']:
                self.users[email]['tokens_used'][month_key] = 0
            self.users[email]['tokens_used'][month_key] += amount
            self.users[email]['total_songs_processed'] += 1
            save_users(self.users)

        remaining = status['tokens_remaining'] - amount
        return True, f"Token used. {remaining} remaining"

    def check_duration_limit(self, email: str, duration_seconds: float) -> Tuple[bool, str]:
        """Check if track duration is within limits"""
        status = self.get_user_status(email)

        if status['max_duration'] is None:
            return True, "No duration limit"

        if duration_seconds > status['max_duration']:
            max_mins = status['max_duration'] // 60
            return False, f"Track exceeds {max_mins}-minute limit for demo mode. Upgrade to process longer tracks."

        return True, "Duration OK"

    def add_bonus_tokens(self, email: str, amount: int) -> Tuple[bool, str]:
        """Add bonus tokens to user account"""
        email = email.strip().lower()

        if email not in self.users:
            return False, "User not found"

        self.users[email]['bonus_tokens'] = self.users[email].get('bonus_tokens', 0) + amount
        save_users(self.users)

        return True, f"Added {amount} bonus tokens"


# ============================================================================
# SINGLETON INSTANCE
# ============================================================================

user_manager = UserManager()


# ============================================================================
# HELPER FUNCTIONS FOR GRADIO
# ============================================================================

def check_can_process(email: str, duration_seconds: float = 0) -> Tuple[bool, str, Dict]:
    """
    Check if user can process a song
    Returns: (can_process, message, status_dict)
    """
    status = user_manager.get_user_status(email)

    # Check tokens
    if status['tokens_remaining'] <= 0 and not status['unlimited']:
        return False, "No tokens remaining. Please upgrade or wait for monthly reset.", status

    # Check duration
    if duration_seconds > 0:
        ok, msg = user_manager.check_duration_limit(email, duration_seconds)
        if not ok:
            return False, msg, status

    return True, "Ready to process", status


def deduct_token(email: str) -> Tuple[bool, str]:
    """Deduct one token after successful processing"""
    return user_manager.use_tokens(email, 1)


def get_status_display(email: str) -> str:
    """Get formatted status for UI display"""
    if not email:
        return "DEMO MODE: 3 free tokens | 5-min track limit | Enter email to track usage"

    status = user_manager.get_user_status(email)

    if status['unlimited']:
        return f"CREATOR: {status['name']} | UNLIMITED ACCESS"

    if status['license_type'] != 'DEMO':
        return f"LICENSED ({status['license_type']}): {status['tokens_remaining']} tokens remaining this month"

    return f"DEMO: {status['tokens_remaining']}/{DEMO_TOKENS} tokens | 5-min limit | Upgrade for full access"


# ============================================================================
# CLI FOR TESTING
# ============================================================================

if __name__ == "__main__":
    import sys

    print("VYNL Token System")
    print("=" * 50)

    if len(sys.argv) < 2:
        print("""
Commands:
    status <email>           Check user status
    create <email> <pass>    Create account
    grant <email> <tokens>   Add tokens to grants file
    activate <email> <key>   Activate license
        """)
        sys.exit(0)

    cmd = sys.argv[1]

    if cmd == "status" and len(sys.argv) >= 3:
        email = sys.argv[2]
        status = user_manager.get_user_status(email)
        print(json.dumps(status, indent=2))

    elif cmd == "create" and len(sys.argv) >= 4:
        email, password = sys.argv[2], sys.argv[3]
        ok, msg = user_manager.create_account(email, password)
        print(f"{'Success' if ok else 'Failed'}: {msg}")

    elif cmd == "grant" and len(sys.argv) >= 4:
        email, tokens = sys.argv[2], sys.argv[3]
        # Append to grants file
        with open(TOKENS_FILE, 'a') as f:
            f.write(f"{email},{tokens}\n")
        print(f"Added {tokens} tokens for {email}")

    elif cmd == "activate" and len(sys.argv) >= 4:
        email, key = sys.argv[2], sys.argv[3]
        ok, msg = user_manager.activate_license(email, key)
        print(f"{'Success' if ok else 'Failed'}: {msg}")