# start_mlflow.py
import os
import sys
import subprocess

# 🔥 Set PYTHONPATH to include /app (for sitecustomize.py)
os.environ["PYTHONPATH"] = "/app:" + os.environ.get("PYTHONPATH", "")

# 🔥 Get environment variables
backend_uri = os.getenv("DATABASE_URL")
default_artifact_root = os.getenv("MLFLOW_HF_ARTIFACT_REPO", "hf://models/omarelrayes/mlflow-artifacts")

if not backend_uri:
    print("❌ DATABASE_URL not set!")
    sys.exit(1)

print(f"📊 Backend Store URI: {backend_uri}")
print(f"📦 Default Artifact Root: {default_artifact_root}")

# 🔥 Start MLflow Server as subprocess
# 🔥 sitecustomize.py will auto-register the plugin in this process!
cmd = [
    sys.executable, "-m", "mlflow", "server",
    "--host", "0.0.0.0",
    "--port", "7860",
    "--backend-store-uri", backend_uri,
    "--default-artifact-root", default_artifact_root,
    "--serve-artifacts",
    "--workers", "1",
    "--allowed-hosts", "*",
    "--cors-allowed-origins", "*"
]

print(f"🚀 Starting MLflow Server...")
print(f"📝 Command: {' '.join(cmd)}")
sys.exit(subprocess.call(cmd))