# Generated by Django 4.2.10 on 2024-04-13 17:54

import base64
import io

from django.db import migrations
from PIL import Image

from khoj.utils.helpers import ImageIntentType


def convert_png_images_to_webp(apps, schema_editor):
    # Get the model from the versioned app registry to ensure the correct version is used
    Conversations = apps.get_model("database", "Conversation")
    for conversation in Conversations.objects.all():
        for chat in conversation.conversation_log["chat"]:
            if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE.value:
                # Decode the base64 encoded PNG image
                decoded_image = base64.b64decode(chat["message"])

                # Convert images from PNG to WebP format
                image_io = io.BytesIO(decoded_image)
                with Image.open(image_io) as png_image:
                    webp_image_io = io.BytesIO()
                    png_image.save(webp_image_io, "WEBP")

                    # Encode the WebP image back to base64
                    webp_image_bytes = webp_image_io.getvalue()
                    chat["message"] = base64.b64encode(webp_image_bytes).decode()
                    chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE_V3.value
                    webp_image_io.close()

        # Save the updated conversation history
        conversation.save()


def convert_webp_images_to_png(apps, schema_editor):
    # Get the model from the versioned app registry to ensure the correct version is used
    Conversations = apps.get_model("database", "Conversation")
    for conversation in Conversations.objects.all():
        for chat in conversation.conversation_log["chat"]:
            if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE_V3.value:
                # Decode the base64 encoded PNG image
                decoded_image = base64.b64decode(chat["message"])

                # Convert images from PNG to WebP format
                image_io = io.BytesIO(decoded_image)
                with Image.open(image_io) as png_image:
                    webp_image_io = io.BytesIO()
                    png_image.save(webp_image_io, "PNG")

                    # Encode the WebP image back to base64
                    webp_image_bytes = webp_image_io.getvalue()
                    chat["message"] = base64.b64encode(webp_image_bytes).decode()
                    chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE.value
                    webp_image_io.close()

        # Save the updated conversation history
        conversation.save()


class Migration(migrations.Migration):
    dependencies = [
        ("database", "0034_alter_chatmodeloptions_chat_model"),
    ]

    operations = [
        migrations.RunPython(convert_png_images_to_webp, reverse_code=convert_webp_images_to_png),
    ]
