main.go

package main

import (
	"fmt"
	"os"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/ericm/stonks/api"
	"github.com/ericm/stonks/graph"
	"github.com/piquette/finance-go/datetime"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var (
	interval,
	save,
	remove,
	name *string
	year    *bool
	ytd     *bool
	week    *bool
	version *bool
	extra   *bool
	width   *int
	columns *int
	theme   *string
	days    *int

	configPath string
)

func setDefaults() {
	home, ok := os.LookupEnv("HOME")
	if !ok {
		panic("No HOME env var set")
	}
	configPath = fmt.Sprintf("%s/.config", home)
	viper.AddConfigPath(configPath)
	viper.SetConfigName("stonks.yml")
	viper.SetConfigType("yaml")

	viper.SetDefault("favourites", map[string]interface{}{})

	// Config defaults
	viper.SetDefault("config.standalone_height", 12)
	viper.SetDefault("config.favourites_height", 12)
	viper.SetDefault("config.default_theme", graph.LineTheme)

	viper.ReadInConfig()
}

func main() {
	setDefaults()
	rootCmd := &cobra.Command{
		Use:   "stonks",
		Short: "A stock visualizer",
		Long:  "Displays realtime stocks in graph format in a terminal",
		Run: func(cmd *cobra.Command, args []string) {
			if *version {
				fmt.Println(api.Version)
				return
			}

			if len(*remove) > 0 {
				saveCmd := strings.ToLower(*remove)
				favourites, ok := viper.Get("favourites").(map[string]interface{})
				if !ok {
					fmt.Println("Read config error")
					os.Exit(1)
				}
				delete(favourites, saveCmd)
				viper.Set("favourites", favourites)
				viper.WriteConfig()
				return
			}

			if len(*save) > 0 {
				saveCmd := strings.ToUpper(*save)
				if _, err := api.GetChart(saveCmd, datetime.FifteenMins, nil, nil, false); err != nil {
					fmt.Println(err.Error())
					os.Exit(1)
				}
				favourites, ok := viper.Get("favourites").(map[string]interface{})
				if !ok {
					fmt.Println("Read config error")
					os.Exit(1)
				}
				nameCmd := saveCmd
				if len(*name) > 0 {
					nameCmd = *name
				}
				favourites[strings.ReplaceAll(saveCmd, ".", "_")] = nameCmd
				viper.Set("favourites", favourites)
				if err := viper.WriteConfig(); err != nil {
					err = viper.WriteConfigAs(path.Join(configPath, "stonks.yml"))
					if err != nil {
						fmt.Println(err.Error())
						os.Exit(1)
					}
				}
				return
			}

			chartThemeType := graph.ChartThemeType(viper.GetInt("config.default_theme"))

			switch *theme {
			case "line":
				chartThemeType = graph.LineTheme
			case "dot":
				chartThemeType = graph.DotTheme
			case "icon":
				chartThemeType = graph.IconTheme
			default:
				if len(*theme) > 0 {
					fmt.Println("Unknown theme, must be \"line\", \"dot\" or \"icon\"")
					os.Exit(1)
				}
			}

			chartTheme := graph.NewChartTheme(chartThemeType)
			var chartsToPrint []string = []string{}

			if len(args) == 0 {
				intervalCmd, start, end := parseTimeRange()
				extraCmd := false
				if end == nil {
					extraCmd = *extra
				}
				// Favourites
				favourites, ok := viper.Get("favourites").(map[string]interface{})
				if !ok {
					fmt.Println("Read config error")
					os.Exit(1)
				}
				if len(favourites) == 0 {
					fmt.Println("No favourites added. You can add them in the format 'stonks -s AMD -n \"Advanced Micro Devices\"'")
				}

				keys := make([]string, 0, len(favourites))
				for k := range favourites {
					keys = append(keys, k)
				}
				sort.Strings(keys)
				for _, symbol := range keys {
					chart, err := api.GetChart(strings.ReplaceAll(strings.ToUpper(symbol), "_", "."), intervalCmd, start, end, extraCmd)
					if err != nil {
						fmt.Println(err.Error())
						continue
					}
					g, _ := graph.GenerateGraph(chart, *width, viper.GetInt("config.favourites_height"), chartTheme, time.Local)
					chartsToPrint = append(chartsToPrint, g)
				}
			}

			for _, symbol := range args {
				intervalCmd, start, end := parseTimeRange()
				extraCmd := false
				if end == nil {
					extraCmd = *extra
				}
				chart, err := api.GetChart(strings.ToUpper(symbol), intervalCmd, start, end, extraCmd)
				if err != nil {
					fmt.Println(err.Error())
					os.Exit(1)
				}
				g, _ := graph.GenerateGraph(chart, *width, viper.GetInt("config.standalone_height"), chartTheme, time.Local)
				chartsToPrint = append(chartsToPrint, g)
			}

			PrintCharts(&chartsToPrint, *columns)
		},
	}
	interval = rootCmd.PersistentFlags().StringP("interval", "i", "15m", "stonks -i X[m|h] (eg 15m, 5m, 1h, 1d)")
	year = rootCmd.PersistentFlags().BoolP("year", "y", false, "Display the last year (will set interval to 5d)")
	ytd = rootCmd.PersistentFlags().Bool("ytd", false, "Display the year to date (will set interval to 5d)")
	week = rootCmd.PersistentFlags().BoolP("week", "w", false, "Display the last week (will set interval to 1d)")
	days = rootCmd.PersistentFlags().IntP("days", "d", 0, "24 hour period of stocks from X of days ago.")
	theme = rootCmd.PersistentFlags().StringP("theme", "t", "", "Display theme for the chart (Options: \"line\", \"dot\", \"icon\")")
	save = rootCmd.PersistentFlags().StringP("save", "s", "", "Add an item to the default stonks command. (Eg: -s AMD -n \"Advanced Micro Devices\")")
	remove = rootCmd.PersistentFlags().StringP("remove", "r", "", "Remove an item from favourites")
	name = rootCmd.PersistentFlags().StringP("name", "n", "", "Optional name for a stonk save")
	version = rootCmd.PersistentFlags().BoolP("version", "v", false, "stonks version")
	extra = rootCmd.PersistentFlags().BoolP("extra", "e", false, "Include extra pre + post time. (Only works for day)")
	width = rootCmd.PersistentFlags().IntP("width", "W", 80, "Base width of the graph")
	columns = rootCmd.PersistentFlags().IntP("columns", "C", 1, "Number of columns")

	rootCmd.Execute()
}

func parseTimeRange() (datetime.Interval, *datetime.Datetime, *datetime.Datetime) {
	var (
		intervalCmd datetime.Interval
		start       *datetime.Datetime
		end         *datetime.Datetime
	)
	switch {
	case *year:
		intervalCmd = datetime.FiveDay
		rn := time.Now()
		e := rn.AddDate(-1, 0, 0)
		start = datetime.New(&e)
		end = datetime.New(&rn)
	case *ytd:
		intervalCmd = datetime.FiveDay
		rn := time.Now()
		e := rn.AddDate(0, -int(rn.Month()), -rn.Day())
		start = datetime.New(&e)
		end = datetime.New(&rn)
	case *week:
		intervalCmd = datetime.OneHour
		rn := time.Now()
		e := rn.AddDate(0, 0, -7)
		start = datetime.New(&e)
		end = datetime.New(&rn)
	case interval == nil:
		intervalCmd = datetime.FifteenMins
	default:
		intervalCmd = datetime.Interval(*interval)
	}

	if *days > 0 {
		s := time.Now().AddDate(0, 0, *days*-1)
		y, m, d := s.Date()
		s = time.Date(y, m, d, 0, 0, 0, 0, s.Location())

		start = datetime.New(&s)
		e := time.Date(y, m, d, 23, 0, 0, 0, s.Location())
		end = datetime.New(&e)
	}
	return intervalCmd, start, end
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func PrintCharts(chartsToPrint *[]string, columns int) {
	var out [][]string

	// Transform chart strings to 2D array with first axis being a chart
	// and second axis being a line
	for i := 0; i < len(*chartsToPrint); i++ {
		out = append(out, strings.Split((*chartsToPrint)[i], "\n"))
	}

	var builder strings.Builder
	var start = 0
	var stop = min(columns, len(*chartsToPrint))

	for start != stop {
		// Loop over lines of a chart string
		for j := 0; j < len(out[0]); j++ {
			// Loop over charts
			for i := start; i < stop; i++ {
				builder.WriteString(out[i][j])
				builder.WriteString("   ")
			}
			builder.WriteString("\n")
		}

		// Move start and stop indices forward. By using min function, we ensure that we do not exceed bounds
		start = min(start+columns, len(*chartsToPrint))
		stop = min(stop+columns, len(*chartsToPrint))
	}

	fmt.Print(builder.String())
}