package main

import (
	"bufio"
	"context"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"github.com/sashabaranov/go-openai"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Category struct {
	Name         string
	Transactions []Transaction
	Total        float64
}

func main() {
	// read csv line by line
	//lines, err := readLinesFromFile("./Transactions-2023-05-03.csv")
	//if err != nil {
	//	panic(err)
	//}

	//transactions, err := getJSONFromFile("./completed.json")
	//if err != nil {
	//	panic(err)
	//}

	transactions, err := getCategories()
	if err != nil {
		panic(err)
	}

	fmt.Println(len(transactions))
	fmt.Println(transactions[0])

	categories := make(map[string]Category)

	for _, transaction := range transactions {
		category, ok := categories[transaction.Category]
		if !ok {
			category = Category{
				Name:         transaction.Category,
				Transactions: []Transaction{},
				Total:        0,
			}
		}

		category.Name = transaction.Category
		category.Transactions = append(category.Transactions, transaction)

		// parse amount from string to float
		amount, err := strconv.ParseFloat(transaction.Amount, 64)
		if err != nil {
			continue
		}
		category.Total += amount
		categories[transaction.Category] = category
	}

	count := 0
	for _, category := range categories {
		fmt.Printf("%s: Count: %v Total: %.2f\n", category.Name, len(category.Transactions), category.Total)
		count += 1
	}
	fmt.Println(count)
	fmt.Println(len(transactions))
}

func getJSONFromFile(filePath string) ([]Transaction, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return []Transaction{}, fmt.Errorf("error opening file: %w", err)
	}
	defer file.Close()

	var transactions []Transaction
	err = json.NewDecoder(file).Decode(&transactions)
	if err != nil {
		return []Transaction{}, fmt.Errorf("error decoding json: %w", err)
	}
	return transactions, nil
}

func getCategories() ([]Transaction, error) {
	transactions, err := readCSV("./Transactions-2023-05-03.csv")
	if err != nil {
		return []Transaction{}, fmt.Errorf("error reading csv: %w", err)
	}

	waitGroup := sync.WaitGroup{}
	transactionChan := make(chan Transaction)

	var finalTransactions []Transaction
	count := 0
	fmt.Printf("Transaction Count: %v\n", len(transactions))
	for _, transaction := range transactions {
		if count != 0 && count%100 == 0 {
			fmt.Printf("At %v waiting...\n", count)
			time.Sleep(60 * time.Second)
			fmt.Println("starting... next 100")
		}
		waitGroup.Add(1)
		go func(line Transaction, wg *sync.WaitGroup, transactionChan chan Transaction, count int) {
			defer wg.Done()
			fmt.Println(line)
			transaction, err := parseTransaction(line.TransactionName)
			if err != nil {
				// continue if parsing error
				fmt.Println(err)
				fmt.Println("done with error")
				return
			}
			fmt.Println("done parsing")
			//fmt.Printf("%+v\n", transaction)
			combinedTransaction := Transaction{
				Category:        transaction.Category,
				SubCategory:     transaction.SubCategory,
				Amount:          line.Amount,
				TransactionName: line.TransactionName,
				IsoDate:         line.IsoDate,
			}
			transactionChan <- combinedTransaction
			fmt.Println("done end of function")
		}(transaction, &waitGroup, transactionChan, count)
		count++
	}
	// Start a goroutine to receive values from the transactionChan
	go func() {
		for combinedTransaction := range transactionChan {
			finalTransactions = append(finalTransactions, combinedTransaction)
		}
	}()
	fmt.Println(count)
	fmt.Println("waiting...")
	waitGroup.Wait()
	close(transactionChan)

	jsonTransactions, err := json.Marshal(finalTransactions)
	if err != nil {
		return []Transaction{}, fmt.Errorf("error marshalling json: %w", err)
	}
	fmt.Println(string(jsonTransactions))
	err = writeBytesToFile("./raw-categories.json", jsonTransactions)
	if err != nil {
		return []Transaction{}, fmt.Errorf("error writing to file: %w", err)
	}
	return finalTransactions, nil
}

func writeBytesToFile(filePath string, bytes []byte) error {
	file, err := os.Create(filePath)
	if err != nil {
		return fmt.Errorf("error creating file: %w", err)
	}
	defer file.Close()

	_, err = file.Write(bytes)
	if err != nil {
		return fmt.Errorf("error writing to file: %w", err)
	}

	return nil
}

func parseTransaction(rawTransaction string) (Transaction, error) {
	// save for later
	// category should be one of the following: "Income", "Food", "Shopping", "Transportation", "Housing", "Utilities", "Insurance", "Medical", "Saving and Investing", "Debt Payments", "Personal Spending", "Entertainment", "Other"
	categories := []string{
		"Donations",
		"Banking/Finance",
		"Mortgage",
		"Credit Cards",
		"Loans",
		"Bills/Utilities",
		"Food/Groceries",
		"Pet Care",
		"Housing",
		"Cloud Services",
		"Health and Beauty",
		"Entertainment/Shopping",
		"Transportation",
		"Uncategorized",
	}

	client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
	messages := []openai.ChatCompletionMessage{
		{
			Role: openai.ChatMessageRoleSystem,
			Content: `
You are a helpful bot that categorizes transactions. You are given a transaction name raw from the bank and you need to categorize it. Please only respond in json format.
category should be one of the following: 
` + strings.Join(categories, ",\n") + `

subCategory should provide more detail on the category. For example, if the category is "Food/Groceries", the subCategory could be "Groceries" or "Restaurants".

Expected output format:
{
"category": "",
"subCategory": "",
"transactionName": ""
}`,
		},
		{
			Role:    openai.ChatMessageRoleUser,
			Content: rawTransaction,
		},
	}

	hasCategory := false
	for !hasCategory {
		resp, err := client.CreateChatCompletion(
			context.Background(),
			openai.ChatCompletionRequest{
				Model:       openai.GPT3Dot5Turbo,
				Messages:    messages,
				MaxTokens:   256,
				Temperature: 0,
			},
		)

		if err != nil {
			// if 429, wait and try again
			if strings.Contains(err.Error(), "429") {
				fmt.Println("429 error, waiting 5 seconds...")
				time.Sleep(5 * time.Second)
				continue
			}
			return Transaction{}, fmt.Errorf("ChatCompletion request error: %w", err)
		}
		fmt.Println(resp.Choices[0].Message.Content)
		transaction, err := parseAIOutput(resp.Choices[0].Message.Content)
		if err != nil {
			return Transaction{}, fmt.Errorf("ChatCompletion parsing error: %w, got response: %v", err, resp.Choices[0].Message.Content)
		}

		// check if category is in the categories slice
		if contains(categories, transaction.Category) {
			fmt.Println("Found category!")
			hasCategory = true
			return transaction, nil
		} else if len(messages) < 5 {
			fmt.Println("Did not find category, trying again...")
			messages = append(messages, openai.ChatCompletionMessage{
				Role:    resp.Choices[0].Message.Role,
				Content: resp.Choices[0].Message.Content,
				Name:    "",
			})
			messages = append(messages, openai.ChatCompletionMessage{
				Role:    openai.ChatMessageRoleSystem,
				Content: "Please provide a valid category.",
				Name:    "",
			})
		} else {
			break
		}
	}
	return Transaction{}, fmt.Errorf("could not find category for transaction: %v\nMessages: %+v", rawTransaction, messages)
}

func contains(categories []string, category string) bool {
	for _, c := range categories {
		if c == category {
			return true
		}
	}
	return false
}

type Transaction struct {
	Category        string `json:"category"`
	SubCategory     string `json:"subCategory"`
	Amount          string `json:"amount"`
	TransactionName string `json:"transactionName"`
	IsoDate         string `json:"isodate"`
}

func parseAIOutput(aiOutput string) (Transaction, error) {
	var transaction Transaction

	aiOutputParts := strings.Split(aiOutput, "{")
	aiJson := "{" + aiOutputParts[len(aiOutputParts)-1]

	err := json.Unmarshal([]byte(aiJson), &transaction)
	if err != nil {
		return Transaction{}, err
	}
	return transaction, nil
}

func readLinesFromFile(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return lines, nil
}

func readCSV(filePath string) ([]Transaction, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return []Transaction{}, fmt.Errorf("error opening file: %w", err)
	}
	defer file.Close()

	reader := csv.NewReader(file)
	records, err := reader.ReadAll()
	if err != nil {
		return []Transaction{}, fmt.Errorf("error reading csv: %w", err)
	}

	var transactions []Transaction
	// Do something with the records
	for _, record := range records {
		//fmt.Println(record)
		isoDate, err := time.Parse("01/02/2006", record[0])
		if err != nil {
			fmt.Println(fmt.Errorf("error parsing date: %w", err))
			continue
		}
		t := Transaction{
			IsoDate:         isoDate.Format("2006-01-02"),
			TransactionName: record[2],
			Amount:          record[7],
		}
		transactions = append(transactions, t)
	}
	return transactions, nil
}