#!/usr/bin/env bash
#
# check-hardware-health.sh — report disk SMART health and hardware temperatures.
# Full documentation: see check-hardware-health.md in this same folder.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib/common.sh
source "${SCRIPT_DIR}/../lib/common.sh"

SCRIPT_HELP=$(cat <<'HELP'
Usage: ./System/check-hardware-health.sh [OPTIONS]

Checks physical disks and their SMART health, reports device/media type,
and reads component temperatures from lm-sensors when available. High
temperatures and detected disk problems are highlighted in red in the report.

Options:
  -o, --output PATH   Save the report to PATH; otherwise ask interactively
    -f, --format FORMAT Report format: txt, html, or json
      --force         Overwrite an existing output file without prompting
  -h, --help          Show this help message and exit

Examples:
  ./System/check-hardware-health.sh
  ./System/check-hardware-health.sh --output "$HOME/hardware-health.html"
    ./System/check-hardware-health.sh --output "$HOME/hardware-health.json"
HELP
)

common::show_help_if_requested "$@"

OUTPUT_PATH=""
OUTPUT_FORMAT=""
FORCE=0
while [[ $# -gt 0 ]]; do
    case "$1" in
        -o|--output)
            [[ $# -ge 2 ]] || common::die "$1 requires a path"
            OUTPUT_PATH="$2"
            shift 2
            ;;
        -f|--format)
            [[ $# -ge 2 ]] || common::die "$1 requires txt, html, or json"
            OUTPUT_FORMAT="${2,,}"
            case "$OUTPUT_FORMAT" in
                txt|html|json) ;;
                *) common::die "unsupported report format: $2 (use txt, html, or json)" ;;
            esac
            shift 2
            ;;
        --force)
            FORCE=1
            shift
            ;;
        *) common::die "unknown argument: $1" ;;
    esac
done

common::require_cmd python3
common::require_cmd lsblk

if [[ "$EUID" -ne 0 ]]; then
    common::log "Warning: without root privileges, SMART data may be unavailable and not all checks will run."
    common::confirm "Continue without root privileges?" || common::die "Hardware health check cancelled." 130
fi

python3 -c "$(cat <<'PYCODE'
import datetime
import html
import json
import os
import pathlib
import pwd
import shutil
import subprocess
import sys
import tempfile


def run(command):
    if shutil.which(command[0]) is None:
        return None
    try:
        result = subprocess.run(command, capture_output=True, text=True, timeout=15)
    except (OSError, subprocess.TimeoutExpired):
        return None
    return result.stdout


def command_json(command):
    output = run(command)
    try:
        return json.loads(output) if output else None
    except json.JSONDecodeError:
        return None


def uptime_info():
    try:
        seconds = int(float(pathlib.Path("/proc/uptime").read_text(encoding="ascii").split()[0]))
    except (OSError, ValueError, IndexError):
        return {"seconds": None, "formatted": "Unavailable"}
    days, remainder = divmod(seconds, 86400)
    hours, remainder = divmod(remainder, 3600)
    minutes = remainder // 60
    return {"seconds": seconds, "formatted": f"{days}d {hours}h {minutes}m"}


def disk_kind(device):
    name = (device.get("kname") or device.get("name") or "").lower()
    transport = (device.get("tran") or "").lower()
    media = "HDD" if device.get("rota") is True else "SSD" if device.get("rota") is False else "storage"
    if name.startswith("nvme") or transport == "nvme":
        return f"NVMe {media}"
    if name.startswith("/dev/hd") and not transport:
        return f"PATA {media}"
    if transport == "sata":
        return f"SATA {media}"
    return f"{transport.upper()} {media}" if transport else media


def smart_check(path):
    if shutil.which("smartctl") is None:
        return {"status": "unavailable", "detail": "smartctl is not installed", "temperature": None}
    output = run(["smartctl", "--json", "--all", path])
    try:
        data = json.loads(output) if output else None
    except json.JSONDecodeError:
        data = None
    if data is None:
        return {"status": "unknown", "detail": "SMART data unavailable; device may require elevated privileges", "temperature": None}

    failures = []
    passed = data.get("smart_status", {}).get("passed")
    if passed is False:
        failures.append("SMART overall-health self-assessment failed")
    nvme = data.get("nvme_smart_health_information_log", {})
    warning = nvme.get("critical_warning")
    if isinstance(warning, int) and warning:
        failures.append(f"NVMe critical warning: {warning}")
    for attribute in data.get("ata_smart_attributes", {}).get("table", []):
        flags = attribute.get("flags", {})
        if flags.get("prefailure") and attribute.get("when_failed") not in (None, "", "-", "Never"):
            failures.append(f"SMART attribute {attribute.get('name', attribute.get('id', 'unknown'))} failed")
    temperature = data.get("temperature", {}).get("current", nvme.get("temperature"))
    try:
        temperature = float(temperature) if temperature is not None else None
    except (TypeError, ValueError):
        temperature = None
    detail = "; ".join(failures) if failures else "SMART health passed" if passed is True else "SMART data has no overall health result"
    return {"status": "failed" if failures else "passed" if passed is True else "unknown", "detail": detail, "temperature": temperature}


def collect_disks():
    data = command_json(["lsblk", "--json", "--bytes", "--paths", "--output", "NAME,KNAME,PATH,TYPE,ROTA,TRAN,MODEL"])
    if not data:
        return [], "Disk inventory unavailable."
    disks = []
    for device in data.get("blockdevices", []):
        if device.get("type") != "disk":
            continue
        name = (device.get("kname") or device.get("name") or "").lower()
        if name.startswith(("loop", "ram", "zram")):
            continue
        path = device.get("path") or device.get("name")
        if path:
            disks.append({"device": path, "model": (device.get("model") or "Unknown model").strip(), "kind": disk_kind(device), "smart": smart_check(path)})
    return disks, None


def sysfs_text(path):
    try:
        return path.read_text(encoding="utf-8").strip()
    except (OSError, UnicodeError):
        return None


def numeric_value(value):
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def collect_batteries():
    batteries = []
    for device in sorted(pathlib.Path("/sys/class/power_supply").glob("BAT*")):
        if sysfs_text(device / "present") == "0":
            continue
        energy_now = numeric_value(sysfs_text(device / "energy_now"))
        energy_full = numeric_value(sysfs_text(device / "energy_full"))
        energy_design = numeric_value(sysfs_text(device / "energy_full_design"))
        charge_now = numeric_value(sysfs_text(device / "charge_now"))
        charge_full = numeric_value(sysfs_text(device / "charge_full"))
        charge_design = numeric_value(sysfs_text(device / "charge_full_design"))

        if energy_full is not None or energy_design is not None or energy_now is not None:
            full_capacity, design_capacity, current_capacity, unit = energy_full, energy_design, energy_now, "uWh"
        else:
            full_capacity, design_capacity, current_capacity, unit = charge_full, charge_design, charge_now, "uAh"

        charge_percent = numeric_value(sysfs_text(device / "capacity"))
        if charge_percent is None and current_capacity is not None and full_capacity:
            charge_percent = current_capacity * 100 / full_capacity
        health_percent = full_capacity * 100 / design_capacity if full_capacity is not None and design_capacity else None
        if health_percent is None:
            reported_health = sysfs_text(device / "health")
            lowered_health = (reported_health or "").lower()
            health_status = "critical" if any(word in lowered_health for word in ("poor", "replace", "critical")) else "warning" if any(word in lowered_health for word in ("fair", "low")) else "unknown"
        else:
            health_status = "critical" if health_percent < 60 else "warning" if health_percent < 80 else "good"

        batteries.append({
            "name": device.name,
            "manufacturer": sysfs_text(device / "manufacturer"),
            "model": sysfs_text(device / "model_name"),
            "status": sysfs_text(device / "status") or "Unknown",
            "charge_percent": round(charge_percent, 1) if charge_percent is not None else None,
            "cycle_count": numeric_value(sysfs_text(device / "cycle_count")),
            "full_capacity": full_capacity,
            "design_capacity": design_capacity,
            "current_capacity": current_capacity,
            "capacity_unit": unit if full_capacity is not None or design_capacity is not None or current_capacity is not None else None,
            "health_percent": round(health_percent, 1) if health_percent is not None else None,
            "reported_health": sysfs_text(device / "health"),
            "health_status": health_status,
        })
    return batteries


def fallback_limits(identity):
    identity = identity.lower()
    if any(word in identity for word in ("nvme", "drive", "disk", "composite")):
        return 60.0, 70.0
    if any(word in identity for word in ("gpu", "amdgpu", "nvidia", "junction", "hotspot")):
        return 85.0, 95.0
    if any(word in identity for word in ("cpu", "k10temp", "core", "package", "tctl", "tdie")):
        return 85.0, 95.0
    return 80.0, 90.0


def valid_limit(value):
    try:
        value = float(value)
    except (TypeError, ValueError):
        return None
    return value if 0 < value < 150 else None


def collect_temperatures():
    data = command_json(["sensors", "-j"])
    if not data:
        return [], "Temperature sensors unavailable (install/configure lm-sensors)."
    readings = []
    for chip, entries in data.items():
        if not isinstance(entries, dict):
            continue
        for label, values in entries.items():
            if not isinstance(values, dict):
                continue
            for key, raw in values.items():
                if not key.startswith("temp") or not key.endswith("_input"):
                    continue
                try:
                    current = float(raw)
                except (TypeError, ValueError):
                    continue
                if not -20 <= current <= 150:
                    continue
                prefix = key[:-6]
                high = valid_limit(values.get(prefix + "_max")) or valid_limit(values.get(prefix + "_high"))
                critical = valid_limit(values.get(prefix + "_crit")) or valid_limit(values.get(prefix + "_emergency"))
                fallback_warning, fallback_critical = fallback_limits(chip + " " + label)
                warning = high or (critical - 5 if critical is not None else fallback_warning)
                critical = critical or fallback_critical
                if critical <= warning:
                    critical = warning + 5
                readings.append({"component": f"{chip} / {label}", "celsius": round(current, 1), "warning_limit": round(warning, 1), "critical_limit": round(critical, 1), "status": "critical" if current >= critical else "warning" if current >= warning else "ok"})
    return readings, None


def collect_report():
    disks, disk_note = collect_disks()
    temperatures, sensor_note = collect_temperatures()
    batteries = collect_batteries()
    for disk in disks:
        value = disk["smart"]["temperature"]
        if value is not None:
            warning, critical = fallback_limits(disk["kind"])
            temperatures.append({"component": disk["device"] + " SMART", "celsius": round(value, 1), "warning_limit": warning, "critical_limit": critical, "status": "critical" if value >= critical else "warning" if value >= warning else "ok"})
    return {"generated_at": datetime.datetime.now().astimezone().isoformat(timespec="seconds"), "uptime": uptime_info(), "disks": disks, "temperatures": temperatures, "batteries": batteries, "collection_notes": [note for note in (disk_note, sensor_note) if note]}


def has_problems(report):
    return (
        any(disk["smart"]["status"] == "failed" for disk in report["disks"])
        or any(item["status"] in ("warning", "critical") for item in report["temperatures"])
        or any(battery["health_status"] in ("warning", "critical") for battery in report["batteries"])
    )


def render_text(report, color=False):
    red, reset = ("\033[31m", "\033[0m") if color else ("", "")
    lines = ["HARDWARE HEALTH REPORT", "======================", f"Generated: {report['generated_at']}", f"System uptime: {report['uptime']['formatted']}", "", "DISKS / SMART", "-------------"]
    if report["disks"]:
        for disk in report["disks"]:
            smart = disk["smart"]
            style = (red, reset) if smart["status"] == "failed" else ("", "")
            lines.append(f"  {style[0]}{disk['device']} | {disk['kind']} | {disk['model']} | {smart['status'].upper()}: {smart['detail']}{style[1]}")
    else:
        lines.append("  No physical disks found")
    lines.extend(["", "COMPONENT TEMPERATURES", "----------------------"])
    for item in report["temperatures"]:
        style = (red, reset) if item["status"] in ("warning", "critical") else ("", "")
        lines.append(f"  {style[0]}{item['component']}: {item['celsius']:.1f} C [{item['status'].upper()}; warning {item['warning_limit']:.1f} C, critical {item['critical_limit']:.1f} C]{style[1]}")
    if not report["temperatures"]:
        lines.append("  No temperature readings available")
    lines.extend(["", "BATTERIES", "---------"])
    if report["batteries"]:
        for battery in report["batteries"]:
            health = f"{battery['health_percent']:.1f}%" if battery["health_percent"] is not None else battery["reported_health"] or "Unavailable"
            charge = f"{battery['charge_percent']:.1f}%" if battery["charge_percent"] is not None else "Unavailable"
            cycles = str(int(battery["cycle_count"])) if battery["cycle_count"] is not None else "Unavailable"
            model = f" | {battery['manufacturer'] or ''} {battery['model'] or ''}".rstrip()
            style = (red, reset) if battery["health_status"] in ("warning", "critical") else ("", "")
            lines.append(f"  {battery['name']}{model} | Charge: {charge} | State: {battery['status']} | Cycles: {cycles}")
            lines.append(f"    {style[0]}Health: {health} ({battery['health_status'].upper()}){style[1]}")
            if battery["full_capacity"] is not None and battery["design_capacity"] is not None:
                lines.append(f"    Full/design capacity: {battery['full_capacity']:.0f}/{battery['design_capacity']:.0f} {battery['capacity_unit']}")
    else:
        lines.append("  No batteries detected")
    if report["collection_notes"]:
        lines.extend(["", "COLLECTION NOTES", *["  " + note for note in report["collection_notes"]]])
    lines.extend(["", "OVERALL: " + ("PROBLEMS DETECTED" if has_problems(report) else "No reported SMART or temperature problems")])
    return "\n".join(lines) + "\n"


def render_html(report):
    def esc(value):
        return html.escape(str(value))
    disk_rows = []
    for disk in report["disks"]:
        smart = disk["smart"]
        status = "problem" if smart["status"] == "failed" else "unknown" if smart["status"] != "passed" else "ok"
        disk_rows.append(f"<tr><td>{esc(disk['device'])}<small>{esc(disk['model'])}</small></td><td>{esc(disk['kind'])}</td><td class=\"{status}\">{esc(smart['status'].upper())}</td><td class=\"{status}\">{esc(smart['detail'])}</td></tr>")
    temp_rows = []
    for item in report["temperatures"]:
        status = "problem" if item["status"] in ("warning", "critical") else "ok"
        temp_rows.append(f"<tr><td>{esc(item['component'])}</td><td class=\"{status}\">{item['celsius']:.1f} C</td><td>{item['warning_limit']:.1f} C</td><td>{item['critical_limit']:.1f} C</td><td class=\"{status}\">{esc(item['status'].upper())}</td></tr>")
    battery_rows = []
    for battery in report["batteries"]:
        health_class = "problem" if battery["health_status"] in ("warning", "critical") else "unknown" if battery["health_status"] == "unknown" else "ok"
        charge = f"{battery['charge_percent']:.1f}%" if battery["charge_percent"] is not None else "Unavailable"
        health = f"{battery['health_percent']:.1f}%" if battery["health_percent"] is not None else battery["reported_health"] or "Unavailable"
        cycles = str(int(battery["cycle_count"])) if battery["cycle_count"] is not None else "Unavailable"
        model = " ".join(value for value in (battery["manufacturer"], battery["model"]) if value)
        battery_rows.append(f"<tr><td>{esc(battery['name'])}<small>{esc(model)}</small></td><td>{esc(charge)}</td><td>{esc(battery['status'])}</td><td>{esc(cycles)}</td><td class=\"{health_class}\">{esc(health)} ({esc(battery['health_status'].upper())})</td></tr>")
    notes = "".join(f"<li>{esc(note)}</li>" for note in report["collection_notes"])
    summary_class = "problem" if has_problems(report) else "ok"
    summary = "Problems detected" if has_problems(report) else "No reported SMART or temperature problems"
    return ("<!doctype html><html lang=\"en\"><head><meta charset=\"utf-8\"><meta name=\"viewport\" content=\"width=device-width, initial-scale=1\"><title>Hardware Health Report</title><style>"
        ":root{color-scheme:light;--ink:#17242b;--muted:#5e7078;--line:#d7e1e3;--paper:#f2f6f5;--panel:#fff;--accent:#087e80;--red:#b42318;--green:#247044;--amber:#8a5a00}*{box-sizing:border-box}body{margin:0;background:var(--paper);color:var(--ink);font:15px/1.5 system-ui,sans-serif}main{max-width:1100px;margin:auto;padding:30px 20px 48px}header{border-bottom:2px solid var(--accent);padding-bottom:16px}h1{margin:0;font-size:28px}header p{color:var(--muted);margin:6px 0 0}section{background:var(--panel);border:1px solid var(--line);padding:18px;margin-top:14px}h2{margin:0 0 14px;font-size:18px}.table-wrap{overflow-x:auto}table{width:100%;border-collapse:collapse;text-align:left}th,td{padding:10px;border-bottom:1px solid var(--line);vertical-align:top}th{color:var(--muted);font-size:12px;text-transform:uppercase}small{display:block;color:var(--muted)}.problem{color:var(--red);font-weight:700}.unknown{color:var(--amber);font-weight:650}.ok{color:var(--green);font-weight:650}.summary{font-size:18px}@media(max-width:640px){main{padding:20px 12px 32px}section{padding:13px}}</style></head><body><main>"
        f"<header><h1>Hardware Health Report</h1><p>Generated {esc(report['generated_at'])} · Uptime {esc(report['uptime']['formatted'])}</p></header><section><p class=\"summary {summary_class}\">{summary}</p></section>"
        "<section><h2>Disks / SMART</h2><div class=\"table-wrap\"><table><thead><tr><th>Device / model</th><th>Type</th><th>SMART</th><th>Details</th></tr></thead><tbody>"
        + ("".join(disk_rows) or "<tr><td colspan=\"4\">No physical disks found</td></tr>")
        + "</tbody></table></div></section><section><h2>Component temperatures</h2><div class=\"table-wrap\"><table><thead><tr><th>Component</th><th>Current</th><th>Warning at</th><th>Critical at</th><th>Status</th></tr></thead><tbody>"
        + ("".join(temp_rows) or "<tr><td colspan=\"5\">No temperature readings available</td></tr>")
        + "</tbody></table></div></section>"
        + "<section><h2>Batteries</h2><div class=\"table-wrap\"><table><thead><tr><th>Battery</th><th>Charge</th><th>State</th><th>Cycles</th><th>Health</th></tr></thead><tbody>"
        + ("".join(battery_rows) or "<tr><td colspan=\"5\">No batteries detected</td></tr>")
        + "</tbody></table></div></section>"
        + (f"<section><h2>Collection notes</h2><ul>{notes}</ul></section>" if notes else "")
        + "</main></body></html>\n")


def ask(prompt, default):
    try:
        answer = input(f"{prompt} [{default}]: ").strip()
    except EOFError:
        return default
    return answer or default


def invoking_user_account():
    sudo_user = os.environ.get("SUDO_USER")
    if sudo_user:
        try:
            return pwd.getpwnam(sudo_user)
        except KeyError:
            pass
    sudo_uid = os.environ.get("SUDO_UID")
    if sudo_uid and sudo_uid.isdigit():
        try:
            return pwd.getpwuid(int(sudo_uid))
        except KeyError:
            pass
    return pwd.getpwuid(os.getuid())


def invoking_user_home():
    return pathlib.Path(invoking_user_account().pw_dir)


def save_report(path_text, format_name, force, interactive, report):
    path = pathlib.Path(path_text).expanduser()
    if not path.parent.is_dir():
        raise OSError(f"output directory does not exist: {path.parent}")
    if path.exists() and not force:
        try:
            overwrite = interactive and input(f"{path} exists. Overwrite? [y/N]: ").strip().lower() in ("y", "yes")
        except EOFError:
            overwrite = False
        if not overwrite:
            print(f"Report not written: {path} already exists.", file=sys.stderr)
            return False
    if format_name == "html":
        content = render_html(report)
    elif format_name == "json":
        content = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
    else:
        content = render_text(report)
    temporary = None
    try:
        with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as output:
            temporary = pathlib.Path(output.name)
            output.write(content)
        sudo_uid = os.environ.get("SUDO_UID")
        if os.geteuid() == 0 and sudo_uid and sudo_uid.isdigit():
            account = invoking_user_account()
            os.chown(temporary, account.pw_uid, account.pw_gid)
        os.replace(temporary, path)
    finally:
        if temporary is not None and temporary.exists():
            temporary.unlink()
    print(f"Report saved: {path}")
    return True


def main():
    output_path, requested_format, force_value = sys.argv[1:]
    interactive = sys.stdin.isatty() and sys.stdout.isatty()
    report = collect_report()
    print(render_text(report, color=sys.stdout.isatty() and not os.environ.get("NO_COLOR")), end="")
    if output_path:
        suffix = pathlib.Path(output_path).suffix.lower().lstrip(".")
        format_name = requested_format or (suffix if suffix in ("txt", "html", "json") else "txt")
        saved = save_report(output_path, format_name, force_value == "1", interactive, report)
        return 1 if not saved and not interactive else 0
    if not interactive:
        return 0
    if input("Save this report to a file? [Y/n]: ").strip().lower() in ("n", "no"):
        return 0
    format_name = requested_format or ask("Report format (txt/html/json)", "txt").lower()
    if format_name not in ("txt", "html", "json"):
        print("Unsupported format; using txt.", file=sys.stderr)
        format_name = "txt"
    path = ask("Output path", str(invoking_user_home() / f"hardware-health.{format_name}"))
    try:
        save_report(path, format_name, force_value == "1", interactive, report)
    except OSError as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1
    return 0


try:
    sys.exit(main())
except OSError as error:
    print(f"Error: {error}", file=sys.stderr)
    sys.exit(1)
PYCODE
)" "$OUTPUT_PATH" "$OUTPUT_FORMAT" "$FORCE"