"""Peewee migrations -- 009_add_models.py.

Some examples (model - class or model name)::

    > Model = migrator.orm['table_name']            # Return model in current state by name
    > Model = migrator.ModelClass                   # Return model in current state by name

    > migrator.sql(sql)                             # Run custom SQL
    > migrator.run(func, *args, **kwargs)           # Run python function with the given args
    > migrator.create_model(Model)                  # Create a model (could be used as decorator)
    > migrator.remove_model(model, cascade=True)    # Remove a model
    > migrator.add_fields(model, **fields)          # Add fields to a model
    > migrator.change_fields(model, **fields)       # Change fields
    > migrator.remove_fields(model, *field_names, cascade=True)
    > migrator.rename_field(model, old_field_name, new_field_name)
    > migrator.rename_table(model, new_table_name)
    > migrator.add_index(model, *col_names, unique=False)
    > migrator.add_not_null(model, *field_names)
    > migrator.add_default(model, field_name, default)
    > migrator.add_constraint(model, name, sql)
    > migrator.drop_index(model, *col_names)
    > migrator.drop_not_null(model, *field_names)
    > migrator.drop_constraints(model, *constraints)

"""

from contextlib import suppress

import peewee as pw
from peewee_migrate import Migrator
import json

from open_webui.utils.misc import parse_ollama_modelfile

with suppress(ImportError):
    import playhouse.postgres_ext as pw_pext


def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
    """Write your migrations here."""

    # Fetch data from 'modelfile' table and insert into 'model' table
    migrate_modelfile_to_model(migrator, database)
    # Drop the 'modelfile' table
    migrator.remove_model("modelfile")


def migrate_modelfile_to_model(migrator: Migrator, database: pw.Database):
    ModelFile = migrator.orm["modelfile"]
    Model = migrator.orm["model"]

    modelfiles = ModelFile.select()

    for modelfile in modelfiles:
        # Extract and transform data in Python

        modelfile.modelfile = json.loads(modelfile.modelfile)
        meta = json.dumps(
            {
                "description": modelfile.modelfile.get("desc"),
                "profile_image_url": modelfile.modelfile.get("imageUrl"),
                "ollama": {"modelfile": modelfile.modelfile.get("content")},
                "suggestion_prompts": modelfile.modelfile.get("suggestionPrompts"),
                "categories": modelfile.modelfile.get("categories"),
                "user": {**modelfile.modelfile.get("user", {}), "community": True},
            }
        )

        info = parse_ollama_modelfile(modelfile.modelfile.get("content"))

        # Insert the processed data into the 'model' table
        Model.create(
            id=f"ollama-{modelfile.tag_name}",
            user_id=modelfile.user_id,
            base_model_id=info.get("base_model_id"),
            name=modelfile.modelfile.get("title"),
            meta=meta,
            params=json.dumps(info.get("params", {})),
            created_at=modelfile.timestamp,
            updated_at=modelfile.timestamp,
        )


def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
    """Write your rollback migrations here."""

    recreate_modelfile_table(migrator, database)
    move_data_back_to_modelfile(migrator, database)
    migrator.remove_model("model")


def recreate_modelfile_table(migrator: Migrator, database: pw.Database):
    query = """
    CREATE TABLE IF NOT EXISTS modelfile (
        user_id TEXT,
        tag_name TEXT,
        modelfile JSON,
        timestamp BIGINT
    )
    """
    migrator.sql(query)


def move_data_back_to_modelfile(migrator: Migrator, database: pw.Database):
    Model = migrator.orm["model"]
    Modelfile = migrator.orm["modelfile"]

    models = Model.select()

    for model in models:
        # Extract and transform data in Python
        meta = json.loads(model.meta)

        modelfile_data = {
            "title": model.name,
            "desc": meta.get("description"),
            "imageUrl": meta.get("profile_image_url"),
            "content": meta.get("ollama", {}).get("modelfile"),
            "suggestionPrompts": meta.get("suggestion_prompts"),
            "categories": meta.get("categories"),
            "user": {k: v for k, v in meta.get("user", {}).items() if k != "community"},
        }

        # Insert the processed data back into the 'modelfile' table
        Modelfile.create(
            user_id=model.user_id,
            tag_name=model.id,
            modelfile=modelfile_data,
            timestamp=model.created_at,
        )