#!/usr/bin/env bash

set -o pipefail

# tmux requires unrecognized OSC sequences to be wrapped with DCS tmux;
# <sequence> ST, and for all ESCs in <sequence> to be replaced with ESC ESC. It
# only accepts ESC backslash for ST. We use TERM instead of TMUX because TERM
# gets passed through ssh.
function print_osc() {
    if [[ $TERM == screen* || $TERM == tmux* ]]; then
        printf "\033Ptmux;\033\033]"
    else
        printf "\033]"
    fi
}

# More of the tmux workaround described above.
function print_st() {
    if [[ $TERM == screen* || $TERM == tmux* ]]; then
        printf "\a\033\\"
    else
        printf "\a"
    fi
}

function load_version() {
    if [ -z ${IMGCAT_BASE64_VERSION+x} ]; then
        IMGCAT_BASE64_VERSION=$(base64 --version 2>&1)
        export IMGCAT_BASE64_VERSION
    fi
}

function b64_encode() {
    load_version
    if [[ $IMGCAT_BASE64_VERSION =~ GNU ]]; then
        # Disable line wrap
        base64 -w0
    else
        base64
    fi
}

function b64_decode() {
    load_version
    if [[ $IMGCAT_BASE64_VERSION =~ fourmilab ]]; then
        BASE64ARG=-d
    elif [[ $IMGCAT_BASE64_VERSION =~ GNU ]]; then
        BASE64ARG=-di
    else
        BASE64ARG=-D
    fi
    base64 $BASE64ARG
}

# print_image filename inline base64contents print_filename width height preserve_aspect_ratio
#   filename: Filename to convey to client
#   inline: 0 or 1, if set to 1, the file will be displayed inline, otherwise, it will be downloaded
#   base64contents: Base64-encoded contents
#   print_filename: 0 or 1, if set to 1, print the filename after outputting the image
#   width: set output width of the image in character cells, pixels or percent
#   height: set output height of the image in character cells, pixels or percent
#   preserve_aspect_ratio: 0 or 1, if set to 1, fill the specified width and height as much as possible without stretching the image
#   file: Empty string or file type like "application/json" or ".js".
function print_image() {
    print_osc
    printf "1337;File=inline=%s" "$2"
    printf ";size=%d" $(printf "%s" "$3" | b64_decode | wc -c)
    [ -n "$1" ] && printf ";name=%s" "$(printf "%s" "$1" | b64_encode)"
    [ -n "$5" ] && printf ";width=%s" "$5"
    [ -n "$6" ] && printf ";height=%s" "$6"
    [ -n "$7" ] && printf ";preserveAspectRatio=%s" "$7"
    [ -n "$8" ] && printf ";type=%s" "$8"
    printf ":%s" "$3"
    print_st
    printf '\n'
    [ "$4" == "1" ] && echo "$1"
    has_image_displayed=t
}

function error() {
    errcho "ERROR: $*"
}

function errcho() {
    echo "$@" >&2
}

function show_help() {
    errcho
    errcho "Usage: imgcat [-p] [-n] [-W width] [-H height] [-r] [-s] [-u] [-t file-type] [-f] filename ..."
    errcho "       cat filename | imgcat [-W width] [-H height] [-r] [-s]"
    errcho
    errcho "Display images inline in the iTerm2 using Inline Images Protocol"
    errcho
    errcho "Options:"
    errcho
    errcho "    -h, --help                      Display help message"
    errcho "    -p, --print                     Enable printing of filename or URL after each image"
    errcho "    -n, --no-print                  Disable printing of filename or URL after each image"
    errcho "    -u, --url                       Interpret following filename arguments as remote URLs"
    errcho "    -f, --file                      Interpret following filename arguments as regular Files"
    errcho "    -t, --type file-type            Provides a type hint"
    errcho "    -r, --preserve-aspect-ratio     When scaling image preserve its original aspect ratio"
    errcho "    -s, --stretch                   Stretch image to specified width and height (this option is opposite to -r)"
    errcho "    -W, --width N                   Set image width to N character cells, pixels or percent (see below)"
    errcho "    -H, --height N                  Set image height to N character cells, pixels or percent (see below)"
    errcho
    errcho "    If you don't specify width or height an appropriate value will be chosen automatically."
    errcho "    The width and height are given as word 'auto' or number N followed by a unit:"
    errcho "        N      character cells"
    errcho "        Npx    pixels"
    errcho "        N%     percent of the session's width or height"
    errcho "        auto   the image's inherent size will be used to determine an appropriate dimension"
    errcho
    errcho "    If a type is provided, it is used as a hint to disambiguate."
    errcho "    The file type can be a mime type like text/markdown, a language name like Java, or a file extension like .c"
    errcho "    The file type can usually be inferred from the extension or its contents. -t is most useful when"
    errcho "    a filename is not available, such as whe input comes from a pipe."
    errcho
    errcho "Examples:"
    errcho
    errcho "    $ imgcat -W 250px -H 250px -s avatar.png"
    errcho "    $ cat graph.png | imgcat -W 100%"
    errcho "    $ imgcat -p -W 500px -u http://host.tld/path/to/image.jpg -W 80 -f image.png"
    errcho "    $ cat url_list.txt | xargs imgcat -p -W 40 -u"
    errcho "    $ imgcat -t application/json config.json"
    errcho
}

function check_dependency() {
    if ! (builtin command -V "$1" >/dev/null 2>&1); then
        error "missing dependency: can't find $1"
        exit 1
    fi
}

# verify that value is in the image sizing unit format: N / Npx / N% / auto
function validate_size_unit() {
    if [[ ! "$1" =~ ^(:?[0-9]+(:?px|%)?|auto)$ ]]; then
        error "Invalid image sizing unit - '$1'"
        show_help
        exit 1
    fi
}

## Main

if [ -t 0 ]; then
    has_stdin=f
else
    has_stdin=t
fi

# Show help if no arguments and no stdin.
if [ $has_stdin = f ] && [ $# -eq 0 ]; then
    show_help
    exit
fi

check_dependency base64
check_dependency wc
file_type=""

# Look for command line flags.
while [ $# -gt 0 ]; do
    case "$1" in
    -h | --h | --help)
        show_help
        exit
        ;;
    -p | --p | --print)
        print_filename=1
        ;;
    -n | --n | --no-print)
        print_filename=0
        ;;
    -W | --W | --width)
        validate_size_unit "$2"
        width="$2"
        shift
        ;;
    -H | --H | --height)
        validate_size_unit "$2"
        height="$2"
        shift
        ;;
    -r | --r | --preserve-aspect-ratio)
        preserve_aspect_ratio=1
        ;;
    -s | --s | --stretch)
        preserve_aspect_ratio=0
        ;;
    -f | --f | --file)
        has_stdin=f
        is_url=f
        ;;
    -u | --u | --url)
        check_dependency curl
        has_stdin=f
        is_url=t
        ;;
    -t | --t | --type)
         file_type="$2"
         shift
         ;;
    -*)
        error "Unknown option flag: $1"
        show_help
        exit 1
        ;;
    *)
        if [ "$is_url" == "t" ]; then
            encoded_image=$(curl -fs "$1" | b64_encode) || {
                error "Could not retrieve image from URL $1, error_code: $?"
                exit 2
            }
        elif [ -r "$1" ]; then
            encoded_image=$(cat "$1" | b64_encode)
        else
            error "imgcat: $1: No such file or directory"
            exit 2
        fi
        has_stdin=f
        print_image "$1" 1 "$encoded_image" "$print_filename" "$width" "$height" "$preserve_aspect_ratio" "$file_type"
        ;;
    esac
    shift
done

# Read and print stdin
if [ $has_stdin = t ]; then
    print_image "" 1 "$(cat | b64_encode)" 0 "$width" "$height" "$preserve_aspect_ratio" "$file_type"
fi

if [ "$has_image_displayed" != "t" ]; then
    error "No image provided. Check command line options."
    show_help
    exit 1
fi

exit 0