315
            

             package rbprompt

/*
FMT strings:

%u = username
%h = hostname
%i = ip address as colours
%e = random emote
%l = line number

%<n>p = last <n> directories in the current working directory
%p = current working directory, short for %0p

*/

import (
	"fmt"
	"math"
	"math/rand"
	"strings"
)

var (
	// EmoteEyes contains characters that can be used as eyes in emotes.
	// If this is empty, using %e in the prompt will cause an error.
	EmoteEyes = ":;="
	// EmoteMouths contains characters that can be used as mouths in emotes.
	// If this is empty, using %e in the prompt will cause an error.
	EmoteMouths = ")3>]DPO"
)

func randInt(max int) int {
	return rand.Int() % max
}

func emote() string {
	lenEyes := len(EmoteEyes)
	lenMouths := len(EmoteMouths)
	if lenEyes == 0 || lenMouths == 0 {
		return ""
	}
	return string(EmoteEyes[randInt(lenEyes)]) + string(EmoteMouths[randInt(lenMouths)])
}

type Rainbow struct {
	R, G, B, S int
}

func (r *Rainbow) init() {
	r.R = 255
	r.G = 0
	r.B = 0
	r.S = 55 - randInt(20)
}

func (r *Rainbow) copy() Rainbow {
	return Rainbow{r.R, r.G, r.B, r.S}
}

func (r *Rainbow) next() {
	if r.S <= 0 || (r.R == r.G && r.G == r.B) {
		r.init()
		return
	}
	for v := []*int{&r.R, &r.G, &r.B}; ; v[0], v[1], v[2] = v[1], v[2], v[0] {
		if *v[0] == 255 && *v[2] == 000 {
			*v[1] += r.S
			if *v[1] > 255 {
				*v[0] += 255 - *v[1]
				*v[1] = 255
			}
			break
		} else if *v[1] == 255 && *v[2] == 000 {
			*v[0] -= r.S
			if *v[0] < 0 {
				*v[2] += 0 - *v[0]
				*v[0] = 0
			}
			break
		}
	}
}

func (r *Rainbow) escape() string {
	return fmt.Sprintf(ESCAPESTR+"38;2;%d;%d;%dm", r.R, r.G, r.B)
}

func (r *Rainbow) colour(in string) (out string) {
	for i := range in {
		out += r.escape() + in[i:i+1]
		r.next()
	}
	out += "\033[0m"
	return
}

// PS is the main struct for the prompt. It contains all the information needed to generate the prompt.
// After setting User, HostName, Home, PWD, and Fmt, call Next() to generate the prompt.
// Then call String() to get the prompt string and print it.
type PS struct {
	// LineNo is the line number, defaults to 0, then increments by 1 each time Next() is called.
	LineNo uint
	// if NoColor is true, the prompt will be printed without any colors.
	NoColor bool
	// Current is the current prompt string, calculated by the last call to Next().
	// Same as String().
	Current string
	// Fmt is the format string for the prompt. It can contain the following:
	// %u = username
	// %h = hostname
	// %i = ip address as colours
	// %e = random emote
	// %l = line number
	// %<n>p = bottom <n> directories in PWD
	// %p = entire PWD, short for %0p
	// adding #xxxxxx (hex colour) after any argument will set the colour of that argument. e.g. %u#ff0000 will set the username to red.
	Fmt string
	// Set User to the username of the current user, required if Fmt contains "%u".
	User string
	// Set Home to the home directory of the current user, required if Fmt contains "%p" or %<n>p.
	Home string
	// Set HostName to the hostname of the current machine, required if Fmt contains "%h".
	HostName string
	// Set PWD to the current working directory, required if Fmt contains "%p" or %<n>p.
	PWD string
	// Set HostIP to the IP address of the current machine, required if Fmt contains "%i".
	HostIP [4]byte
	// Rainbow is the rainbow struct used to generate the rainbow effect.
	Rainbow Rainbow
}

// String returns the prompt calculated by the last call to Next.
func (p *PS) String() string {
	return p.Current
}

var ESCAPESTR = "\033["

func (p *PS) ip2col() string {
	if p.NoColor {
		return ""
	}
	n := func() func() int {
		i := 0
		return func() int {
			i++
			return int(p.HostIP[i-1])*216/256 + 16
		}
	}()
	return fmt.Sprintf(ESCAPESTR+"48;5;%d;38;5;%dm▄"+ESCAPESTR+"48;5;%d;38;5;%dm▄"+ESCAPESTR+"0m", n(), n(), n(), n())
}

func (p *PS) pwd(n int) string {
	pwd := p.PWD
	trimmed := false
	if strings.HasPrefix(pwd, p.Home) {
		trimmed = true
		pwd = strings.TrimPrefix(pwd, p.Home)
		if len(pwd) == 0 {
			pwd = "~"
		} else {
			if pwd[0] != '/' {
				pwd = "/" + pwd
			}
			pwd = "~" + pwd
		}
	}
	if n <= 0 {
		return pwd
	}
	a := strings.Split(pwd, "/")
	if !trimmed {
		a = append([]string{"/"}, a...)
	}
	ll := len(a) - 1
	s := ""
	for j := 0; j < n && j <= ll; j++ {
		s = a[ll-j] + " " + s
	}
	return strings.TrimSpace(s)
}

func errGen(fnName string) func(s string) error {
	return func(s string) error {
		return fmt.Errorf("rbprompt: %s: %s", fnName, s)
	}
}

type hex struct {
	rgb [3]byte
}

func (h *hex) readByte(i int, c byte) bool {
	if c >= '0' && c <= '9' {
		c = c - '0'
	} else if c >= 'a' && c <= 'f' {
		c = c - 'a' + 10
	} else {
		return false
	}
	h.rgb[i/2] |= byte(c) << (4 * (1 - i%2))
	return true
}
func (h *hex) escape(s string) string {
	return fmt.Sprintf(ESCAPESTR+"38;2;%d;%d;%dm%s"+ESCAPESTR+"0m", h.rgb[0], h.rgb[1], h.rgb[2], s)
}

func (p *PS) doFmt(rbfn *func(string) string, dst *string, f *string, i *int) error {
	errFn := errGen("PS.doFmt")
	s := ""
	var j int = *i
	*i++
	if *i >= len(*f) {
		return errFn("unexpected end of format string")
	}
	switch (*f)[*i] {
	case 'u':
		s = p.User
		break
	case 'h':
		s = p.HostName
		break
	case 'i':
		*dst += p.ip2col()
		*i++
		return nil
	case 'e':
		if ss := emote(); len(ss) == 2 {
			s = ss
		} else {
			return errFn("cannot use %e in format string if emote list is empty")
		}
		break
	case 'p':
		s = p.pwd(0)
		break
	case 'l':
		s = fmt.Sprint(p.LineNo)
		break
	case '%':
		s = "%"
		break
	default:
		for l := 0; *i < len(*f); *i++ {
			if (*f)[*i] >= '0' && (*f)[*i] <= '9' {
				l = l*10 + int((*f)[*i]-'0')
				continue
			} else if (*f)[*i] == 'p' {
				s = p.pwd(l)
			}
			break
		}
		break
	}
	if s == "" {
		return errFn("invalid format " + (*f)[j:*i])
	}
	if *i < len(*f)-1 && (*f)[*i+1] == '#' {
		*i++
		var rgb hex
		k := 0
		for ; *i < len(*f)-1 && k < 6; k++ {
			*i++
			if !rgb.readByte(k, (*f)[*i]) {
				goto parseError
			}
		}
		if k < 6 {
			goto parseError
		}
		*dst += rgb.escape(s)
	} else {
		*dst += (*rbfn)(s)
	}
	*i++
	return nil
parseError:
	*i++
	return errFn("parse error: invalid hex colour " + (*f)[j:*i])
}

// Next generates the next prompt string based on the format string and the current state of the PS struct.
func (p *PS) Next() error {
	p.Rainbow.next()
	r := p.Rainbow.copy()
	r.S += int(10 * (math.Sin(float64(p.LineNo) * math.Pi / 10)))

	p.Current = ""
	f := p.Fmt
	var i int
	var err error
	rbfn := r.colour
	if p.NoColor {
		rbfn = func(in string) string {
			return in
		}
	}
	for {
		i = strings.Index(f, "%")
		if i < 0 {
			p.Current += rbfn(f)
			break
		}
		if i > 0 {
			p.Current += rbfn(f[:i])
		}
		if err = p.doFmt(&rbfn, &p.Current, &f, &i); err != nil {
			return err
		}
		f = f[i:]
	}
	p.LineNo++
	return nil
}