"""migrate existing file artifacts to database

Revision ID: b2c3d4e5f6g7
Revises: a1b2c3d4e5f6
Create Date: 2026-01-21

This migration automatically moves existing experiment artifacts from
file-based storage to database storage. It also populates metadata tables
(feature_selections, feature_selection_ranks, model_selections,
model_selection_scores) from folder data if they don't exist.
"""

import io
import json
import os
import logging
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.orm import Session

# revision identifiers, used by Alembic.
revision: str = "b2c3d4e5f6g7"
down_revision: Union[str, None] = "a1b2c3d4e5f6"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

logger = logging.getLogger("alembic.runtime.migration")


# =============================================================================
# Metadata table population functions
# =============================================================================

def _get_or_create_model(session, Model, model_name: str):
    """Get or create a Model record."""
    # Determine model_type from name
    model_type_map = {
        "xgb": "xgboost",
        "lgb": "lightgbm",
        "catboost": "catboost",
        "rf": "random_forest",
        "svm": "svm",
        "lr": "logistic_regression",
        "mlp": "mlp",
        "keras": "keras",
    }
    model_type = model_type_map.get(model_name.split(".")[0], model_name.split(".")[0])

    existing = session.query(Model).filter(Model.name == model_name).first()
    if existing:
        return existing

    model = Model(name=model_name, model_type=model_type)
    session.add(model)
    session.flush()
    return model


def _get_or_create_feature(session, Feature, feature_name: str):
    """Get or create a Feature record."""
    existing = session.query(Feature).filter(Feature.name == feature_name).first()
    if existing:
        return existing

    feature = Feature(name=feature_name)
    session.add(feature)
    session.flush()
    return feature


def _populate_model_selection_scores(session, experiment, target, model_selection, target_dir, Model, ModelSelectionScore, joblib, pd):
    """Populate model_selection_scores from scores_tracking.csv."""
    scores_path = f"{target_dir}/scores_tracking.csv"
    if not os.path.exists(scores_path):
        return

    try:
        df = pd.read_csv(scores_path)

        for _, row in df.iterrows():
            model_name = row.get("Model") or row.get("model")
            if not model_name:
                continue

            # Check if score already exists
            model = _get_or_create_model(session, Model, model_name)
            existing = session.query(ModelSelectionScore).filter(
                ModelSelectionScore.model_id == model.id,
                ModelSelectionScore.model_selection_id == model_selection.id,
            ).first()

            if existing:
                continue

            # Load best_params if available
            best_params = None
            params_path = f"{target_dir}/{model_name}/best_params.json"
            if os.path.exists(params_path):
                try:
                    with open(params_path, "r") as f:
                        best_params = json.load(f)
                except:
                    pass

            # Create score record
            score = ModelSelectionScore(
                model_id=model.id,
                model_selection_id=model_selection.id,
                best_params=best_params,
                rmse=row.get("RMSE") or row.get("rmse"),
                mae=row.get("MAE") or row.get("mae"),
                r2=row.get("R2") or row.get("r2"),
                logloss=row.get("LOGLOSS") or row.get("logloss"),
                accuracy=row.get("ACCURACY") or row.get("accuracy"),
                precision=row.get("PRECISION") or row.get("precision"),
                recall=row.get("RECALL") or row.get("recall"),
                f1=row.get("F1") or row.get("f1"),
                roc_auc=row.get("ROC_AUC") or row.get("roc_auc"),
                avg_precision=row.get("AVG_PRECISION") or row.get("avg_precision"),
                training_time=row.get("TRAINING_TIME") or row.get("training_time"),
            )
            session.add(score)

        session.flush()
        logger.info(f"  Populated model_selection_scores for {target.name}")
    except Exception as e:
        logger.warning(f"  Failed to populate model_selection_scores for {target.name}: {e}")


def _populate_model_selection(session, experiment, target, target_dir, Model, ModelSelection, ModelSelectionScore, joblib, pd):
    """Populate model_selection from folder data."""
    # Check if model_selection already exists
    existing = session.query(ModelSelection).filter(
        ModelSelection.experiment_id == experiment.id,
        ModelSelection.target_id == target.id,
    ).first()

    if existing:
        # Still populate scores if missing
        if not existing.model_selection_scores:
            _populate_model_selection_scores(session, experiment, target, existing, target_dir, Model, ModelSelectionScore, joblib, pd)
        return existing

    # Load best_params
    best_model_params = None
    best_params_path = f"{target_dir}/best_params.json"
    if os.path.exists(best_params_path):
        try:
            with open(best_params_path, "r") as f:
                best_model_params = json.load(f)
        except:
            pass

    # Load thresholds
    best_thresholds = None
    thresholds_path = f"{target_dir}/thresholds.pkl"
    if os.path.exists(thresholds_path):
        try:
            best_thresholds = joblib.load(thresholds_path)
        except:
            pass

    # Determine best model from scores_tracking.csv
    best_model_id = None
    scores_path = f"{target_dir}/scores_tracking.csv"
    if os.path.exists(scores_path):
        try:
            df = pd.read_csv(scores_path)
            # Find best model (lowest logloss or rmse)
            if "LOGLOSS" in df.columns or "logloss" in df.columns:
                col = "LOGLOSS" if "LOGLOSS" in df.columns else "logloss"
                best_row = df.loc[df[col].idxmin()]
            elif "RMSE" in df.columns or "rmse" in df.columns:
                col = "RMSE" if "RMSE" in df.columns else "rmse"
                best_row = df.loc[df[col].idxmin()]
            else:
                best_row = df.iloc[0] if len(df) > 0 else None

            if best_row is not None:
                model_name = best_row.get("Model") or best_row.get("model")
                if model_name:
                    model = _get_or_create_model(session, Model, model_name)
                    best_model_id = model.id
        except:
            pass

    # Create model_selection
    model_selection = ModelSelection(
        experiment_id=experiment.id,
        target_id=target.id,
        best_model_params=best_model_params,
        best_thresholds=best_thresholds,
        best_model_id=best_model_id,
    )
    session.add(model_selection)
    session.flush()

    # Populate scores
    _populate_model_selection_scores(session, experiment, target, model_selection, target_dir, Model, ModelSelectionScore, joblib, pd)

    logger.info(f"  Created model_selection for {target.name}")
    return model_selection


def _populate_feature_selection_ranks(session, feature_selection, target_dir, Feature, FeatureSelectionRank, pd):
    """Populate feature_selection_ranks from feature_scores_*.csv files."""
    feature_selection_dir = f"{target_dir}/feature_selection"
    if not os.path.exists(feature_selection_dir):
        return

    feature_score_methods = [
        ("Chi2", "Chi2", "chi2"),
        ("Person's_R", "Person's R", "pearson"),
        ("ANOVA", "ANOVA", "anova"),
        ("Spearman's_R", "Spearman's R", "spearman"),
        ("Kendall's_Tau", "Kendall's Tau", "kendall"),
        ("MI", "MI", "mutual_info"),
        ("FI", "FI", "feature_importance"),
        ("RFE", "RFE", "rfe"),
        ("SFS", "SFS", "sfs"),
    ]

    for old_file_method, new_file_method, db_method in feature_score_methods:
        # Try new naming convention first, then old
        scores_path = f"{feature_selection_dir}/{new_file_method}.csv"
        if not os.path.exists(scores_path):
            scores_path = f"{feature_selection_dir}/feature_scores_{old_file_method}.csv"
        if not os.path.exists(scores_path):
            continue

        try:
            df = pd.read_csv(scores_path)

            for _, row in df.iterrows():
                feature_name = row.get("Feature") or row.get("feature")
                if not feature_name:
                    continue

                feature = _get_or_create_feature(session, Feature, feature_name)

                # Check if rank already exists
                existing = session.query(FeatureSelectionRank).filter(
                    FeatureSelectionRank.feature_id == feature.id,
                    FeatureSelectionRank.feature_selection_id == feature_selection.id,
                    FeatureSelectionRank.method == db_method,
                ).first()

                if existing:
                    continue

                rank = FeatureSelectionRank(
                    feature_id=feature.id,
                    feature_selection_id=feature_selection.id,
                    method=db_method,
                    score=row.get("Score") or row.get("score"),
                    pvalue=row.get("PValue") or row.get("pvalue"),
                    rank=row.get("Rank") or row.get("rank"),
                    support=row.get("Support") or row.get("support"),
                )
                session.add(rank)

            session.flush()
            logger.info(f"  Populated feature_selection_ranks ({db_method})")
        except Exception as e:
            logger.warning(f"  Failed to populate ranks for {file_method}: {e}")


def _populate_feature_selection(session, experiment, target, target_dir, Feature, FeatureSelection, FeatureSelectionRank, feature_selection_association, joblib, pd):
    """Populate feature_selection and feature_selection_association from features.pkl."""
    # Check if feature_selection already exists
    existing = session.query(FeatureSelection).filter(
        FeatureSelection.experiment_id == experiment.id,
        FeatureSelection.target_id == target.id,
    ).first()

    if existing:
        # Still populate ranks if missing
        if not existing.feature_selection_ranks:
            _populate_feature_selection_ranks(session, existing, target_dir, Feature, FeatureSelectionRank, pd)
        return existing

    features_path = f"{target_dir}/features.pkl"
    if not os.path.exists(features_path):
        return None

    try:
        feature_names = joblib.load(features_path)
        if not isinstance(feature_names, list):
            return None

        # Create feature_selection
        feature_selection = FeatureSelection(
            experiment_id=experiment.id,
            target_id=target.id,
            best_features_path=features_path,
        )
        session.add(feature_selection)
        session.flush()

        # Create feature records and associations
        for feature_name in feature_names:
            feature = _get_or_create_feature(session, Feature, feature_name)

            # Insert into association table
            session.execute(
                feature_selection_association.insert().values(
                    feature_selection_id=feature_selection.id,
                    feature_id=feature.id,
                ).prefix_with("IGNORE")
            )

        session.flush()
        logger.info(f"  Created feature_selection for {target.name} with {len(feature_names)} features")

        # Populate ranks
        _populate_feature_selection_ranks(session, feature_selection, target_dir, Feature, FeatureSelectionRank, pd)

        return feature_selection
    except Exception as e:
        logger.warning(f"  Failed to populate feature_selection for {target.name}: {e}")
        return None


def _populate_metadata_tables(session, experiment, experiment_path, Feature, FeatureSelection, FeatureSelectionRank, Model, ModelSelection, ModelSelectionScore, feature_selection_association, joblib, pd):
    """Populate all metadata tables from folder data."""
    for target in experiment.targets:
        target_dir = f"{experiment_path}/{target.name}"
        if not os.path.exists(target_dir):
            continue

        # Populate feature_selection, feature_selection_association, and ranks
        _populate_feature_selection(
            session, experiment, target, target_dir,
            Feature, FeatureSelection, FeatureSelectionRank,
            feature_selection_association, joblib, pd
        )

        # Populate model_selection and scores
        _populate_model_selection(
            session, experiment, target, target_dir,
            Model, ModelSelection, ModelSelectionScore, joblib, pd
        )


# =============================================================================
# Artifact serialization functions
# =============================================================================

def _serialize_joblib(obj) -> bytes:
    """Serialize an object using joblib."""
    import joblib
    buffer = io.BytesIO()
    joblib.dump(obj, buffer)
    return buffer.getvalue()


def _serialize_dataframe(df) -> bytes:
    """Serialize a DataFrame to parquet bytes."""
    # Convert object columns with non-serializable types to strings
    df_copy = df.copy()
    for col in df_copy.columns:
        if df_copy[col].dtype == "object":
            df_copy[col] = df_copy[col].astype(str)
    buffer = io.BytesIO()
    df_copy.to_parquet(buffer, compression="snappy", index=True)
    return buffer.getvalue()


def _serialize_keras(model) -> bytes:
    """Serialize a Keras model to bytes."""
    buffer = io.BytesIO()
    model.save(buffer)
    return buffer.getvalue()


def _save_artifact(session, ExperimentArtifact, experiment_id, artifact_type, artifact_name, data, serialization_format, target_id=None):
    """Save an artifact directly to the database."""
    # Check if exists
    existing = session.query(ExperimentArtifact).filter(
        ExperimentArtifact.experiment_id == experiment_id,
        ExperimentArtifact.artifact_type == artifact_type,
        ExperimentArtifact.artifact_name == artifact_name,
        ExperimentArtifact.target_id == target_id,
    ).first()

    if existing:
        existing.data = data
        existing.serialization_format = serialization_format
    else:
        artifact = ExperimentArtifact(
            experiment_id=experiment_id,
            target_id=target_id,
            artifact_type=artifact_type,
            artifact_name=artifact_name,
            data=data,
            serialization_format=serialization_format,
        )
        session.add(artifact)


def _save_dataframe_record(session, ExperimentData, experiment_id, data_type, data, row_count, column_count, target_id=None, model_selection_score_id=None):
    """Save a DataFrame record directly to the database."""
    # Check if exists
    if data_type == "prediction" and model_selection_score_id:
        existing = session.query(ExperimentData).filter(
            ExperimentData.model_selection_score_id == model_selection_score_id,
            ExperimentData.data_type == "prediction",
        ).first()
    else:
        existing = session.query(ExperimentData).filter(
            ExperimentData.experiment_id == experiment_id,
            ExperimentData.data_type == data_type,
            ExperimentData.target_id == target_id,
        ).first()

    if existing:
        existing.data = data
        existing.row_count = row_count
        existing.column_count = column_count
    else:
        record = ExperimentData(
            experiment_id=experiment_id,
            target_id=target_id,
            model_selection_score_id=model_selection_score_id,
            data_type=data_type,
            data=data,
            row_count=row_count,
            column_count=column_count,
        )
        session.add(record)


def upgrade() -> None:
    """Migrate existing file-based artifacts to database and populate metadata tables."""
    bind = op.get_bind()
    session = Session(bind=bind)

    try:
        import joblib
        import pandas as pd
        from lecrapaud.models import (
            Experiment, ExperimentArtifact, ExperimentData,
            Feature, FeatureSelection, FeatureSelectionRank,
            Model, ModelSelection, ModelSelectionScore,
        )
        from lecrapaud.models.feature_selection import lecrapaud_feature_selection_association

        # Find all experiments with a path set
        experiments = (
            session.query(Experiment)
            .filter(Experiment.path.isnot(None))
            .all()
        )

        if not experiments:
            logger.info("No experiments with file paths found, nothing to migrate")
            return

        logger.info(f"Found {len(experiments)} experiment(s) to migrate")

        for experiment in experiments:
            experiment_path = str(experiment.path) if experiment.path else None

            if not experiment_path or not os.path.exists(experiment_path):
                logger.warning(
                    f"Experiment {experiment.id} ({experiment.name}): path does not exist, skipping"
                )
                continue

            logger.info(f"Migrating experiment {experiment.id}: {experiment.name}")

            try:
                # 1. Populate metadata tables (feature_selections, model_selections, etc.)
                _populate_metadata_tables(
                    session, experiment, experiment_path,
                    Feature, FeatureSelection, FeatureSelectionRank,
                    Model, ModelSelection, ModelSelectionScore,
                    lecrapaud_feature_selection_association, joblib, pd
                )

                # 2. Migrate binary artifacts and dataframes
                _migrate_experiment(session, experiment, experiment_path, ExperimentArtifact, ExperimentData, joblib, pd)

                # Commit after each experiment to avoid large transactions
                session.commit()
                logger.info(f"  Committed experiment {experiment.id}")
            except Exception as e:
                logger.error(f"Failed to migrate experiment {experiment.id}: {e}")
                session.rollback()
                # Continue with other experiments

        logger.info("Migration complete")

    except ImportError as e:
        logger.warning(f"Could not import required modules, skipping migration: {e}")
    except Exception as e:
        logger.error(f"Migration failed: {e}")
        session.rollback()
        raise
    finally:
        session.close()


def _migrate_experiment(session, experiment, experiment_path, ExperimentArtifact, ExperimentData, joblib, pd):
    """Migrate all artifacts for a single experiment."""
    data_dir = f"{experiment_path}/data"
    preprocessing_dir = f"{experiment_path}/preprocessing"

    # 1. Migrate DataFrames
    dataframe_files = [
        ("full", f"{data_dir}/full.pkl"),
        ("train", f"{data_dir}/train.pkl"),
        ("val", f"{data_dir}/val.pkl"),
        ("test", f"{data_dir}/test.pkl"),
        ("train_scaled", f"{preprocessing_dir}/train_scaled.pkl"),
        ("val_scaled", f"{preprocessing_dir}/val_scaled.pkl"),
        ("test_scaled", f"{preprocessing_dir}/test_scaled.pkl"),
    ]

    for data_type, file_path in dataframe_files:
        if os.path.exists(file_path):
            try:
                df = joblib.load(file_path)
                if isinstance(df, pd.DataFrame):
                    data = _serialize_dataframe(df)
                    _save_dataframe_record(
                        session, ExperimentData,
                        experiment_id=experiment.id,
                        data_type=data_type,
                        data=data,
                        row_count=len(df),
                        column_count=len(df.columns),
                    )
                    logger.info(f"  Migrated DataFrame: {data_type}")
            except Exception as e:
                logger.warning(f"  Failed to migrate {data_type}: {e}")

    # 2. Migrate global scaler_x
    scaler_x_path = f"{preprocessing_dir}/scaler_x.pkl"
    if os.path.exists(scaler_x_path):
        try:
            scaler = joblib.load(scaler_x_path)
            data = _serialize_joblib(scaler)
            _save_artifact(
                session, ExperimentArtifact,
                experiment_id=experiment.id,
                artifact_type="scaler",
                artifact_name="scaler_x",
                data=data,
                serialization_format="joblib",
            )
            logger.info(f"  Migrated scaler: scaler_x")
        except Exception as e:
            logger.warning(f"  Failed to migrate scaler_x: {e}")

    # 3. Migrate PCAs and transformers
    artifacts = [
        ("pca", "pcas", f"{preprocessing_dir}/pcas.pkl"),
        ("pca", "pcas_cross_sectional", f"{preprocessing_dir}/pcas_cross_sectional.pkl"),
        ("pca", "pcas_temporal", f"{preprocessing_dir}/pcas_temporal.pkl"),
        ("transformer", "column_transformer", f"{preprocessing_dir}/column_transformer.pkl"),
    ]

    for artifact_type, artifact_name, file_path in artifacts:
        if os.path.exists(file_path):
            try:
                obj = joblib.load(file_path)
                data = _serialize_joblib(obj)
                _save_artifact(
                    session, ExperimentArtifact,
                    experiment_id=experiment.id,
                    artifact_type=artifact_type,
                    artifact_name=artifact_name,
                    data=data,
                    serialization_format="joblib",
                )
                logger.info(f"  Migrated {artifact_type}: {artifact_name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate {artifact_name}: {e}")

    # 3b. Migrate all_features lists
    feature_list_files = [
        ("features", "all_features_before_encoding", f"{preprocessing_dir}/all_features_before_encoding.pkl"),
        ("features", "all_features_before_selection", f"{preprocessing_dir}/all_features_before_selection.pkl"),
        ("features", "all_features", f"{preprocessing_dir}/all_features.pkl"),
    ]

    for artifact_type, artifact_name, file_path in feature_list_files:
        if os.path.exists(file_path):
            try:
                features = joblib.load(file_path)
                if isinstance(features, list):
                    data = json.dumps(features).encode("utf-8")
                    _save_artifact(
                        session, ExperimentArtifact,
                        experiment_id=experiment.id,
                        artifact_type=artifact_type,
                        artifact_name=artifact_name,
                        data=data,
                        serialization_format="json",
                    )
                    logger.info(f"  Migrated {artifact_type}: {artifact_name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate {artifact_name}: {e}")

    # 4. Migrate feature_summary
    summary_path = f"{experiment_path}/feature_summary.csv"
    if os.path.exists(summary_path):
        try:
            df = pd.read_csv(summary_path)
            data = _serialize_dataframe(df)
            _save_dataframe_record(
                session, ExperimentData,
                experiment_id=experiment.id,
                data_type="feature_summary",
                data=data,
                row_count=len(df),
                column_count=len(df.columns),
            )
            logger.info(f"  Migrated feature_summary")
        except Exception as e:
            logger.warning(f"  Failed to migrate feature_summary: {e}")

    # 5. Migrate per-target artifacts
    for target in experiment.targets:
        target_dir = f"{experiment_path}/{target.name}"
        if not os.path.exists(target_dir):
            continue

        # scaler_y
        scaler_y_path = f"{target_dir}/scaler_y.pkl"
        if os.path.exists(scaler_y_path):
            try:
                scaler = joblib.load(scaler_y_path)
                data = _serialize_joblib(scaler)
                _save_artifact(
                    session, ExperimentArtifact,
                    experiment_id=experiment.id,
                    artifact_type="scaler",
                    artifact_name="scaler_y",
                    data=data,
                    serialization_format="joblib",
                    target_id=target.id,
                )
                logger.info(f"  Migrated scaler_y for {target.name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate scaler_y for {target.name}: {e}")

        # features
        features_path = f"{target_dir}/features.pkl"
        if os.path.exists(features_path):
            try:
                features = joblib.load(features_path)
                data = json.dumps(features).encode("utf-8")
                _save_artifact(
                    session, ExperimentArtifact,
                    experiment_id=experiment.id,
                    artifact_type="features",
                    artifact_name="selected_features",
                    data=data,
                    serialization_format="json",
                    target_id=target.id,
                )
                logger.info(f"  Migrated features for {target.name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate features for {target.name}: {e}")

        # thresholds
        thresholds_path = f"{target_dir}/thresholds.pkl"
        if os.path.exists(thresholds_path):
            try:
                thresholds = joblib.load(thresholds_path)
                data = json.dumps(thresholds).encode("utf-8")
                _save_artifact(
                    session, ExperimentArtifact,
                    experiment_id=experiment.id,
                    artifact_type="thresholds",
                    artifact_name="classification_thresholds",
                    data=data,
                    serialization_format="json",
                    target_id=target.id,
                )
                logger.info(f"  Migrated thresholds for {target.name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate thresholds for {target.name}: {e}")

        # scores_tracking
        scores_path = f"{target_dir}/scores_tracking.csv"
        if os.path.exists(scores_path):
            try:
                df = pd.read_csv(scores_path)
                data = _serialize_dataframe(df)
                _save_dataframe_record(
                    session, ExperimentData,
                    experiment_id=experiment.id,
                    data_type="scores_tracking",
                    data=data,
                    row_count=len(df),
                    column_count=len(df.columns),
                    target_id=target.id,
                )
                logger.info(f"  Migrated scores_tracking for {target.name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate scores_tracking for {target.name}: {e}")

        # feature_scores (stored in feature_selection subdirectory)
        feature_selection_dir = f"{target_dir}/feature_selection"
        feature_score_methods = [
            ("Chi2", "Chi2"),
            ("Person's_R", "Person's R"),
            ("ANOVA", "ANOVA"),
            ("Spearman's_R", "Spearman's R"),
            ("Kendall's_Tau", "Kendall's Tau"),
            ("MI", "MI"),
            ("FI", "FI"),
            ("RFE", "RFE"),
            ("SFS", "SFS"),
        ]
        for db_method, file_method in feature_score_methods:
            # Try both naming conventions
            scores_path = f"{feature_selection_dir}/{file_method}.csv"
            if not os.path.exists(scores_path):
                scores_path = f"{feature_selection_dir}/feature_scores_{db_method}.csv"
            if os.path.exists(scores_path):
                try:
                    df = pd.read_csv(scores_path)
                    data = _serialize_dataframe(df)
                    _save_dataframe_record(
                        session, ExperimentData,
                        experiment_id=experiment.id,
                        data_type=f"feature_scores_{db_method}",
                        data=data,
                        row_count=len(df),
                        column_count=len(df.columns),
                        target_id=target.id,
                    )
                    logger.info(f"  Migrated feature_scores_{db_method} for {target.name}")
                except Exception as e:
                    logger.warning(f"  Failed to migrate feature_scores_{db_method}: {e}")

        # features_final (from features.csv in feature_selection directory)
        features_final_path = f"{feature_selection_dir}/features.csv"
        if os.path.exists(features_final_path):
            try:
                df = pd.read_csv(features_final_path)
                data = _serialize_dataframe(df)
                _save_dataframe_record(
                    session, ExperimentData,
                    experiment_id=experiment.id,
                    data_type="features_final",
                    data=data,
                    row_count=len(df),
                    column_count=len(df.columns),
                    target_id=target.id,
                )
                logger.info(f"  Migrated features_final for {target.name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate features_final: {e}")

        # Models
        for file_name in os.listdir(target_dir):
            file_path = f"{target_dir}/{file_name}"
            if os.path.isdir(file_path):
                continue

            is_keras = file_name.endswith(".keras") or file_name.endswith(".h5")
            is_sklearn = file_name.endswith(".best")

            if not (is_keras or is_sklearn):
                continue

            try:
                if is_keras:
                    from keras.models import load_model
                    model = load_model(file_path)
                    data = _serialize_keras(model)
                    _save_artifact(
                        session, ExperimentArtifact,
                        experiment_id=experiment.id,
                        artifact_type="model",
                        artifact_name=file_name,
                        data=data,
                        serialization_format="keras",
                        target_id=target.id,
                    )
                else:
                    model = joblib.load(file_path)
                    data = _serialize_joblib(model)
                    _save_artifact(
                        session, ExperimentArtifact,
                        experiment_id=experiment.id,
                        artifact_type="model",
                        artifact_name=file_name,
                        data=data,
                        serialization_format="joblib",
                        target_id=target.id,
                    )
                logger.info(f"  Migrated model: {file_name} for {target.name}")
            except Exception as e:
                logger.warning(f"  Failed to migrate model {file_name}: {e}")

    # 6. Migrate predictions (requires model_selection_scores)
    for model_selection in experiment.model_selections:
        target = model_selection.target
        target_dir = f"{experiment_path}/{target.name}"

        for score in model_selection.model_selection_scores:
            model_name = score.model.name
            prediction_path = f"{target_dir}/{model_name}/prediction.csv"

            if os.path.exists(prediction_path):
                try:
                    df = pd.read_csv(prediction_path)
                    data = _serialize_dataframe(df)
                    _save_dataframe_record(
                        session, ExperimentData,
                        experiment_id=experiment.id,
                        data_type="prediction",
                        data=data,
                        row_count=len(df),
                        column_count=len(df.columns),
                        model_selection_score_id=score.id,
                    )
                    logger.info(f"  Migrated prediction for {target.name}/{model_name}")
                except Exception as e:
                    logger.warning(f"  Failed to migrate prediction for {model_name}: {e}")


def downgrade() -> None:
    """
    Downgrade is a no-op - we don't delete the migrated data.
    The files still exist on disk as backup.
    """
    pass
