"""Unwrap CatBoostWrapper from existing catboost model artifacts

Revision ID: c3d4e5f6g7h8
Revises: b2c3d4e5f6g7
Create Date: 2026-01-22

This migration updates existing catboost model artifacts to store the raw
CatBoost model instead of the CatBoostWrapper that was previously used.
"""

import io
import logging
from typing import Sequence, Union

from alembic import op
from sqlalchemy.orm import Session

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

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


# Define a stub CatBoostWrapper class for deserialization purposes
# This class must match the original structure so joblib can unpickle it
class CatBoostWrapper:
    """Stub class for deserializing old CatBoostWrapper artifacts."""

    __slots__ = ("_model", "_extra_attrs")

    def __init__(self, model=None, **kwargs):
        object.__setattr__(self, "_model", model)
        object.__setattr__(self, "_extra_attrs", kwargs)


def _inject_catboost_wrapper():
    """Inject CatBoostWrapper into the module where joblib expects to find it."""
    import lecrapaud.model_selection as model_selection_module
    model_selection_module.CatBoostWrapper = CatBoostWrapper


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


def _deserialize_joblib(data: bytes):
    """Deserialize an object using joblib."""
    import joblib
    buffer = io.BytesIO(data)
    return joblib.load(buffer)


def _is_catboost_wrapper(obj) -> bool:
    """Check if an object is a CatBoostWrapper."""
    # CatBoostWrapper has specific __slots__ = ("_model", "_extra_attrs")
    return (
        hasattr(obj, "_model") and
        hasattr(obj, "_extra_attrs") and
        obj.__class__.__name__ == "CatBoostWrapper"
    )


def _unwrap_catboost(obj):
    """Extract the inner CatBoost model from a CatBoostWrapper."""
    if _is_catboost_wrapper(obj):
        # Access the internal _model attribute
        return object.__getattribute__(obj, "_model")
    return obj


def upgrade() -> None:
    """Unwrap CatBoostWrapper from existing catboost model artifacts."""
    bind = op.get_bind()
    session = Session(bind=bind)

    try:
        # Inject stub CatBoostWrapper so joblib can deserialize old artifacts
        _inject_catboost_wrapper()

        from lecrapaud.models import ExperimentArtifact

        # Find all catboost model artifacts
        catboost_artifacts = (
            session.query(ExperimentArtifact)
            .filter(
                ExperimentArtifact.artifact_type == "model",
                ExperimentArtifact.artifact_name.like("%catboost%"),
                ExperimentArtifact.serialization_format == "joblib",
            )
            .all()
        )

        if not catboost_artifacts:
            logger.info("No catboost model artifacts found, nothing to migrate")
            return

        logger.info(f"Found {len(catboost_artifacts)} catboost artifact(s) to check")

        unwrapped_count = 0
        for artifact in catboost_artifacts:
            try:
                # Deserialize the artifact
                obj = _deserialize_joblib(artifact.data)

                # Check if it's a CatBoostWrapper
                if _is_catboost_wrapper(obj):
                    # Extract the inner model
                    inner_model = _unwrap_catboost(obj)

                    # Re-serialize just the inner model
                    new_data = _serialize_joblib(inner_model)

                    # Update the artifact
                    artifact.data = new_data
                    unwrapped_count += 1
                    logger.info(
                        f"  Unwrapped CatBoostWrapper for artifact {artifact.id} "
                        f"(experiment={artifact.experiment_id}, target={artifact.target_id})"
                    )
                else:
                    logger.info(
                        f"  Artifact {artifact.id} is not a CatBoostWrapper, skipping"
                    )

            except Exception as e:
                logger.warning(f"  Failed to process artifact {artifact.id}: {e}")
                continue

        session.commit()
        logger.info(f"Migration complete: unwrapped {unwrapped_count} CatBoostWrapper artifact(s)")

    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 downgrade() -> None:
    """
    Downgrade is a no-op - we cannot re-wrap models into CatBoostWrapper
    since CatBoostWrapper no longer exists in the codebase.
    """
    pass
