main.go

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"time"
)

var (
	alphavantageKey = flag.String("alpha-vantage-key", "", "Alpha Vantage API key")
	baseCurrency    = flag.String("base-currency", "USD", "Base currency")
	stocks          = flag.String("stocks", "", "Stocks to obtain data for")
	currencies      = flag.String("currencies", "", "Currencies to obtain data for")
	rateLimit       = flag.Int("rate-limit", 1000, "Number of requests per minute permitted to the Alpha Vantage API")
)

type currencyExchangeRate struct {
	FromCurrency string    `json:"1. From_Currency Code"`
	ToCurrency   string    `json:"3. To_Currency Code"`
	ExchangeRate float64   `json:"5. Exchange Rate"`
	Date         time.Time `json:"6. Last Refreshed"`
}

type stock struct {
	Symbol   string    `json:"01. symbol"`
	Price    float64   `json:"05. price"`
	Date     time.Time `json:"07. latest trading day"`
	Currency string
}

type stockSearchResult struct {
	Symbol     string  `json:"1. symbol"`
	Currency   string  `json:"8. currency"`
	MatchScore float64 `json:"9. matchScore,string"`
}

type httpClient struct {
	http.Client
	RequestsPerMinute int
	RequestsMade      int
	LastInterval      time.Time
}

// Add a few seconds to the rate limit interval for extra leeway.
const rateLimitInterval = time.Duration(63) * time.Second

func (client *httpClient) Do(req *http.Request) (*http.Response, error) {
	durationSinceLastInterval := time.Now().Sub(client.LastInterval)
	if durationSinceLastInterval > rateLimitInterval {
		client.LastInterval = time.Now()
		client.RequestsMade = 0
	} else {
		if client.RequestsMade == client.RequestsPerMinute {
			time.Sleep(rateLimitInterval - durationSinceLastInterval)
			client.LastInterval = time.Now()
			client.RequestsMade = 0
		}
	}

	resp, err := client.Client.Do(req)
	if err != nil {
		return nil, err
	}
	client.RequestsMade++

	return resp, err
}

func (currency *currencyExchangeRate) UnmarshalJSON(b []byte) error {
	err := UnmarshalJsonWithCustomDate(b, currency, "2006-01-02 15:04:05")
	if err != nil {
		return err
	}

	return nil
}

func (stock *stock) UnmarshalJSON(b []byte) error {
	err := UnmarshalJsonWithCustomDate(b, stock, "2006-01-02")
	if err != nil {
		return err
	}

	return nil
}

// Unmarshal JSON with custom parsing to create `time.Time` fields.
func UnmarshalJsonWithCustomDate[T any](b []byte, struct_ptr *T, dateLayout string) error {
	var unmarshalledJson map[string]string
	err := json.Unmarshal(b, &unmarshalledJson)
	if err != nil {
		return err
	}

	struct_value := reflect.ValueOf(struct_ptr).Elem()
	struct_type := struct_value.Type()
	num_fields := struct_type.NumField()
	json_has_fields := false

	for i := 0; i < num_fields; i++ {
		field := struct_type.Field(i)
		field_name := field.Name
		json_field_name := field.Tag.Get("json")

		if field_value, ok := unmarshalledJson[json_field_name]; ok {
			json_has_fields = true
			st_field := struct_value.FieldByName(field_name)

			if st_field.Kind() == reflect.String {
				st_field.SetString(field_value)
			} else if st_field.Kind() == reflect.Float64 {
				fl, err := strconv.ParseFloat(field_value, 64)
				if err != nil {
					return err
				}

				st_field.SetFloat(fl)
			} else if st_field.Type() == reflect.TypeOf(time.Time{}) {
				date, err := time.Parse(dateLayout, field_value)
				if err != nil {
					return err
				}

				st_field.Set(reflect.ValueOf(date))
			}
		}
	}

	if !json_has_fields {
		return errors.New("JSON has no fields")
	}

	return nil
}

func main() {
	flag.Parse()

	fmt.Println("; Generated by https://github.com/bnjmnt4n/hledger-prices on", time.Now().Format("2006/01/02 15:04:05"))
	fmt.Println()

	stock_symbols := cleanStringSlice(strings.Split(*stocks, ","))
	currency_symbols := cleanStringSlice(strings.Split(*currencies, ","))

	client := &httpClient{Client: http.Client{}, RequestsPerMinute: *rateLimit}

	// Get stock prices.
	if len(stock_symbols) > 0 {
		fmt.Println("; Stocks")
		fmt.Println()

		results_chan := make(chan struct {
			stock
			error
		})
		go getStockPrices(client, stock_symbols, results_chan)

		for result := range results_chan {
			stock, error := result.stock, result.error
			if error != nil {
				fmt.Printf("P %s %s %s %f ; %s\n", time.Now().Format("2006/01/02"), stock.Symbol, *baseCurrency, 0.0, error)
			} else {
				// Add currency of current stock to list of currencies to be fetched.
				currency_symbols = append(currency_symbols, stock.Currency)

				fmt.Printf("P %s %s %s %f\n", stock.Date.Format("2006/01/02"), stock.Symbol, stock.Currency, stock.Price)
			}
		}
	}

	// Get currency exchange rates.
	currency_symbols = cleanStringSlice(currency_symbols)
	if len(currency_symbols) > 0 {
		if len(stock_symbols) > 0 {
			fmt.Println()
		}
		fmt.Println("; Currencies")
		fmt.Println()

		results_chan := make(chan struct {
			currencyExchangeRate
			error
		})
		go getCurrencies(client, currency_symbols, results_chan)

		for result := range results_chan {
			exchange_rate, error := result.currencyExchangeRate, result.error
			if error != nil {
				fmt.Printf("P %s %s %s %f ; %s\n", time.Now().Format("2006/01/02"), exchange_rate.FromCurrency, *baseCurrency, 0.0, error)
			} else {
				fmt.Printf("P %s %s %s %f\n", exchange_rate.Date.Format("2006/01/02"), exchange_rate.FromCurrency, exchange_rate.ToCurrency, exchange_rate.ExchangeRate)
			}
		}
	}

	if len(stock_symbols) == 0 && len(currency_symbols) == 0 {
		fmt.Println("; No stocks/currencies provided.")
		fmt.Println("; Run `hledger-prices -h` to view accepted options.")
	}
}

func getAlphavantageData(client *httpClient, url string, output any) error {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		bytes, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return fmt.Errorf("Status code %d: Failed to read body: %w", resp.StatusCode, err)
		}

		return fmt.Errorf("Status code %d: %s", resp.StatusCode, string(bytes))
	}

	var data map[string]any
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return err
	}

	// API returns error message in JSON field `Note` or `Error Message`.
	if error_note, ok := data["Note"]; ok {
		return fmt.Errorf("Error: %s", error_note)
	}
	if error_note, ok := data["Error Message"]; ok {
		return fmt.Errorf("Error: %s", error_note)
	}

	// JSON returned only contains a single item with an inner object.
	for _, data := range data {
		json_bytes, err := json.Marshal(data)
		if err != nil {
			return err
		}

		err = json.Unmarshal(json_bytes, output)
		if err != nil {
			return err
		}

		return nil
	}

	return errors.New("Response had no data")
}

func getCurrencies(client *httpClient, currencies []string, results_chan chan<- struct {
	currencyExchangeRate
	error
}) {
	for _, currency_sym := range currencies {
		url := fmt.Sprintf("https://www.alphavantage.co/query?function=CURRENCY_EXCHANGE_RATE&from_currency=%s&to_currency=%s&apikey=%s", currency_sym, *baseCurrency, *alphavantageKey)
		exchange_rate := new(currencyExchangeRate)
		err := getAlphavantageData(client, url, exchange_rate)

		if err != nil {
			exchange_rate.FromCurrency = currency_sym
			results_chan <- struct {
				currencyExchangeRate
				error
			}{*exchange_rate, fmt.Errorf("Failed to get currency %s: %w", currency_sym, err)}
		} else {
			results_chan <- struct {
				currencyExchangeRate
				error
			}{*exchange_rate, nil}
		}
	}

	close(results_chan)
}

func getStockPrices(client *httpClient, stock_symbols []string, results_chan chan<- struct {
	stock
	error
}) {
	for _, stock_sym := range stock_symbols {
		url := fmt.Sprintf("https://www.alphavantage.co/query?function=GLOBAL_QUOTE&symbol=%s&apikey=%s", stock_sym, *alphavantageKey)
		stock_data := new(stock)
		err := getAlphavantageData(client, url, stock_data)

		if stock_data.Symbol == "" {
			stock_data.Symbol = stock_sym
		}
		symbol, _, found := strings.Cut(stock_data.Symbol, ".")
		stock_data.Symbol = symbol
		if err != nil {
			results_chan <- struct {
				stock
				error
			}{*stock_data, fmt.Errorf("Failed to get stock price for %s: %w", stock_sym, err)}
		} else {
			// Default: assume that currency is USD for all stocks traded in the US market.
			stock_data.Currency = "USD"
			if found {
				// Otherwise, fetch stock currency if stock is traded in other exchanges.
				url := fmt.Sprintf("https://www.alphavantage.co/query?function=SYMBOL_SEARCH&keywords=%s&apikey=%s", stock_sym, *alphavantageKey)
				stock_search_results := new([]stockSearchResult)
				err := getAlphavantageData(client, url, stock_search_results)
				if err != nil {
					results_chan <- struct {
						stock
						error
					}{*stock_data, fmt.Errorf("Failed to get stock currency for %s: %w", stock_sym, err)}
					continue
				}

				for _, stock_search_result := range *stock_search_results {
					if stock_search_result.Symbol == stock_sym || stock_search_result.MatchScore == 1.0 {
						stock_data.Currency = stock_search_result.Currency
						break
					}
				}
			}

			results_chan <- struct {
				stock
				error
			}{*stock_data, nil}
		}
	}

	close(results_chan)
}

// Trims whitespace from each string in a slice, sorts the slice and removes duplicate strings.
func cleanStringSlice(strs []string) []string {
	result := make([]string, 0, len(strs))

	for i, str := range strs {
		strs[i] = strings.TrimSpace(str)
	}
	sort.Strings(strs)

	var prev_str string
	for _, str := range strs {
		if str != "" && str != prev_str {
			result = append(result, str)
		}

		prev_str = str
	}

	return result
}