# Generated by Django 6.0.2 on 2026-02-22

from decimal import Decimal

from django.db import migrations, models


def populate_weekly_activity_cost(apps, schema_editor):
    """Populate cost field on WeeklyActivity from SessionItem costs.

    Sums costs from all SessionItems (not just user_message) grouped by week,
    including all session types (sessions and subagents). Updates both per-project
    and global (project=NULL) rows.
    """
    from django.db.models import Sum
    from django.db.models.functions import TruncWeek

    SessionItem = apps.get_model("core", "SessionItem")
    WeeklyActivity = apps.get_model("core", "WeeklyActivity")

    # Compute costs per project per week from SessionItems (all session types)
    rows = (
        SessionItem.objects.filter(
            cost__isnull=False,
            timestamp__isnull=False,
        )
        .values(
            project_id=models.F("session__project_id"),
            week_start=TruncWeek("timestamp"),
        )
        .annotate(total_cost=Sum("cost"))
    )

    global_costs = {}  # week -> Decimal cost

    for row in rows:
        week = row["week_start"].date()
        project_id = row["project_id"]
        cost = row["total_cost"] or Decimal(0)

        if not cost:
            continue

        # Update per-project row
        WeeklyActivity.objects.filter(
            project_id=project_id, date=week
        ).update(cost=cost)

        # If no per-project row exists (cost-only items in a week with no user messages),
        # create one
        if not WeeklyActivity.objects.filter(project_id=project_id, date=week).exists():
            WeeklyActivity.objects.create(
                project_id=project_id, date=week, count=0, cost=cost
            )

        global_costs[week] = global_costs.get(week, Decimal(0)) + cost

    # Update global rows
    for week, cost in global_costs.items():
        updated = WeeklyActivity.objects.filter(
            project__isnull=True, date=week
        ).update(cost=cost)
        if not updated:
            WeeklyActivity.objects.create(
                project=None, date=week, count=0, cost=cost
            )


def populate_daily_activity_cost(apps, schema_editor):
    """Populate cost field on DailyActivity from SessionItem costs.

    Sums costs from all SessionItems (not just user_message) grouped by day,
    including all session types (sessions and subagents). Updates both per-project
    and global (project=NULL) rows.
    """
    from django.db.models import Sum
    from django.db.models.functions import TruncDate

    SessionItem = apps.get_model("core", "SessionItem")
    DailyActivity = apps.get_model("core", "DailyActivity")

    # Compute costs per project per day from SessionItems (all session types)
    rows = (
        SessionItem.objects.filter(
            cost__isnull=False,
            timestamp__isnull=False,
        )
        .values(
            project_id=models.F("session__project_id"),
            day=TruncDate("timestamp"),
        )
        .annotate(total_cost=Sum("cost"))
    )

    global_costs = {}  # date -> Decimal cost

    for row in rows:
        day = row["day"]
        if hasattr(day, "date"):
            day = day.date()
        project_id = row["project_id"]
        cost = row["total_cost"] or Decimal(0)

        if not cost:
            continue

        # Update per-project row
        DailyActivity.objects.filter(
            project_id=project_id, date=day
        ).update(cost=cost)

        # If no per-project row exists (cost-only items in a day with no user messages),
        # create one
        if not DailyActivity.objects.filter(project_id=project_id, date=day).exists():
            DailyActivity.objects.create(
                project_id=project_id, date=day, count=0, cost=cost
            )

        global_costs[day] = global_costs.get(day, Decimal(0)) + cost

    # Update global rows
    for day, cost in global_costs.items():
        updated = DailyActivity.objects.filter(
            project__isnull=True, date=day
        ).update(cost=cost)
        if not updated:
            DailyActivity.objects.create(
                project=None, date=day, count=0, cost=cost
            )


class Migration(migrations.Migration):

    dependencies = [
        ('core', '0030_recompute_activity_exclude_subagents'),
    ]

    operations = [
        # Rename week -> date on WeeklyActivity (to match abstract PeriodicActivity)
        migrations.RenameField(
            model_name='weeklyactivity',
            old_name='week',
            new_name='date',
        ),
        # Update unique constraint name (week -> date)
        migrations.RemoveConstraint(
            model_name='weeklyactivity',
            name='unique_project_week',
        ),
        migrations.AddConstraint(
            model_name='weeklyactivity',
            constraint=models.UniqueConstraint(
                fields=['project', 'date'],
                name='unique_project_weeklyactivity',
            ),
        ),
        # Update unique constraint name on DailyActivity
        migrations.RemoveConstraint(
            model_name='dailyactivity',
            name='unique_project_date',
        ),
        migrations.AddConstraint(
            model_name='dailyactivity',
            constraint=models.UniqueConstraint(
                fields=['project', 'date'],
                name='unique_project_dailyactivity',
            ),
        ),
        # Update related_name on project FK
        migrations.AlterField(
            model_name='weeklyactivity',
            name='project',
            field=models.ForeignKey(
                blank=True,
                null=True,
                on_delete=models.deletion.CASCADE,
                related_name='%(class)s_set',
                to='core.project',
            ),
        ),
        migrations.AlterField(
            model_name='dailyactivity',
            name='project',
            field=models.ForeignKey(
                blank=True,
                null=True,
                on_delete=models.deletion.CASCADE,
                related_name='%(class)s_set',
                to='core.project',
            ),
        ),
        # Add cost field to WeeklyActivity
        migrations.AddField(
            model_name='weeklyactivity',
            name='cost',
            field=models.DecimalField(decimal_places=6, default=0, max_digits=12),
        ),
        # Add cost field to DailyActivity
        migrations.AddField(
            model_name='dailyactivity',
            name='cost',
            field=models.DecimalField(decimal_places=6, default=0, max_digits=12),
        ),
        # Populate cost data from existing SessionItems
        migrations.RunPython(populate_weekly_activity_cost, migrations.RunPython.noop),
        migrations.RunPython(populate_daily_activity_cost, migrations.RunPython.noop),
    ]
