#!/bin/bash

if [ "$EUID" -ne 0 ]; then
    echo "Error: This script must be run as root (e.g., sudo $0)" >&2
    exit 1
fi

if ! command -v smartctl >/dev/null 2>&1; then
    echo "Error: smartctl is not installed." >&2
    exit 1
fi

OUTPUT="${1:-diskreport.html}"

escape_html() {
    sed 's/&/\&amp;/g; s/</\&lt;/g; s/>/\&gt;/g'
}

cat << EOF > "$OUTPUT"
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>DiskReport</title>
    <style>
        :root {
            --bg-color: #0f0f0f;
            --box-bg: #1a1a1a;
            --raw-bg: #0a0a0a;
            --text-color: #d3d7cf;
            --pass-color: #8ae234;
            --fail-color: #ef2929;
            --warn-color: #fce94f;
            --link-color: #729fcf;
            --border-color: #333333;
            --muted-color: #888a85;
        }
        
        * {
            scrollbar-width: thin;
            scrollbar-color: var(--border-color) var(--bg-color);
        }

        ::-webkit-scrollbar {
            width: 10px;
            height: 10px;
        }
        ::-webkit-scrollbar-track {
            background: var(--bg-color);
        }
        ::-webkit-scrollbar-thumb {
            background: var(--border-color);
            border-radius: 5px;
        }
        ::-webkit-scrollbar-thumb:hover {
            background: var(--muted-color);
        }

        html, body {
            min-height: 100vh;
            margin: 0;
            padding: 0;
            background-color: var(--bg-color);
            color: var(--text-color);
            font-family: ui-monospace, "Cascadia Code", "Source Code Pro", Menlo, Consolas, "Courier New", monospace;
            font-size: 16px;
            line-height: 1.5;
        }
        body {
            display: flex;
            justify-content: center;
            align-items: flex-start;
            padding: 2rem;
            box-sizing: border-box;
        }
        .container {
            margin: auto;
            width: 100%;
            max-width: 800px;
        }
        .header-box {
            border: 2px solid var(--pass-color);
            text-align: center;
            padding: 1.5rem;
            border-radius: 8px;
            margin-bottom: 2rem;
            background: var(--box-bg);
            box-shadow: 0 4px 6px rgba(0,0,0,0.3);
        }
        .header-title {
            color: var(--pass-color);
            margin: 0 0 0.5rem 0;
            font-size: 2rem;
            letter-spacing: 2px;
            text-transform: uppercase;
        }
        .header-date {
            color: var(--muted-color);
            font-size: 0.9rem;
        }
        .drive-box { 
            margin-bottom: 2rem; 
            background: var(--box-bg);
            padding: 1.5rem;
            border-radius: 8px;
            border: 1px solid var(--border-color);
            box-shadow: 0 4px 6px rgba(0,0,0,0.3);
        }
        .pass { color: var(--pass-color); font-weight: bold; }
        .fail { color: var(--fail-color); font-weight: bold; }
        .warn { color: var(--warn-color); font-weight: bold; }
        pre { 
            margin: 0; 
            white-space: pre-wrap; 
            word-wrap: break-word; 
        }
        summary { 
            cursor: pointer; 
            color: var(--link-color); 
            margin-top: 1.5rem; 
            outline: none;
            font-weight: bold;
            user-select: none;
            transition: color 0.2s ease;
        }
        summary:hover { color: #ffffff; }
        .raw-data { 
            color: var(--muted-color); 
            font-size: 14px; 
            margin-top: 1rem;
            padding: 1rem; 
            background: var(--raw-bg);
            border-left: 3px solid var(--border-color); 
            overflow-x: auto;
            white-space: pre;
        }
        
        .raw-data::-webkit-scrollbar-track {
            background: var(--raw-bg);
        }
    </style>
</head>
<body>
<div class="container">

<div class="header-box">
    <h1 class="header-title">DiskReport</h1>
    <div class="header-date">GENERATED: $(date)</div>
</div>
EOF

for disk in $(lsblk -nd -o NAME,TYPE | awk '$2=="disk" {print $1}'); do
    dev="/dev/$disk"
    
    if smartctl -i "$dev" >/dev/null 2>&1; then
        raw_output=$(smartctl -a "$dev")
        
        model=$(echo "$raw_output" | grep -iE "^Device Model:|^Product:|^Model Number:|^Vendor:" | head -n 1 | awk -F: '{print $2}' | xargs)
        if [ -z "$model" ]; then
            model=$(lsblk -n -d -o MODEL "$dev" | xargs)
        fi
        
        interface=$(lsblk -n -d -o TRAN "$dev" | tr 'a-z' 'A-Z' | xargs)
        if [ -z "$interface" ] || [ "$interface" = "UNKNOWN" ]; then
            if echo "$raw_output" | grep -qi "Transport protocol:"; then
                interface=$(echo "$raw_output" | grep -i "Transport protocol:" | awk -F: '{print $2}' | xargs | cut -d' ' -f1)
            elif echo "$raw_output" | grep -qi "Protocol:.*NVM Express"; then
                interface="NVMe"
            elif echo "$raw_output" | grep -qi "SATA Version is:"; then
                interface=$(echo "$raw_output" | grep -i "SATA Version is:" | awk -F: '{print $2}' | xargs | cut -d, -f1)
            elif echo "$raw_output" | grep -qi "SCSI"; then
                interface="SCSI"
            else
                interface="Unknown"
            fi
        fi

        health=$(echo "$raw_output" | grep -iE "test result:|Health Status:" | awk -F: '{print $2}' | xargs)
        if [[ "$health" == *"PASSED"* || "$health" == *"OK"* ]]; then
            health_fmt="<span class='pass'>OK</span>"
        else
            health_fmt="<span class='fail'>${health:-FAIL/UNKNOWN}</span>"
        fi

        health_pct=""
        pct_used=$(echo "$raw_output" | grep -i "^Percentage Used:" | awk '{print $3}' | tr -d '%')
        if [ -z "$pct_used" ]; then
            pct_used=$(echo "$raw_output" | grep -i "Percentage used endurance indicator:" | awk '{print $5}' | tr -d '%')
        fi
        
        if [ -n "$pct_used" ] && [[ "$pct_used" =~ ^[0-9]+$ ]]; then
            health_pct=$((100 - pct_used))
        else
            wearout=$(echo "$raw_output" | grep -iE "^[ ]*(231|233|202|177|169)[ ]" | head -n 1 | awk '{print $4}')
            if [ -n "$wearout" ] && [[ "$wearout" =~ ^[0-9]+$ ]]; then
                health_pct=$((10#$wearout))
            fi
        fi

        if [ -n "$health_pct" ]; then
            if [ "$health_pct" -ge 80 ]; then health_str="<span class='pass'>${health_pct}% (Excellent)</span>"
            elif [ "$health_pct" -ge 50 ]; then health_str="<span class='warn'>${health_pct}% (Good)</span>"
            elif [ "$health_pct" -ge 20 ]; then health_str="<span class='warn'>${health_pct}% (Fair)</span>"
            else health_str="<span class='fail'>${health_pct}% (Poor)</span>"
            fi
        else
            if [[ "$health" == *"PASSED"* || "$health" == *"OK"* ]]; then
                health_str="<span class='pass'>100% (Excellent)</span>"
            else
                health_str="<span class='fail'>0% (Failing)</span>"
            fi
        fi

        perf_pct=100
        perf_reason=""
        
        crc_err=$(echo "$raw_output" | grep -iE "^[ ]*199[ ]+UDMA_CRC_Error_Count" | awk '{print $10}')
        if [[ -n "$crc_err" && "$crc_err" -gt 0 ]]; then
            perf_pct=80
            perf_reason=" (Degraded - Cable/Interface Errors)"
        fi
        
        nvme_err=$(echo "$raw_output" | grep -i "^Media and Data Integrity Errors:" | awk '{print $6}')
        if [[ -n "$nvme_err" && "$nvme_err" -gt 0 ]]; then
            perf_pct=70
            perf_reason=" (Degraded - Integrity Errors)"
        fi
        
        scsi_defects=$(echo "$raw_output" | grep -i "Elements in grown defect list:" | awk -F: '{print $2}' | xargs)
        if [[ -n "$scsi_defects" && "$scsi_defects" -gt 0 ]]; then
            perf_pct=70
            perf_reason=" (Degraded - Grown Defects)"
        fi
        
        realloc=$(echo "$raw_output" | grep -iE "^[ ]*5[ ]+Reallocated_Sector_Ct" | awk '{print $10}')
        if [[ -n "$realloc" && "$realloc" -gt 0 ]]; then
            if [ "$perf_pct" -gt 80 ]; then perf_pct=80; fi
            perf_reason=" (Degraded - Reallocated Sectors)"
        fi

        if [ "$perf_pct" -eq 100 ]; then
            perf_str="<span class='pass'>100% (Excellent)</span>"
        elif [ "$perf_pct" -ge 80 ]; then
            perf_str="<span class='warn'>${perf_pct}%${perf_reason}</span>"
        else
            perf_str="<span class='fail'>${perf_pct}%${perf_reason}</span>"
        fi

        temp=$(echo "$raw_output" | grep -iE "^Temperature:|^Current Drive Temperature:" | head -n 1 | awk -F: '{print $2}' | xargs)

        pwr_raw=$(echo "$raw_output" | grep -iE "Power On Hours|Accumulated power on time|number of hours powered up")
        if echo "$pwr_raw" | grep -q "minutes"; then
            pwr=$(echo "$pwr_raw" | awk -F'hours:minutes ' '{print $2}' | cut -d: -f1)
            pwr="${pwr} Hours"
        elif echo "$pwr_raw" | grep -qi "number of hours powered up"; then
            pwr=$(echo "$pwr_raw" | awk -F= '{print $2}' | awk '{print $1}')
            pwr="${pwr} Hours"
        elif [ -n "$pwr_raw" ]; then
            pwr=$(echo "$pwr_raw" | awk -F: '{print $2}' | xargs | sed 's/,//g')
            pwr="${pwr} Hours"
        else
            pwr="N/A"
        fi

        written="N/A"
        if echo "$raw_output" | grep -q "Data Units Written:"; then
            written=$(echo "$raw_output" | grep "Data Units Written:" | awk -F'[][]' '{print $2}')
        elif echo "$raw_output" | grep -qi "gigabytes processed \[consume\]:"; then
            gb_written=$(echo "$raw_output" | grep -i "gigabytes processed \[consume\]:" | awk '{print $4}')
            written=$(awk -v gb="$gb_written" 'BEGIN { printf "%.2f TB", gb/1000 }')
        elif echo "$raw_output" | grep -qiE "^[ ]*241[ ]+Total_LBAs_Written"; then
            lbas=$(echo "$raw_output" | grep -iE "^[ ]*241[ ]+Total_LBAs_Written" | awk '{print $10}')
            written=$(awk -v lbas="$lbas" 'BEGIN { printf "%.2f TB", (lbas * 512) / 1000000000000 }')
        fi

        safe_dev=$(echo "$dev" | escape_html)
        safe_model=$(echo "${model:-Unknown}" | escape_html)
        safe_interface=$(echo "${interface:-Unknown}" | escape_html)
        safe_temp=$(echo "${temp:-N/A}" | escape_html)
        safe_pwr=$(echo "${pwr:-N/A}" | escape_html)
        safe_written=$(echo "${written:-N/A}" | escape_html)
        safe_raw=$(echo "$raw_output" | grep -v "Copyright" | grep -v "smartmontools.org" | escape_html)

        cat << EOF >> "$OUTPUT"
<div class='drive-box'>
<pre>
==========================================================
 DRIVE        : $safe_dev
 MODEL        : $safe_model
 INTERFACE    : $safe_interface
----------------------------------------------------------
 HEALTH       : $health_str
 PERFORMANCE  : $perf_str
 STATUS       : $health_fmt
 TEMP         : $safe_temp
 POWER ON     : $safe_pwr
 WRITTEN      : $safe_written
==========================================================
</pre>
<details>
  <summary>raw</summary>
  <pre class='raw-data'>$safe_raw</pre>
</details>
</div>
EOF
    fi
done

cat << 'EOF' >> "$OUTPUT"
</div>
</body>
</html>
EOF

echo "DiskReport generated: $(pwd)/$OUTPUT"
