import os
import mlflow
from datetime import datetime

# -----------------------------
# AUTO CONFIG (JupyterHub)
# -----------------------------
mlflow.set_tracking_uri(
    os.environ.get(
        "MLFLOW_TRACKING_URI",
        "http://mlflow.mlops.svc.cluster.local:5000"
    )
)

STUDENT_NAME = os.environ.get("JUPYTERHUB_USER", "Student")

# One experiment per student (recommended)
experiment_name = f"{STUDENT_NAME}_AUTO_TRACKING"
artifact_location = f"s3://niet-minio-bucket/{STUDENT_NAME}"

# Create or load experiment
try:
    experiment_id = mlflow.create_experiment(
        name=experiment_name,
        artifact_location=artifact_location
    )
    print(f"Created experiment: {experiment_name}")
except:
    experiment = mlflow.get_experiment_by_name(experiment_name)
    experiment_id = experiment.experiment_id
    print(f"Using existing experiment: {experiment_name}")

mlflow.set_experiment(experiment_name)

# -----------------------------
# ENABLE FULL AUTO-LOGGING
# -----------------------------
mlflow.autolog(log_models=True)

# -----------------------------
# AUTO RUN CONTEXT
# -----------------------------
with mlflow.start_run(run_name=f"{STUDENT_NAME}_AUTO_RUN_{datetime.now().strftime('%H%M%S')}"):

    # Minimal manual metadata (optional)
    mlflow.set_tag("student_name", STUDENT_NAME)
    mlflow.set_tag("mode", "auto_tracking")
    mlflow.set_tag("platform", "jupyterhub")

    # -----------------------------
    # STUDENT WORK STARTS HERE
    # (ANY ML / GENAI / DATA CODE)
    # -----------------------------

    # Example: student creates outputs normally
    os.makedirs("outputs", exist_ok=True)

    with open("outputs/result.txt", "w") as f:
        f.write("This file was automatically captured by MLflow.")

    # Auto-log all outputs
    mlflow.log_artifacts("outputs")

    run_id = mlflow.active_run().info.run_id

# -----------------------------
# CONFIRMATION
# -----------------------------
print("\n AUTO TRACKING ENABLED SUCCESSFULLY")
print(f"Experiment ID : {experiment_id}")
print(f"Run ID        : {run_id}")
print("\nMinIO Path:")
print(f"{STUDENT_NAME}/{experiment_id}/{run_id}/artifacts/")