"""
Report what the AI chat has actually cost, using the token counts already
stored on every assistant message.

    python manage.py chat_usage              # last 30 days
    python manage.py chat_usage --days 90
    python manage.py chat_usage --by-session # worst offenders first
"""
from datetime import timedelta

from django.core.management.base import BaseCommand
from django.db.models import Sum, Count
from django.utils import timezone

from apps.chat.models import ChatMessage, ChatSession

# USD per 1M tokens. Blended rate used because we store a single total per
# message rather than an input/output split.
PRICING = {
    'gpt-4o':      4.00,
    'gpt-4o-mini': 0.25,
}


class Command(BaseCommand):
    help = 'Show chat token usage and estimated spend.'

    def add_arguments(self, parser):
        parser.add_argument('--days', type=int, default=30)
        parser.add_argument('--by-session', action='store_true',
                            help='List the most expensive sessions.')

    def handle(self, *args, **options):
        from apps.chat.openai_client import MODEL

        since = timezone.now() - timedelta(days=options['days'])
        qs = ChatMessage.objects.filter(created_at__gte=since, tokens__gt=0)

        totals = qs.aggregate(t=Sum('tokens'), n=Count('id'))
        tokens = totals['t'] or 0
        calls  = totals['n'] or 0

        rate = PRICING.get(MODEL, PRICING['gpt-4o'])
        cost = tokens / 1_000_000 * rate

        self.stdout.write(f'Window:         last {options["days"]} days')
        self.stdout.write(f'Model in use:   {MODEL}')
        self.stdout.write(f'API calls:      {calls}')
        self.stdout.write(f'Total tokens:   {tokens:,}')
        if calls:
            self.stdout.write(f'Avg per reply:  {tokens // calls:,} tokens')
        self.stdout.write(self.style.SUCCESS(f'Estimated cost: about ${cost:,.2f}'))

        cheap = tokens / 1_000_000 * PRICING['gpt-4o-mini']
        if MODEL != 'gpt-4o-mini':
            self.stdout.write(f'Same traffic on gpt-4o-mini: about ${cheap:,.2f}')

        if options['by_session']:
            self.stdout.write('')
            self.stdout.write('Most expensive sessions:')
            rows = (
                ChatSession.objects
                .filter(messages__created_at__gte=since)
                .annotate(t=Sum('messages__tokens'), n=Count('messages'))
                .order_by('-t')[:15]
            )
            for s in rows:
                t = s.t or 0
                self.stdout.write(
                    f'  #{s.pk:<5} {(s.topic or s.visitor_name or "Anonymous")[:24]:<26} '
                    f'{s.n:>3} msgs  {t:>8,} tokens  ~${t / 1_000_000 * rate:,.2f}'
                )