"""MLflow model registry helpers."""

import mlflow
import mlflow.pyfunc
import mlflow.sklearn
import mlflow.xgboost
import pandas as pd
from mlflow.tracking import MlflowClient
from sqlalchemy import Engine, text


def register_model(
    run_id: str,
    model_name: str,
    model_uri_suffix: str,
    stage: str = "Production",
) -> None:
    """Register a logged model and transition it to the given stage."""
    client = MlflowClient()
    model_uri = f"runs:/{run_id}/{model_uri_suffix}"

    # Create registered model if it doesn't exist
    try:
        client.create_registered_model(model_name)
    except Exception:
        pass  # Already exists

    version = client.create_model_version(
        name=model_name,
        source=model_uri,
        run_id=run_id,
    )
    client.transition_model_version_stage(
        name=model_name,
        version=version.version,
        stage=stage,
        archive_existing_versions=True,
    )
    print(f"  Registered '{model_name}' version {version.version} → {stage}")


def load_production_model(model_name: str):
    """Load the Production-stage model from the MLflow registry."""
    return mlflow.pyfunc.load_model(f"models:/{model_name}/Production")


def write_scores_to_warehouse(
    engine: Engine,
    scores_df: pd.DataFrame,
) -> int:
    """
    Bulk-update dim_customers with ML-derived scores.

    scores_df must have column 'customer_unique_id' plus any subset of:
    churn_probability, segment_label, rfm_segment, clv_score.
    Uses a single connection so temp table is visible across all statements.
    """
    possible_cols = ["customer_unique_id", "churn_probability", "segment_label", "rfm_segment", "clv_score"]
    cols = [c for c in possible_cols if c in scores_df.columns]
    df = scores_df[cols].drop_duplicates("customer_unique_id").copy()

    with engine.begin() as conn:
        conn.execute(text("DROP TABLE IF EXISTS _scores_staging"))
        conn.execute(text("""
            CREATE TEMP TABLE _scores_staging (
                customer_unique_id TEXT,
                churn_probability  NUMERIC(5,4),
                segment_label      INTEGER,
                rfm_segment        TEXT,
                clv_score          NUMERIC(10,4)
            )
        """))
        # Insert using pandas within same connection
        ncols = len(cols)
        chunk = max(1, 65535 // ncols)
        df.to_sql(
            "_scores_staging",
            con=conn,
            if_exists="append",
            index=False,
            method="multi",
            chunksize=chunk,
        )
        conn.execute(text("""
            UPDATE warehouse.dim_customers dc
            SET
                churn_probability = COALESCE(s.churn_probability, dc.churn_probability),
                segment_label     = COALESCE(s.segment_label, dc.segment_label),
                rfm_segment       = COALESCE(s.rfm_segment, dc.rfm_segment),
                clv_score         = COALESCE(s.clv_score, dc.clv_score),
                _updated_at       = now()
            FROM _scores_staging s
            WHERE dc.customer_unique_id = s.customer_unique_id
        """))
        result = conn.execute(text("SELECT COUNT(*) FROM _scores_staging"))
        return result.scalar()