import os
import sys
import platform
import mlflow
from typing import Dict, Any, Optional

def setup_mlflow(config: dict) -> str:
    """Sets up local MLflow tracking server URI and artifact directory."""
    tracking_uri = config.get("tracking", {}).get("tracking_uri", "sqlite:///mlflow.db")
    mlflow.set_tracking_uri(tracking_uri)

    # Disable annoying hints
    os.environ["MLFLOW_DISABLE_AGENT_HINT"] = "1"
    
    # Pre-create standard experiments
    experiments = config.get("tracking", {}).get("experiments", {
        "ocr_comparison": "ocr_comparison",
        "doc_classification": "doc_classification",
        "field_extraction": "field_extraction",
        "robustness_tests": "robustness_tests",
        "end_to_end": "end_to_end",
        "threshold_analysis": "threshold_analysis"
    })
    for exp_key, exp_name in experiments.items():
        try:
            mlflow.create_experiment(exp_name)
        except Exception:
            pass  # Already exists

    return tracking_uri

def get_standard_run_tags(stage: str, model_family: str, purpose: str = "baseline") -> Dict[str, str]:
    """Generates standard run tags required by project specs."""
    return {
        "stage": stage,
        "model_family": model_family,
        "dataset_version": "1.0.0",
        "author": "doc_intel_agent",
        "purpose": purpose,
        "os": platform.system(),
        "python_version": platform.python_version()
    }

def get_standard_run_params(seed: int = 42, **kwargs) -> Dict[str, Any]:
    """Generates standard parameter dictionary."""
    params = {
        "seed": seed,
        "hardware": "cuda" if False else "cpu",
        "python": platform.python_version()
    }
    params.update(kwargs)
    return params