Ml Ops Engineer

#!/usr/bin/env python3
"""Detect data and model drift by comparing reference and current distributions.

Reads two CSV files (reference/baseline and current/production) and
computes distribution shift per feature using statistical tests.  Reports
drift scores, flags drifted features, and suggests remediation.

Usage:
    python drift_detector.py --reference train_data.csv --current prod_data.csv
    python drift_detector.py --reference baseline.csv --current latest.csv --threshold 0.1 --json
    python drift_detector.py --reference ref.csv --current cur.csv --columns feature_1 feature_2 feature_3
"""

import argparse
import csv
import json
import math
import os
import sys
from collections import Counter


# ---------------------------------------------------------------------------
# Statistical helpers (standard library only)
# ---------------------------------------------------------------------------

def _is_numeric(value: str) -> bool:
    try:
        float(value)
        return True
    except (ValueError, TypeError):
        return False


def _to_floats(values: list) -> list:
    return [float(v) for v in values if _is_numeric(str(v)) and str(v).strip()]


def _mean(vals: list) -> float:
    return sum(vals) / len(vals) if vals else 0.0


def _std(vals: list) -> float:
    if len(vals) < 2:
        return 0.0
    m = _mean(vals)
    return math.sqrt(sum((x - m) ** 2 for x in vals) / (len(vals) - 1))


def _ks_statistic(ref: list, cur: list) -> float:
    """Approximate Kolmogorov-Smirnov statistic for two samples."""
    all_vals = sorted(set(ref + cur))
    if not all_vals:
        return 0.0

    n_ref, n_cur = len(ref), len(cur)
    if n_ref == 0 or n_cur == 0:
        return 0.0

    ref_sorted = sorted(ref)
    cur_sorted = sorted(cur)

    max_diff = 0.0
    i, j = 0, 0
    for val in all_vals:
        while i < n_ref and ref_sorted[i] <= val:
            i += 1
        while j < n_cur and cur_sorted[j] <= val:
            j += 1
        diff = abs(i / n_ref - j / n_cur)
        if diff > max_diff:
            max_diff = diff

    return max_diff


def _psi(ref: list, cur: list, bins: int = 10) -> float:
    """Population Stability Index for numeric features."""
    if not ref or not cur:
        return 0.0

    all_vals = ref + cur
    mn, mx = min(all_vals), max(all_vals)
    if mn == mx:
        return 0.0

    step = (mx - mn) / bins
    epsilon = 1e-6

    def _bin_counts(vals):
        counts = [0] * bins
        for v in vals:
            idx = min(int((v - mn) / step), bins - 1)
            counts[idx] += 1
        total = len(vals)
        return [(c / total) + epsilon for c in counts]

    ref_pcts = _bin_counts(ref)
    cur_pcts = _bin_counts(cur)

    psi_val = 0.0
    for r, c in zip(ref_pcts, cur_pcts):
        psi_val += (c - r) * math.log(c / r)

    return psi_val


def _chi_square_categorical(ref: list, cur: list) -> float:
    """Chi-square divergence for categorical features."""
    ref_counts = Counter(ref)
    cur_counts = Counter(cur)
    all_keys = set(ref_counts.keys()) | set(cur_counts.keys())

    n_ref, n_cur = len(ref), len(cur)
    if n_ref == 0 or n_cur == 0:
        return 0.0

    chi2 = 0.0
    for key in all_keys:
        expected = ref_counts.get(key, 0) / n_ref
        observed = cur_counts.get(key, 0) / n_cur
        if expected > 0:
            chi2 += (observed - expected) ** 2 / expected

    return chi2


# ---------------------------------------------------------------------------
# Drift analysis
# ---------------------------------------------------------------------------

def analyze_feature(name: str, ref_values: list, cur_values: list, threshold: float) -> dict:
    ref_non_null = [v for v in ref_values if v is not None and str(v).strip()]
    cur_non_null = [v for v in cur_values if v is not None and str(v).strip()]

    ref_numeric = _to_floats(ref_non_null)
    cur_numeric = _to_floats(cur_non_null)
    is_numeric = (len(ref_numeric) > len(ref_non_null) * 0.8) if ref_non_null else False

    result = {
        "feature": name,
        "data_type": "numeric" if is_numeric else "categorical",
        "ref_count": len(ref_non_null),
        "cur_count": len(cur_non_null),
    }

    if is_numeric and ref_numeric and cur_numeric:
        result["ref_mean"] = round(_mean(ref_numeric), 4)
        result["cur_mean"] = round(_mean(cur_numeric), 4)
        result["ref_std"] = round(_std(ref_numeric), 4)
        result["cur_std"] = round(_std(cur_numeric), 4)
        result["mean_shift"] = round(result["cur_mean"] - result["ref_mean"], 4)

        ks = _ks_statistic(ref_numeric, cur_numeric)
        psi = _psi(ref_numeric, cur_numeric)

        result["ks_statistic"] = round(ks, 4)
        result["psi"] = round(psi, 4)
        result["drift_score"] = round(max(ks, psi), 4)
        result["is_drifted"] = result["drift_score"] > threshold

        # PSI interpretation
        if psi < 0.1:
            result["psi_interpretation"] = "stable"
        elif psi < 0.2:
            result["psi_interpretation"] = "moderate_shift"
        else:
            result["psi_interpretation"] = "significant_shift"
    else:
        chi2 = _chi_square_categorical(ref_non_null, cur_non_null)
        result["chi_square"] = round(chi2, 4)
        result["drift_score"] = round(chi2, 4)
        result["is_drifted"] = chi2 > threshold

        # Check for new categories
        ref_cats = set(ref_non_null)
        cur_cats = set(cur_non_null)
        new_cats = cur_cats - ref_cats
        missing_cats = ref_cats - cur_cats
        if new_cats:
            result["new_categories"] = list(new_cats)[:10]
        if missing_cats:
            result["missing_categories"] = list(missing_cats)[:10]

    if result["is_drifted"]:
        if is_numeric:
            result["recommendation"] = "Investigate distribution shift; consider retraining if model performance degrades."
        else:
            result["recommendation"] = "New/missing categories detected; update encoding and retrain."

    return result


def detect_drift(ref_data: list, cur_data: list, columns: list = None, threshold: float = 0.1) -> dict:
    if columns is None:
        columns = list(ref_data[0].keys()) if ref_data else []

    results = []
    for col in columns:
        ref_vals = [row.get(col) for row in ref_data]
        cur_vals = [row.get(col) for row in cur_data]
        results.append(analyze_feature(col, ref_vals, cur_vals, threshold))

    drifted = [r for r in results if r["is_drifted"]]
    results.sort(key=lambda x: x["drift_score"], reverse=True)

    return {
        "total_features": len(results),
        "drifted_features": len(drifted),
        "drift_rate": round(len(drifted) / len(results) * 100, 1) if results else 0,
        "threshold": threshold,
        "alert": len(drifted) > len(results) * 0.3,
        "alert_message": f"{len(drifted)}/{len(results)} features drifted (>{threshold} threshold). Consider model retraining." if len(drifted) > len(results) * 0.3 else "",
        "features": results,
    }


def main():
    parser = argparse.ArgumentParser(description="Detect data and model drift between reference and current datasets.")
    parser.add_argument("--reference", required=True, help="Path to reference/baseline CSV file")
    parser.add_argument("--current", required=True, help="Path to current/production CSV file")
    parser.add_argument("--columns", nargs="*", help="Specific columns to check (default: all)")
    parser.add_argument("--threshold", type=float, default=0.1, help="Drift threshold (default: 0.1)")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    args = parser.parse_args()

    for path, label in [(args.reference, "Reference"), (args.current, "Current")]:
        if not os.path.exists(path):
            print(f"Error: {label} file not found: {path}", file=sys.stderr)
            sys.exit(1)

    with open(args.reference, "r", newline="") as f:
        ref_data = list(csv.DictReader(f))
    with open(args.current, "r", newline="") as f:
        cur_data = list(csv.DictReader(f))

    if not ref_data or not cur_data:
        print("Error: Both files must contain data rows.", file=sys.stderr)
        sys.exit(1)

    result = detect_drift(ref_data, cur_data, args.columns, args.threshold)

    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print("Drift Detection Report")
        print("=" * 65)
        alert_str = " [ALERT]" if result["alert"] else ""
        print(f"Features: {result['total_features']}  |  Drifted: {result['drifted_features']} ({result['drift_rate']}%)  |  Threshold: {result['threshold']}{alert_str}")
        if result["alert_message"]:
            print(f"\n  {result['alert_message']}")
        print()
        print(f"{'Feature':<25} {'Type':<12} {'Drift Score':>12} {'Drifted':<8}")
        print("-" * 65)
        for f in result["features"]:
            marker = "[!!]" if f["is_drifted"] else "[ ]"
            print(f"  {f['feature']:<23} {f['data_type']:<12} {f['drift_score']:>12.4f} {marker}")
            if f.get("recommendation"):
                print(f"    -> {f['recommendation']}")

    sys.exit(1 if result["alert"] else 0)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Track model versions, metadata, and lifecycle stage in a local JSON registry.

Maintains a structured model registry.  Each entry records model name,
version, metrics, parameters, stage (staging/production/archived), and
timestamps.  Supports registering, promoting, listing, and comparing models.

Usage:
    python model_registry.py register --name fraud_detector --version v2.3 --metrics '{"f1":0.91,"auc":0.95}' --params '{"n_estimators":200}'
    python model_registry.py promote --name fraud_detector --version v2.3 --stage production
    python model_registry.py list --name fraud_detector
    python model_registry.py list --stage production --json
    python model_registry.py compare --name fraud_detector --versions v2.2 v2.3
"""

import argparse
import json
import os
import sys
from datetime import datetime


DEFAULT_REGISTRY = "model_registry.json"
VALID_STAGES = {"development", "staging", "production", "archived", "retired"}


def _load_registry(path: str) -> list:
    if not os.path.exists(path):
        return []
    with open(path, "r") as f:
        return json.load(f)


def _save_registry(registry: list, path: str):
    with open(path, "w") as f:
        json.dump(registry, f, indent=2)


def _find_entry(registry: list, name: str, version: str):
    for entry in registry:
        if entry["name"] == name and entry["version"] == version:
            return entry
    return None


def cmd_register(args):
    registry = _load_registry(args.registry)

    existing = _find_entry(registry, args.name, args.version)
    if existing:
        print(f"Error: Model '{args.name}' version '{args.version}' already registered. Use 'promote' to update stage.", file=sys.stderr)
        sys.exit(1)

    try:
        metrics = json.loads(args.metrics) if args.metrics else {}
    except json.JSONDecodeError:
        print("Error: --metrics must be valid JSON.", file=sys.stderr)
        sys.exit(1)

    try:
        params = json.loads(args.params) if args.params else {}
    except json.JSONDecodeError:
        print("Error: --params must be valid JSON.", file=sys.stderr)
        sys.exit(1)

    entry = {
        "name": args.name,
        "version": args.version,
        "stage": args.stage or "development",
        "registered_at": datetime.now().strftime("%Y-%m-%dT%H:%M:%S"),
        "updated_at": datetime.now().strftime("%Y-%m-%dT%H:%M:%S"),
        "metrics": metrics,
        "parameters": params,
        "tags": [t.strip() for t in args.tags.split(",")] if args.tags else [],
        "description": args.description or "",
        "artifact_path": args.artifact or "",
    }

    registry.append(entry)
    _save_registry(registry, args.registry)

    if args.json:
        print(json.dumps(entry, indent=2))
    else:
        print(f"Registered: {args.name} {args.version} (stage: {entry['stage']})")
        if metrics:
            print(f"  Metrics: {json.dumps(metrics)}")


def cmd_promote(args):
    registry = _load_registry(args.registry)
    entry = _find_entry(registry, args.name, args.version)

    if not entry:
        print(f"Error: Model '{args.name}' version '{args.version}' not found.", file=sys.stderr)
        sys.exit(1)

    if args.stage not in VALID_STAGES:
        print(f"Error: Invalid stage '{args.stage}'. Valid: {', '.join(sorted(VALID_STAGES))}.", file=sys.stderr)
        sys.exit(1)

    old_stage = entry["stage"]

    # If promoting to production, archive the current production version
    if args.stage == "production":
        for e in registry:
            if e["name"] == args.name and e["stage"] == "production" and e["version"] != args.version:
                e["stage"] = "archived"
                e["updated_at"] = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")

    entry["stage"] = args.stage
    entry["updated_at"] = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
    _save_registry(registry, args.registry)

    if args.json:
        print(json.dumps(entry, indent=2))
    else:
        print(f"Promoted: {args.name} {args.version}: {old_stage} -> {args.stage}")


def cmd_list(args):
    registry = _load_registry(args.registry)
    if not registry:
        print("Registry is empty.")
        return

    filtered = registry
    if args.name:
        filtered = [e for e in filtered if e["name"] == args.name]
    if args.stage:
        filtered = [e for e in filtered if e["stage"] == args.stage]

    filtered.sort(key=lambda e: e["updated_at"], reverse=True)

    if args.json:
        print(json.dumps(filtered, indent=2))
    else:
        print(f"{'Name':<25} {'Version':<10} {'Stage':<14} {'Updated':<20} {'Metrics'}")
        print("-" * 90)
        for e in filtered:
            metrics_str = ", ".join(f"{k}={v}" for k, v in e.get("metrics", {}).items())
            print(f"{e['name']:<25} {e['version']:<10} {e['stage']:<14} {e['updated_at']:<20} {metrics_str}")


def cmd_compare(args):
    registry = _load_registry(args.registry)
    entries = [e for e in registry if e["name"] == args.name and e["version"] in args.versions]

    if len(entries) < 2:
        print(f"Error: Need at least 2 versions to compare. Found {len(entries)} for '{args.name}'.", file=sys.stderr)
        sys.exit(1)

    all_metrics = set()
    all_params = set()
    for e in entries:
        all_metrics.update(e.get("metrics", {}).keys())
        all_params.update(e.get("parameters", {}).keys())

    if args.json:
        print(json.dumps({"name": args.name, "entries": entries, "metric_keys": sorted(all_metrics)}, indent=2))
    else:
        print(f"Model Comparison: {args.name}")
        print("=" * 60)

        header = f"{'Attribute':<20}"
        for e in entries:
            header += f" {e['version']:<18}"
        print(header)
        print("-" * len(header))

        print(f"{'stage':<20}" + "".join(f" {e['stage']:<18}" for e in entries))
        print(f"{'updated_at':<20}" + "".join(f" {e['updated_at']:<18}" for e in entries))

        if all_metrics:
            print("\nMetrics:")
            for m in sorted(all_metrics):
                vals = [e.get("metrics", {}).get(m) for e in entries]
                best = max((v for v in vals if v is not None), default=None)
                row = f"  {m:<18}"
                for v in vals:
                    marker = " *" if v == best and v is not None else ""
                    row += f" {str(v if v is not None else '-'):<16}{marker}"
                print(row)
            print("  * = best")

        if all_params:
            print("\nParameters:")
            for p in sorted(all_params):
                row = f"  {p:<18}"
                for e in entries:
                    val = e.get("parameters", {}).get(p, "-")
                    row += f" {str(val):<18}"
                print(row)


def main():
    parser = argparse.ArgumentParser(description="Manage a local model version registry.")
    parser.add_argument("--registry", default=DEFAULT_REGISTRY, help=f"Path to registry file (default: {DEFAULT_REGISTRY})")
    parser.add_argument("--json", action="store_true", help="Output as JSON")

    sub = parser.add_subparsers(dest="command", help="Command")

    reg = sub.add_parser("register", help="Register a new model version")
    reg.add_argument("--name", required=True, help="Model name")
    reg.add_argument("--version", required=True, help="Model version")
    reg.add_argument("--metrics", help="Metrics as JSON string")
    reg.add_argument("--params", help="Parameters as JSON string")
    reg.add_argument("--stage", default="development", help="Initial stage")
    reg.add_argument("--tags", help="Comma-separated tags")
    reg.add_argument("--description", help="Model description")
    reg.add_argument("--artifact", help="Path to model artifact")

    promo = sub.add_parser("promote", help="Change model stage")
    promo.add_argument("--name", required=True, help="Model name")
    promo.add_argument("--version", required=True, help="Model version")
    promo.add_argument("--stage", required=True, help=f"Target stage: {', '.join(sorted(VALID_STAGES))}")

    lst = sub.add_parser("list", help="List registered models")
    lst.add_argument("--name", help="Filter by model name")
    lst.add_argument("--stage", help="Filter by stage")

    cmp = sub.add_parser("compare", help="Compare model versions")
    cmp.add_argument("--name", required=True, help="Model name")
    cmp.add_argument("--versions", nargs="+", required=True, help="Versions to compare")

    args = parser.parse_args()
    if not args.command:
        parser.print_help()
        sys.exit(1)

    if args.command == "register":
        cmd_register(args)
    elif args.command == "promote":
        cmd_promote(args)
    elif args.command == "list":
        cmd_list(args)
    elif args.command == "compare":
        cmd_compare(args)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Validate ML pipeline definitions for completeness, ordering, and best practices.

Reads a pipeline definition file (JSON) and checks for missing stages,
incorrect ordering, missing validation gates, absent monitoring hooks,
and configuration issues.

Usage:
    python pipeline_validator.py --pipeline pipeline.json
    python pipeline_validator.py --pipeline pipeline.json --strict --json

Pipeline definition format:
{
    "name": "fraud_detection_training",
    "stages": [
        {
            "name": "data_ingestion",
            "type": "data",
            "inputs": ["raw_transactions"],
            "outputs": ["cleaned_data"],
            "timeout_minutes": 30
        },
        {
            "name": "feature_engineering",
            "type": "transform",
            "inputs": ["cleaned_data"],
            "outputs": ["feature_set"],
            "validation": {"null_check": true, "schema_check": true}
        },
        {
            "name": "model_training",
            "type": "train",
            "inputs": ["feature_set"],
            "outputs": ["trained_model"],
            "parameters": {"algorithm": "xgboost"}
        },
        {
            "name": "model_evaluation",
            "type": "evaluate",
            "inputs": ["trained_model", "test_data"],
            "outputs": ["metrics"],
            "gate": {"metric": "f1", "threshold": 0.85}
        },
        {
            "name": "model_deployment",
            "type": "deploy",
            "inputs": ["trained_model", "metrics"],
            "outputs": ["serving_endpoint"],
            "rollback": true
        }
    ],
    "schedule": "0 2 * * 1",
    "notifications": {"on_failure": "ml-team@company.com"}
}
"""

import argparse
import json
import os
import re
import sys


# ---------------------------------------------------------------------------
# Validation rules
# ---------------------------------------------------------------------------

REQUIRED_STAGE_TYPES = {"data", "transform", "train", "evaluate"}
RECOMMENDED_STAGE_TYPES = {"deploy", "monitor"}
VALID_STAGE_TYPES = {"data", "transform", "train", "evaluate", "deploy", "monitor", "validate", "register", "test"}

EXPECTED_ORDER = ["data", "transform", "train", "evaluate", "deploy", "monitor"]
CRON_PATTERN = re.compile(r"^(\*|[0-9,\-\/]+)\s+(\*|[0-9,\-\/]+)\s+(\*|[0-9,\-\/]+)\s+(\*|[0-9,\-\/]+)\s+(\*|[0-9,\-\/]+)$")


def _validate_pipeline(pipeline: dict, strict: bool = False) -> list:
    issues = []
    name = pipeline.get("name", "(unnamed)")

    # Top-level checks
    if not pipeline.get("name"):
        issues.append({"severity": "error", "rule": "MISSING_NAME", "message": "Pipeline is missing a 'name' field."})

    stages = pipeline.get("stages", [])
    if not stages:
        issues.append({"severity": "error", "rule": "NO_STAGES", "message": "Pipeline has no stages defined."})
        return issues

    # Check required stage types
    stage_types = [s.get("type", "unknown") for s in stages]
    for req in REQUIRED_STAGE_TYPES:
        if req not in stage_types:
            issues.append({
                "severity": "error",
                "rule": "MISSING_REQUIRED_STAGE",
                "message": f"Pipeline is missing a required '{req}' stage.",
            })

    for rec in RECOMMENDED_STAGE_TYPES:
        if rec not in stage_types and strict:
            issues.append({
                "severity": "warning",
                "rule": "MISSING_RECOMMENDED_STAGE",
                "message": f"Pipeline is missing recommended '{rec}' stage.",
            })

    # Per-stage validation
    all_outputs = set()
    stage_names = []

    for i, stage in enumerate(stages):
        sname = stage.get("name", f"stage_{i}")
        stype = stage.get("type", "unknown")
        stage_names.append(sname)

        # Required fields
        if not stage.get("name"):
            issues.append({"severity": "error", "rule": "STAGE_MISSING_NAME", "message": f"Stage #{i+1} is missing a name."})

        if stype not in VALID_STAGE_TYPES:
            issues.append({"severity": "warning", "rule": "UNKNOWN_STAGE_TYPE", "message": f"Stage '{sname}' has unknown type '{stype}'."})

        # Inputs/outputs
        inputs = stage.get("inputs", [])
        outputs = stage.get("outputs", [])

        if not outputs:
            issues.append({"severity": "warning", "rule": "NO_OUTPUTS", "message": f"Stage '{sname}' defines no outputs."})

        # Check that inputs are produced by earlier stages (except first stage)
        if i > 0 and inputs:
            for inp in inputs:
                if inp not in all_outputs:
                    issues.append({
                        "severity": "warning",
                        "rule": "UNRESOLVED_INPUT",
                        "message": f"Stage '{sname}' requires input '{inp}' which is not produced by any earlier stage.",
                    })

        for out in outputs:
            if out in all_outputs:
                issues.append({
                    "severity": "warning",
                    "rule": "DUPLICATE_OUTPUT",
                    "message": f"Output '{out}' is produced by multiple stages.",
                })
            all_outputs.add(out)

        # Evaluation gate
        if stype == "evaluate" and not stage.get("gate"):
            issues.append({
                "severity": "warning" if not strict else "error",
                "rule": "MISSING_EVAL_GATE",
                "message": f"Evaluation stage '{sname}' has no quality gate defined. Models may deploy without validation.",
            })

        # Deploy should have rollback
        if stype == "deploy" and not stage.get("rollback"):
            issues.append({
                "severity": "warning",
                "rule": "NO_ROLLBACK",
                "message": f"Deploy stage '{sname}' has no rollback configuration.",
            })

        # Timeout
        if strict and not stage.get("timeout_minutes"):
            issues.append({
                "severity": "info",
                "rule": "NO_TIMEOUT",
                "message": f"Stage '{sname}' has no timeout configured.",
            })

    # Ordering check
    order_positions = {}
    for i, stype in enumerate(stage_types):
        if stype in EXPECTED_ORDER:
            order_positions[stype] = i

    for a, b in zip(EXPECTED_ORDER, EXPECTED_ORDER[1:]):
        if a in order_positions and b in order_positions:
            if order_positions[a] > order_positions[b]:
                issues.append({
                    "severity": "error",
                    "rule": "STAGE_ORDER",
                    "message": f"Stage type '{a}' appears after '{b}'; expected order: {' -> '.join(EXPECTED_ORDER)}.",
                })

    # Duplicate stage names
    seen_names = set()
    for sn in stage_names:
        if sn in seen_names:
            issues.append({"severity": "error", "rule": "DUPLICATE_STAGE_NAME", "message": f"Duplicate stage name: '{sn}'."})
        seen_names.add(sn)

    # Schedule validation
    schedule = pipeline.get("schedule")
    if schedule and not CRON_PATTERN.match(schedule.strip()):
        issues.append({"severity": "warning", "rule": "INVALID_SCHEDULE", "message": f"Schedule '{schedule}' does not match cron format."})

    # Notifications
    if not pipeline.get("notifications"):
        issues.append({"severity": "warning", "rule": "NO_NOTIFICATIONS", "message": "Pipeline has no notification configuration for failures."})

    return issues


def validate(pipeline: dict, strict: bool = False) -> dict:
    issues = _validate_pipeline(pipeline, strict)
    errors = sum(1 for i in issues if i["severity"] == "error")
    warnings = sum(1 for i in issues if i["severity"] == "warning")
    infos = sum(1 for i in issues if i["severity"] == "info")

    return {
        "pipeline_name": pipeline.get("name", "(unnamed)"),
        "total_stages": len(pipeline.get("stages", [])),
        "total_issues": len(issues),
        "errors": errors,
        "warnings": warnings,
        "info": infos,
        "valid": errors == 0,
        "issues": issues,
    }


def main():
    parser = argparse.ArgumentParser(description="Validate ML pipeline definitions for completeness and best practices.")
    parser.add_argument("--pipeline", required=True, help="Path to pipeline definition JSON file")
    parser.add_argument("--strict", action="store_true", help="Enable strict validation (recommended and info checks)")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    args = parser.parse_args()

    if not os.path.exists(args.pipeline):
        print(f"Error: File not found: {args.pipeline}", file=sys.stderr)
        sys.exit(1)

    with open(args.pipeline, "r") as f:
        pipeline = json.load(f)

    if not isinstance(pipeline, dict):
        print("Error: Pipeline definition must be a JSON object.", file=sys.stderr)
        sys.exit(1)

    result = validate(pipeline, args.strict)

    if args.json:
        print(json.dumps(result, indent=2))
    else:
        status = "PASS" if result["valid"] else "FAIL"
        print("ML Pipeline Validation Report")
        print("=" * 60)
        print(f"Pipeline: {result['pipeline_name']}  |  Stages: {result['total_stages']}")
        print(f"Status: [{status}]  |  Errors: {result['errors']}  Warnings: {result['warnings']}  Info: {result['info']}")
        print()

        if not result["issues"]:
            print("Pipeline definition passes all validation checks.")
        else:
            for issue in result["issues"]:
                sev = issue["severity"].upper()
                print(f"  [{sev}] {issue['rule']}")
                print(f"    {issue['message']}")

    sys.exit(1 if result["errors"] > 0 else 0)


if __name__ == "__main__":
    main()