"""
Backfill token data from AITaskLog to CreditUsageLog
"""
from django.core.management.base import BaseCommand
from igny8_core.ai.models import AITaskLog
from igny8_core.modules.billing.models import CreditUsageLog
from datetime import timedelta


class Command(BaseCommand):
    help = 'Backfill token data from AITaskLog to CreditUsageLog'

    def handle(self, *args, **options):
        self.stdout.write("=== Token Data Backfill ===\n")
        
        # Get AITaskLog entries with token data
        ai_logs = AITaskLog.objects.filter(
            tokens__gt=0,
            status='success'
        ).select_related('account').order_by('created_at')
        
        self.stdout.write(f"Found {ai_logs.count()} AITaskLog entries with tokens\n")
        
        updated_count = 0
        skipped_count = 0
        no_match_count = 0
        
        for ai_log in ai_logs:
            # Find matching CreditUsageLog within 10 second window
            time_start = ai_log.created_at - timedelta(seconds=10)
            time_end = ai_log.created_at + timedelta(seconds=10)
            
            # Try to find exact match
            credit_log = CreditUsageLog.objects.filter(
                account=ai_log.account,
                created_at__gte=time_start,
                created_at__lte=time_end,
                tokens_input=0,
                tokens_output=0
            ).order_by('created_at').first()
            
            if credit_log:
                # AITaskLog has total tokens, estimate 40/60 split for input/output
                # This is approximate but better than 0
                total_tokens = ai_log.tokens
                estimated_input = int(total_tokens * 0.4)
                estimated_output = total_tokens - estimated_input
                
                credit_log.tokens_input = estimated_input
                credit_log.tokens_output = estimated_output
                credit_log.save(update_fields=['tokens_input', 'tokens_output'])
                
                updated_count += 1
                if updated_count % 50 == 0:
                    self.stdout.write(f"  Updated {updated_count} records...")
            else:
                no_match_count += 1
        
        self.stdout.write(self.style.SUCCESS(f"\n✅ Backfill complete!"))
        self.stdout.write(f"  Updated: {updated_count}")
        self.stdout.write(f"  No match: {no_match_count}")
        self.stdout.write(f"  Total processed: {ai_logs.count()}")
