#!/usr/bin/perl -w
#
# Usage example:
#    PS1=`colorstring -ps Cyan` Hello, `colorstring -ps Green`"world ==>"
#
# Written <2006-10-04, Steven J. DeRose, sderose@acm.org.
# 2008-02-11 sjd: Add -perl, perl -w.
# 2008-09-03 sjd: BSD. Improve doc, error-checking, fix bug in -all.
#
use strict;
use Getopt::Long;

my $version="2008-09-03";

my $all = 0;
my $help = 0;
my $list = 0;
my $message = "";
my $perl = 0;
my $print = 0;
my $ps = 0;
my $quiet = 0;
my $verbose = 0;
my $warnmsg = "";

my %colorTable = ();
setupColors();

# Process options
Getopt::Long::Configure ("ignore_case");
my $result = GetOptions(
    "all"               => \$all,
    "h|help|?"          => \$help,
    "list"              => \$list,
    "m=s"               => \$message,
	"perl!"             => \$perl,
    "print"             => \$print,
    "ps"                => \$ps,
    "q|quiet!"          => \$quiet,
    "v|verbose+"        => \$verbose,
    "version"           => sub {
        die "Version of $version, by Steven J. DeRose, sderose\@acm.org.\n";
    },
    "w=s"               => \$warnmsg,
    );

if ($help) { showUsage(); exit; }
($result) || die "Bad options.\n";


################################################################################

# For '-list', show all the colors
if ($list) {
    foreach my $c (sort keys %colorTable) {
        print chr(27) . $colorTable{$c} . $c
            . chr(27) . $colorTable{"fg_default"}
            . chr(27) . $colorTable{"bg_default"} . ", ";
    }
    print "\n";
    exit;
}

(scalar @ARGV) || die "No color specified.\n";

# For '-all', copy stdin coloring each line.
# Support a list of colors to alternate among.
if ($all) {
    my $reset = $colorTable{"fg_default"};
    my $bg_reset = $colorTable{"bg_default"};
    my @clist = ();
    while ($ARGV[0]) {
        my $colorName = shift;
        if ($colorName eq "usa")  {
            push(@clist, $colorTable{"fg2_red"});
            push(@clist, $colorTable{"fg2_default"});
            push(@clist, $colorTable{"fg2_blue"});
        }
        elsif ($colorName eq "christmas")  {
            push(@clist, $colorTable{"fg2_red"});
            push(@clist, $colorTable{"fg2_green"});
        }
        elsif ($colorName eq "italy")  {
            push(@clist, $colorTable{"fg2_red"});
            push(@clist, $colorTable{"fg2_green"});
            push(@clist, $colorTable{"fg2_default"});
        }
        elsif ($colorName eq "rainbow")  {
            push(@clist, $colorTable{"fg2_red"});
            push(@clist, $colorTable{"fg_red"});
            push(@clist, $colorTable{"fg2_yellow"});
            push(@clist, $colorTable{"fg2_green"});
            push(@clist, $colorTable{"fg2_blue"});
            push(@clist, $colorTable{"fg2_magenta"});
            push(@clist, $colorTable{"fg_magenta"});
        }
        else {
            my $seq = $colorTable{$colorName};
            if (!$seq) { $seq = $colorTable{"fg2_$colorName"}; }
            if (!$seq) {
            	die "colorstring: Unknown color '$colorName'.\n";
            }
            push(@clist, $seq); 
        }
    }
    #warn "Color sequence: " . join(" ",@clist) . "\n";

    my $n = 0;
    while (my $l = <>) {
        chomp($l);
        print chr(27) . $clist[$n] . $l
            . chr(27) . $reset
            . chr(27) . $bg_reset . "\n";
        $n++;
        if ($n >= scalar @clist) { $n = 0; }
    }
    exit;
}


# Look up the color name they requested
my $colorName = lc(shift);
my $escString = $colorTable{$colorName};
if (!$escString) {
    $escString = $colorTable{"fg2_$colorName"};
}
($escString) || 
    die "colorstring: Unknown color key '$ARGV[0]'. Use -h for help.\n";


if ($ps) {
    $escString = "\\e" . $escString;
}
elsif ($perl) {
    $escString = "   \$c$colorName = chr(27) . \"" . $escString . "\";\n";
}
elsif ($print) {
    $escString = "   ESC " . $escString . "\n";
}
elsif ($message || $warnmsg) {
    my $s1 = chr(27) . $escString;
    my $s3 = chr(27) . $colorTable{"cancel"} . "\n";
    $escString = "$s1$message$s3\n";
    if ($warnmsg) {
        print STDERR "$s1$warnmsg$s3\n";
    }
    if (!$message) { exit; }
}
else {
    $escString = chr(27) . $escString;
}

print $escString;

exit;


################################################################################

# Create a hash of named colors, etc. Don't include the initial escape,
# because it has to be different in prompts than elsewhere.
sub setupColors {
   %colorTable = (
      # Foreground colors:
      "fg_black"    => "[0;30m",
      "fg_red"      => "[0;31m",
      "fg_green"    => "[0;32m",
      "fg_yellow"   => "[0;33m",
      "fg_blue"     => "[0;34m",
      "fg_magenta"  => "[0;35m",
      "fg_cyan"     => "[0;36m",
      "fg_lightgray"=> "[0;37m",
      "fg_lightgrey"=> "[0;37m", # be international
      "fg_default"  => "[0;39m",

      # Bright foreground colors:
      "fg2_darkgray"=> "[1;30m",
      "fg2_darkgrey"=> "[1;30m",
      "fg2_red"     => "[1;31m",
      "fg2_green"   => "[1;32m",
      "fg2_yellow"  => "[1;33m",
      "fg2_blue"    => "[1;34m",
      "fg2_magenta" => "[1;35m",
      "fg2_cyan"    => "[1;36m",
      "fg2_white"   => "[1;37m",
      "fg2_default" => "[0;39m",

      # Background colors:
      "bg_black"    => "[40m",
      "bg_red"      => "[41m",
      "bg_green"    => "[42m",
      "bg_yellow"   => "[43m",
      "bg_blue"     => "[44m",
      "bg_magenta"  => "[45m",
      "bg_cyan"     => "[46m",
      "bg_lightgray"=> "[47m",
      "bg_lightgrey"=> "[47m",
      "bg_default"  => "[49m",

      # 'off' ones are for xterm. Other terminals vary.
      "cancel"      => "[m",
      "bold"        => "[1m",
      "bold_off"    => "[22m",
      "inverse"     => "[7m",
      "inverse_off" => "[27m",
      "ul"          => "[4m",
      "ul_off"      => "[24m",
      "blink"       => "[5m",   # cons25/vt100 only.
      "blink_off"   => "[25m",
      "invisible"   => "[8m",   # not in some terminal emulators
   );
}


################################################################################

sub showUsage {
    print "
Usage:  colorstring [options] colorname
    Returns the terminal escape string for the given color. The main colors are:
    black red green yellow blue magenta cyan lightgray default. Each comes in 'fg_' for 
    regular foreground, 'fg2_' for bright foreground, and 'bg_' for background.
    Case is ignored for color names, and 'FG2_' may be omitted (e.g. 'red' works).
Options:
    -all      show all of stdin in the specified color
                  (with -all, multiple 'colorname's will alternate, or you
                  can say 'usa', 'christmas', 'italy', or 'rainbow').
    -list     show names and escapes for all known colors, separated by commas.
    -m 'msg'  send this message to terminal (stdout) in the specified color.
    -w 'msg'  send this message to stderr in the specified color.
    -perl     return Perl code to generate and assign the color string.
    -print    print out the color string requested.
    -ps       return color command with '\\e' to put in a Bash prompt string.
    -q        suppress most messages.
    -version  display version information and exit ($version, sjd).
Related commands: 
   Several *nix commands have a '--color=auto' option, including ls and grep.
";
}