# Generated by Django 5.2 on 2025-01-05 11:58

from pathlib import Path

import lamindb_setup as ln_setup
import psycopg2
from django.db import migrations


def get_artifact_path_psycopg2(artifact_id):
    """Get artifact path using psycopg2."""
    query = """
        SELECT
            s.root || '/.lamindb/' || a.uid || a.suffix AS full_path
        FROM
            lamindb_artifact a
            JOIN lamindb_storage s ON a.storage_id = s.id
        WHERE
            a.id = %s
    """

    with psycopg2.connect(ln_setup.settings.instance.db) as conn:
        with conn.cursor() as cur:
            cur.execute(query, (artifact_id,))
            return cur.fetchone()[0]


def transfer_source_code(apps, schema_editor):
    from lamindb._finish import notebook_to_script

    Transform = apps.get_model("lamindb", "Transform")
    transforms = Transform.objects.filter(
        _source_code_artifact__isnull=False,
    ).select_related("_source_code_artifact")

    for transform in transforms:
        print(f"migrating source code of transform {transform}")
        artifact = transform._source_code_artifact
        print("artifact", artifact.uid)

        path_str = get_artifact_path_psycopg2(artifact.id)
        print(ln_setup.settings.storage.root_as_str)
        print(path_str)
        if path_str.startswith(ln_setup.settings.storage.root_as_str):
            path = (
                ln_setup.settings.storage.root
                / f".lamindb/{artifact.uid}{artifact.suffix}"
            )
        else:
            path = ln_setup.core.upath.UPath(path_str)
        if path.exists():
            if path_str.startswith("s3://"):
                local_path = Path(f"temp{path.suffix}")
                path.download_to(local_path)
            else:
                local_path = path

            if artifact.suffix == ".ipynb":
                transform.source_code = notebook_to_script(transform, local_path)
            else:
                transform.source_code = local_path.read_text()
            transform.hash = artifact.hash
            path.unlink()
        else:
            print(f"path did not exist: {path_str}")
        transform._source_code_artifact = None
        transform.save()
        artifact.delete()


class Migration(migrations.Migration):
    dependencies = [
        ("lamindb", "0069_squashed"),
    ]

    operations = [
        migrations.RunPython(transfer_source_code),
    ]
