#!/usr/bin/env tclsh
# ssd-info, a utility that reports the remaining write resource of an SSD.
# Copyright (c) 2015-2016, 2020-2021, 2024 D. Bohdan.
# License: MIT.

package require Tcl 8.6 9

namespace eval ssd-info {
    variable version 0.3.0
    variable message {%1$s (%3$s) has %2$s%% of its write resource remaining.}
}

# Get the value of the SMART vendor attribute $name for $device.
proc ::ssd-info::get-attribute {device name} {
    set lines [split [exec smartctl -A $device] \n]
    set line [lsearch -inline -glob $lines *$name*]
    set value [string trimleft [lindex $line 3] 0]

    if {![string is integer -strict $value]} {
        error [list $device has no attribute $name]
    }
    return $value
}

# Get SSD wearout value using one of the two common vendor attributes.
proc ::ssd-info::get-wearout {device} {
    # Samsung can be incorrectly detected as Intel.  In that case the wear
    # value may be wrong.
    foreach {model attr} {
        Generic
        Remaining_Lifetime_Perc

        SandForce
        SSD_Life_Left

        Intel
        Media_Wearout_Indicator

        Crucial/Micron
        Percent_Lifetime_Remain
    } {
        try {
            return [list [get-attribute $device $attr] $model]
        } on error _ {}
    }

    error {unknown SSD model}
}

# Linear interpolation.
proc ::ssd-info::interpolate-color {color1 color2 {x 0.5}} {
    set result {}
    foreach v1 $color1 v2 $color2 {
        lappend result [expr {
            round($v1 * (1 - $x) + $v2 * $x)
        }]
    }
    return $result
}

# Draw a progress bar-like gradient with text ${value}% over it.
proc ::ssd-info::draw-bar {canvas value color {steps 10} {font barFont}} {
    set width [$canvas cget -width]
    set height [$canvas cget -height]

    set barWidth [expr { $value * $width / 100.0 }]
    set stepSize [expr { $barWidth / $steps }]
    set color1 {}
    foreach x $color {
        lappend color1 [expr { $x / 2.0 }]
    }
    set color2 $color

    # Create a gradient out of rectangles.
    for {set step 0} {$step < $steps} {incr step} {
        set color [interpolate-color \
                $color1 \
                $color2 \
                [expr { (1.0 * $step) / ($steps - 1) }]]

            $canvas create rectangle \
                [expr { $step * $stepSize }] \
                0 \
                [expr { ($step + 1) * $stepSize }] \
                $height \
                -width 0 \
                -fill [format #%02x%02x%02x {*}$color]
    }

    $canvas create text \
            [expr { $barWidth / 2 }] \
            [expr { $height / 2 }] \
            -text [expr {round($value)}]% \
            -font $font \
            -fill white
}

# Display a GUI showing device wear.
proc ::ssd-info::gui {device wear model} {
    variable message

    wm title . ssd-info
    wm resizable . false false

    ::ttk::frame .frame

    canvas .frame.canvas -width 800 -height 50
    ::ttk::label .frame.status \
        -text [format $message $device $wear $model]

    set font [font actual .frame.canvas]
    dict set font \
        -size [expr { 2 * [dict get $font -size] }]
    font create barFont {*}$font

    ::ssd-info::draw-bar \
            .frame.canvas \
            $wear \
            [interpolate-color \
                {128 0 0} \
                {0 128 0} \
                [expr { $wear / 100.0 }] \
            ] \
            100

    bind . <Escape> { exit 0 }

    pack .frame.canvas
    pack .frame.status

    pack .frame
}

# Produce text output reporting device wear.
proc ::ssd-info::report {device wear model} {
    variable message

    puts [format $message $device $wear $model]
}

# From https://wiki.tcl-lang.org/page/main+script.
proc ::ssd-info::main-script? {} {
    global argv0
    if {[info exists argv0]
        && [file exists [info script]] && [file exists $argv0]} {
        file stat $argv0 argv0Info
        file stat [info script] scriptInfo

        expr {
            $argv0Info(dev) == $scriptInfo(dev)
            && $argv0Info(ino) == $scriptInfo(ino)
        }
    } else {
        return 0
    }
}

proc ::ssd-info::main {argv0 argv} {
    # Parse the command line.
    set helpFlags {/? -? -h -help --help}
    set noGUI --?no-?gui

    set help [expr { $argv in $helpFlags }]
    lassign $argv device flag

    if {[llength $argv] ni {1 2}
        || ![file exists $device]
        || ($flag ne {} && ![regexp -- $noGUI $flag])} {
        puts "usage: [file tail $argv0] device \[--no-gui\]"
        exit [expr { !$help }]
    }

    # Get the value.
    lassign [::ssd-info::get-wearout $device] wear model
    if {$wear < 0 || $wear > 100} {
        error [list wear level $wear is outside of range 0-100]
    }

    # Report the result.
    if {$flag eq {}} {
        set argv [list -- {*}$argv]
        package require Tk

        rename ::send {}

        gui $device $wear $model
    } else {
        report $device $wear $model
    }

}

if {[::ssd-info::main-script?]} {
    ::ssd-info::main $argv0 $argv
}

proc ::ssd-info::get-attribute {device name} {
    set name [regsub -all {[][*+?{}()<>|.^$\\]} $name {\\&}]
    if {![regsub [format {.*\n\s*\S+\s+%s\s+\S+\s+(\S+).*} $name]\
            [exec smartctl -A $device] {\1} result]} {
        error "$device has no attribute $name"
    }
    return $result
}

# Get the value of the SMART vendor attribute $name for $device.
proc ::ssd-info::get-attribute {device name} {
    set value [string trimleft \
            [exec smartctl -A $device | awk "/$name/ { print \$4 }"] 0]
    if {![string is integer -strict $value]} {
        error "$device has no attribute $name"
    }
    return $value
}