#!/usr/bin/env python3

import re
import sys
import os

MDSTAT_PATH = "/proc/mdstat"

def parse_mdstat():
    if not os.path.exists(MDSTAT_PATH):
        return [], "CRITICAL - /proc/mdstat not found", 2

    arrays = {}
    current_array = None

    try:
        with open(MDSTAT_PATH, "r") as f:
            for line in f:
                # RAID device line (e.g., "md127 : active raid1 sda1[0] sdb1[1]")
                match = re.match(r'^(md\d+)\s*:\s*(.*)', line)
                if match:
                    current_array = match.group(1)
                    arrays[current_array] = {'status': match.group(2), 'lines': [line.strip()]}
                elif current_array:
                    arrays[current_array]['lines'].append(line.strip())
    except Exception as e:
        return [], f"CRITICAL - Error reading /proc/mdstat: {e}", 2

    return arrays, None, None

def analyze_arrays(arrays):
    total = len(arrays)
    degraded = []
    recovering = []
    resyncing = []

    for md, data in arrays.items():
        combined = " ".join(data['lines'])
        if re.search(r'\[.*_.*\]', combined):
            degraded.append(md)
        elif "recovery" in combined:
            recovering.append(md)
        elif "resync" in combined:
            resyncing.append(md)

    return total, degraded, recovering, resyncing

def main():
    arrays, error_msg, exit_code = parse_mdstat()

    if error_msg:
        print(error_msg)
        sys.exit(exit_code)

    total, degraded, recovering, resyncing = analyze_arrays(arrays)

    if recovering:
        print(f"WARNING - Checked {total} arrays, recovering: {', '.join(recovering)}")
        sys.exit(1)
    elif resyncing:
        print(f"WARNING - Checked {total} arrays, resync: {', '.join(resyncing)}")
        sys.exit(1)
    elif not degraded:
        print(f"OK - Checked {total} arrays. All arrays healthy.")
        sys.exit(0)
    else:
        print(f"CRITICAL - Checked {total} arrays, degraded: {', '.join(degraded)}")
        sys.exit(2)

if __name__ == "__main__":
    main()