import django.db.models.deletion
import uuid
from django.conf import settings
from django.db import migrations, models
from pgvector.django import VectorField


def _create_vector_extension(apps, schema_editor):
    if schema_editor.connection.vendor != 'postgresql':
        return
    with schema_editor.connection.cursor() as cursor:
        cursor.execute('CREATE EXTENSION IF NOT EXISTS vector')


def _drop_vector_extension(apps, schema_editor):
    if schema_editor.connection.vendor != 'postgresql':
        return
    with schema_editor.connection.cursor() as cursor:
        cursor.execute('DROP EXTENSION IF EXISTS vector')


class Migration(migrations.Migration):

    initial = True

    dependencies = [
        ('orgs', '0001_initial'),
        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
    ]

    operations = [
        migrations.RunPython(
            code=_create_vector_extension,
            reverse_code=_drop_vector_extension,
        ),
        migrations.CreateModel(
            name='AgentModel',
            fields=[
                ('id', models.BigAutoField(primary_key=True, serialize=False)),
                ('uuid', models.UUIDField(default=uuid.uuid4, editable=False, unique=True)),
                ('name', models.CharField(max_length=255)),
                ('version', models.CharField(max_length=50)),
                ('path', models.CharField(blank=True, default='', max_length=1024)),
            ],
            options={
                'verbose_name': 'Model',
                'verbose_name_plural': 'Models',
            },
        ),
        migrations.CreateModel(
            name='Agent',
            fields=[
                ('created_at', models.DateTimeField(auto_now_add=True, verbose_name='Created At')),
                ('updated_at', models.DateTimeField(auto_now=True, verbose_name='Updated At')),
                ('id', models.BigAutoField(primary_key=True, serialize=False)),
                ('uuid', models.UUIDField(default=uuid.uuid4, editable=False, unique=True)),
                ('status', models.CharField(choices=[('idle', 'Idle'), ('running', 'Running'), ('paused', 'Paused'), ('completed', 'Completed'), ('failed', 'Failed')], default='idle', max_length=20)),
                ('description', models.TextField(blank=True, default='')),
                ('started_at', models.DateTimeField(blank=True, null=True)),
                ('completed_at', models.DateTimeField(blank=True, null=True)),
                ('organization', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='agents', to='orgs.organization')),
                ('model', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='agents', to='mlstore.agentmodel')),
            ],
            options={
                'verbose_name': 'Agent Instance',
                'verbose_name_plural': 'Agent Instances',
            },
        ),
        migrations.CreateModel(
            name='AgentRun',
            fields=[
                ('created_at', models.DateTimeField(auto_now_add=True, verbose_name='Created At')),
                ('updated_at', models.DateTimeField(auto_now=True, verbose_name='Updated At')),
                ('id', models.BigAutoField(primary_key=True, serialize=False)),
                ('uuid', models.UUIDField(default=uuid.uuid4, editable=False, unique=True)),
                ('status', models.CharField(choices=[('queued', 'Queued'), ('running', 'Running'), ('completed', 'Completed'), ('failed', 'Failed')], default='queued', max_length=20)),
                ('input_data', models.JSONField(blank=True, default=dict)),
                ('output_data', models.JSONField(blank=True, default=dict)),
                ('error_message', models.TextField(blank=True, default='')),
                ('started_at', models.DateTimeField(blank=True, null=True)),
                ('completed_at', models.DateTimeField(blank=True, null=True)),
                ('agent', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='runs', to='mlstore.agent')),
                ('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='agent_runs', to=settings.AUTH_USER_MODEL)),
            ],
            options={
                'verbose_name': 'Agent Run',
                'verbose_name_plural': 'Agent Runs',
            },
        ),
        migrations.CreateModel(
            name='AgentEvent',
            fields=[
                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
                ('uuid', models.UUIDField(default=uuid.uuid4, editable=False, unique=True)),
                ('event_type', models.CharField(choices=[('started', 'Started'), ('message', 'Message'), ('progress', 'Progress'), ('completed', 'Completed'), ('error', 'Error'), ('step', 'Step')], max_length=20)),
                ('content', models.JSONField()),
                ('timestamp', models.DateTimeField(auto_now_add=True)),
                ('execution', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='events', to='mlstore.agentrun')),
            ],
            options={
                'verbose_name': 'Agent Event',
                'verbose_name_plural': 'Agent Events',
                'ordering': ['timestamp'],
            },
        ),
        migrations.CreateModel(
            name='RoleRagDocument',
            fields=[
                ('created_at', models.DateTimeField(auto_now_add=True, verbose_name='Created At')),
                ('updated_at', models.DateTimeField(auto_now=True, verbose_name='Updated At')),
                ('id', models.BigAutoField(primary_key=True, serialize=False)),
                ('uuid', models.UUIDField(default=uuid.uuid4, editable=False, unique=True)),
                ('content', models.TextField()),
                ('content_hash', models.CharField(db_index=True, max_length=64)),
                ('embedding', VectorField(blank=True, dimensions=1536, null=True)),
                ('metadata', models.JSONField(blank=True, default=dict)),
                ('chunk_index', models.IntegerField(default=0)),
                ('is_active', models.BooleanField(default=True)),
                ('role', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='rag_documents', to='orgs.role')),
                ('training_file', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='rag_documents', to='orgs.trainingfile')),
            ],
            options={
                'verbose_name': 'Role RAG Document',
                'verbose_name_plural': 'Role RAG Documents',
            },
        ),
    ]