# Generated by Django 5.2.4 on 2025-07-04 14:25
import uuid
from django.conf import settings
from django.db import migrations, models
import django.db.models.deletion
import psql_partition.manager.manager
import psql_partition.backend.migrations.operations.create_partitioned_model
import psql_partition.models.partitioned
import psql_partition.types
from glitchtip.model_utils import TestDefaultPartition

# --- SQL Definitions for the Advanced Path ---

# SQL for TransactionEvent
CREATE_ADVANCED_TRANSACTION_EVENT_SQL = """
    CREATE TABLE "performance_transactionevent" (
        "event_id" UUID NOT NULL,
        "organization_id" INTEGER NOT NULL,
        "start_timestamp" TIMESTAMPTZ NOT NULL,
        "group_id" BIGINT NOT NULL,
        "trace_id" UUID NOT NULL,
        "timestamp" TIMESTAMPTZ NULL,
        "duration" INTEGER NOT NULL,
        "data" JSONB NOT NULL,
        "tags" JSONB NOT NULL,
        PRIMARY KEY ("event_id", "organization_id", "start_timestamp")
    ) PARTITION BY HASH (organization_id);

    CREATE TABLE performance_transactionevent_p0 PARTITION OF performance_transactionevent FOR VALUES WITH (MODULUS 4, REMAINDER 0) PARTITION BY RANGE (start_timestamp);
    CREATE TABLE performance_transactionevent_p1 PARTITION OF performance_transactionevent FOR VALUES WITH (MODULUS 4, REMAINDER 1) PARTITION BY RANGE (start_timestamp);
    CREATE TABLE performance_transactionevent_p2 PARTITION OF performance_transactionevent FOR VALUES WITH (MODULUS 4, REMAINDER 2) PARTITION BY RANGE (start_timestamp);
    CREATE TABLE performance_transactionevent_p3 PARTITION OF performance_transactionevent FOR VALUES WITH (MODULUS 4, REMAINDER 3) PARTITION BY RANGE (start_timestamp);
"""
DROP_TRANSACTION_EVENT_SQL = 'DROP TABLE IF EXISTS "performance_transactionevent";'

# SQL for TransactionGroupAggregate
CREATE_ADVANCED_TRANSACTION_GROUP_AGGREGATE_SQL = """
    CREATE TABLE "performance_transactiongroupaggregate" (
        "group_id" BIGINT NOT NULL,
        "organization_id" INTEGER NOT NULL,
        "date" TIMESTAMPTZ NOT NULL,
        "count" INTEGER NOT NULL,
        "total_duration" DOUBLE PRECISION NOT NULL,
        "sum_of_squares_duration" DOUBLE PRECISION NOT NULL,
        "histogram" JSONB NOT NULL,
        PRIMARY KEY ("group_id", "organization_id", "date")
    ) PARTITION BY HASH (organization_id);

    CREATE TABLE performance_transactiongroupaggregate_p0 PARTITION OF performance_transactiongroupaggregate FOR VALUES WITH (MODULUS 4, REMAINDER 0) PARTITION BY RANGE (date);
    CREATE TABLE performance_transactiongroupaggregate_p1 PARTITION OF performance_transactiongroupaggregate FOR VALUES WITH (MODULUS 4, REMAINDER 1) PARTITION BY RANGE (date);
    CREATE TABLE performance_transactiongroupaggregate_p2 PARTITION OF performance_transactiongroupaggregate FOR VALUES WITH (MODULUS 4, REMAINDER 2) PARTITION BY RANGE (date);
    CREATE TABLE performance_transactiongroupaggregate_p3 PARTITION OF performance_transactiongroupaggregate FOR VALUES WITH (MODULUS 4, REMAINDER 3) PARTITION BY RANGE (date);
"""
DROP_TRANSACTION_GROUP_AGGREGATE_SQL = (
    'DROP TABLE IF EXISTS "performance_transactiongroupaggregate";'
)

always_run_operations = [
    migrations.AddField(
        model_name="transactiongroup",
        name="is_deleted",
        field=models.BooleanField(default=False),
    ),
    migrations.DeleteModel("TransactionEvent"),
]

# --- Define the operations for the simple path ---
base_operations = [
    psql_partition.backend.migrations.operations.create_partitioned_model.PostgresCreatePartitionedModel(
        name="TransactionEvent",
        fields=[
            (
                "pk",
                models.CompositePrimaryKey(
                    "event_id", "organization", "start_timestamp", primary_key=True
                ),
            ),
            ("event_id", models.UUIDField(default=uuid.uuid4, editable=False)),
            ("trace_id", models.UUIDField(db_index=True)),
            (
                "start_timestamp",
                models.DateTimeField(
                    db_index=True,
                    help_text="Datetime reported by client as the time the measurement started",
                ),
            ),
            (
                "timestamp",
                models.DateTimeField(
                    blank=True,
                    help_text="Datetime reported by client as the time the measurement finished",
                    null=True,
                ),
            ),
            (
                "duration",
                models.PositiveIntegerField(db_index=True, help_text="Milliseconds"),
            ),
            (
                "data",
                models.JSONField(help_text="General event data that is searchable"),
            ),
            ("tags", models.JSONField(default=dict)),
            (
                "group",
                models.ForeignKey(
                    on_delete=django.db.models.deletion.CASCADE,
                    to="performance.transactiongroup",
                ),
            ),
            (
                "organization",
                models.ForeignKey(
                    on_delete=django.db.models.deletion.CASCADE,
                    to="organizations_ext.organization",
                ),
            ),
        ],
        options={"ordering": ["-start_timestamp"]},
        partitioning_options={
            "method": psql_partition.types.PostgresPartitioningMethod["RANGE"],
            "key": ["start_timestamp"],
        },
        bases=(
            psql_partition.models.partitioned.PostgresPartitionedModel,
            models.Model,
        ),
    ),
    TestDefaultPartition(model_name="TransactionEvent", name="default"),
    psql_partition.backend.migrations.operations.create_partitioned_model.PostgresCreatePartitionedModel(
        name="TransactionGroupAggregate",
        fields=[
            (
                "pk",
                models.CompositePrimaryKey(
                    "group", "organization", "date", primary_key=True
                ),
            ),
            ("date", models.DateTimeField()),
            ("count", models.PositiveIntegerField()),
            (
                "total_duration",
                models.FloatField(
                    default=0.0,
                    help_text="Sum of all transaction durations (in ms) for calculating the mean.",
                ),
            ),
            (
                "sum_of_squares_duration",
                models.FloatField(
                    default=0.0,
                    help_text="Sum of squares of durations, for calculating standard deviation.",
                ),
            ),
            (
                "histogram",
                models.JSONField(
                    default=dict,
                    help_text="Stores a fixed-bucket histogram for percentile approximation.",
                ),
            ),
            (
                "group",
                models.ForeignKey(
                    on_delete=django.db.models.deletion.CASCADE,
                    to="performance.transactiongroup",
                ),
            ),
            (
                "organization",
                models.ForeignKey(
                    on_delete=django.db.models.deletion.CASCADE,
                    to="organizations_ext.organization",
                ),
            ),
        ],
        options={"abstract": False},
        partitioning_options={
            "method": psql_partition.types.PostgresPartitioningMethod["RANGE"],
            "key": ["date"],
        },
        bases=(
            psql_partition.models.partitioned.PostgresPartitionedModel,
            models.Model,
        ),
    ),
    TestDefaultPartition(model_name="TransactionGroupAggregate", name="default"),
    migrations.AlterModelManagers(
        name="transactionevent",
        managers=[
            ("objects", psql_partition.manager.manager.PostgresManager()),
        ],
    ),
    migrations.AlterModelManagers(
        name="transactiongroupaggregate",
        managers=[
            ("objects", psql_partition.manager.manager.PostgresManager()),
        ],
    ),
    # Django bug? makemigrations will try to alter it like this, which is nonsensible.
    migrations.AlterField(
        model_name="transactionevent",
        name="pk",
        field=models.CompositePrimaryKey(
            "event_id",
            "organization",
            "start_timestamp",
            blank=True,
            editable=False,
            primary_key=True,
            serialize=False,
        ),
    ),
    migrations.AlterField(
        model_name="transactiongroupaggregate",
        name="pk",
        field=models.CompositePrimaryKey(
            "group",
            "organization",
            "date",
            blank=True,
            editable=False,
            primary_key=True,
            serialize=False,
        ),
    ),
]

# --- Conditionally build the final operations list ---
if settings.GLITCHTIP_ADVANCED_PARTITIONING:
    # For the advanced path, we use SeparateDatabaseAndState for each model
    operations = always_run_operations + [
        # The AddField operation is common to both paths
        # Handle TransactionEvent
        migrations.SeparateDatabaseAndState(
            state_operations=base_operations,
            database_operations=[
                migrations.RunSQL(
                    sql=CREATE_ADVANCED_TRANSACTION_EVENT_SQL,
                    reverse_sql=DROP_TRANSACTION_EVENT_SQL,
                ),
                migrations.RunSQL(
                    sql=CREATE_ADVANCED_TRANSACTION_GROUP_AGGREGATE_SQL,
                    reverse_sql=DROP_TRANSACTION_GROUP_AGGREGATE_SQL,
                ),
            ],
        ),
    ]
else:
    # For the simple path, use the auto-generated operations directly
    operations = always_run_operations + base_operations


class Migration(migrations.Migration):
    dependencies = [
        ("organizations_ext", "0010_alter_organization_id"),
        ("performance", "0001_squashed_0014_initial"),
    ]

    operations = operations