path: root/backend/internal/database/stocks.go

package database

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"

	pb "github.com/ansg191/ibd-trader-backend/api/gen/idb/stock/v1"
	"github.com/ansg191/ibd-trader-backend/internal/analyzer"
	"github.com/ansg191/ibd-trader-backend/internal/utils"

	"github.com/Rhymond/go-money"
)

var ErrStockNotFound = errors.New("stock not found")

func GetStock(ctx context.Context, exec Executor, symbol string) (Stock, error) {
	row := exec.QueryRowContext(ctx, `
SELECT symbol, name, ibd_url
FROM stocks
WHERE symbol = $1;
`, symbol)

	var stock Stock
	if err := row.Scan(&stock.Symbol, &stock.Name, &stock.IBDUrl); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return Stock{}, ErrStockNotFound
		}
		return Stock{}, err
	}

	return stock, nil
}

func AddStock(ctx context.Context, exec Executor, stock Stock) error {
	_, err := exec.ExecContext(ctx, `
INSERT INTO stocks (symbol, name, ibd_url)
VALUES ($1, $2, $3)
ON CONFLICT (symbol)
    DO UPDATE SET name    = $2,
                  ibd_url = $3;`, stock.Symbol, stock.Name, stock.IBDUrl)
	return err
}

func AddRanking(ctx context.Context, exec Executor, symbol string, ibd50, cap20 int) error {
	if ibd50 > 0 {
		_, err := exec.ExecContext(ctx, `
INSERT INTO stock_rank (symbol, rank_type, rank)
VALUES ($1, $2, $3)`, symbol, "ibd50", ibd50)
		if err != nil {
			return err
		}
	}
	if cap20 > 0 {
		_, err := exec.ExecContext(ctx, `
INSERT INTO stock_rank (symbol, rank_type, rank)
VALUES ($1, $2, $3)`, symbol, "cap20", cap20)
		if err != nil {
			return err
		}
	}
	return nil
}

func AddStockInfo(ctx context.Context, exec TransactionExecutor, info *StockInfo) (string, error) {
	tx, err := exec.BeginTx(ctx, nil)
	if err != nil {
		return "", err
	}
	defer func(tx *sql.Tx) {
		_ = tx.Rollback()
	}(tx)

	// Add raw chart analysis
	row := tx.QueryRowContext(ctx, `
INSERT INTO chart_analysis (raw_analysis)
VALUES ($1)
RETURNING id;`, info.ChartAnalysis)

	var chartAnalysisID string
	if err = row.Scan(&chartAnalysisID); err != nil {
		return "", err
	}

	// Add stock info
	row = tx.QueryRowContext(ctx,
		`
INSERT INTO ratings (symbol, composite, eps, rel_str, group_rel_str, smr, acc_dis, chart_analysis, price)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
RETURNING id;`,
		info.Symbol,
		info.Ratings.Composite,
		info.Ratings.EPS,
		info.Ratings.RelStr,
		info.Ratings.GroupRelStr,
		info.Ratings.SMR,
		info.Ratings.AccDis,
		chartAnalysisID,
		info.Price.Display(),
	)

	var ratingsID string
	if err = row.Scan(&ratingsID); err != nil {
		return "", err
	}

	return ratingsID, tx.Commit()
}

func GetStockInfo(ctx context.Context, exec Executor, id string) (*StockInfo, error) {
	row := exec.QueryRowContext(ctx, `
SELECT r.symbol,
       s.name,
       ca.raw_analysis,
       r.composite,
       r.eps,
       r.rel_str,
       r.group_rel_str,
       r.smr,
       r.acc_dis,
       r.price
FROM ratings r
         INNER JOIN stocks s on r.symbol = s.symbol
         INNER JOIN chart_analysis ca on r.chart_analysis = ca.id
WHERE r.id = $1;`, id)

	var info StockInfo
	var priceStr string
	err := row.Scan(
		&info.Symbol,
		&info.Name,
		&info.ChartAnalysis,
		&info.Ratings.Composite,
		&info.Ratings.EPS,
		&info.Ratings.RelStr,
		&info.Ratings.GroupRelStr,
		&info.Ratings.SMR,
		&info.Ratings.AccDis,
		&priceStr,
	)
	if err != nil {
		return nil, err
	}

	info.Price, err = utils.ParseMoney(priceStr)
	if err != nil {
		return nil, err
	}

	return &info, nil
}

func AddAnalysis(
	ctx context.Context,
	exec Executor,
	ratingId string,
	analysis *analyzer.Analysis,
) (id string, err error) {
	err = exec.QueryRowContext(ctx, `
UPDATE chart_analysis ca
SET processed  = true,
    action     = $2,
    price      = $3,
    reason     = $4,
    confidence = $5
FROM ratings r
WHERE r.id = $1
  AND r.chart_analysis = ca.id
RETURNING ca.id;`,
		ratingId,
		analysis.Action,
		analysis.Price.Display(),
		analysis.Reason,
		analysis.Confidence,
	).Scan(&id)
	return id, err
}

type Stock struct {
	Symbol string
	Name   string
	IBDUrl string
}

type StockInfo struct {
	Symbol        string
	Name          string
	ChartAnalysis string
	Ratings       Ratings
	Price         *money.Money
}

type Ratings struct {
	Composite   uint8
	EPS         uint8
	RelStr      uint8
	GroupRelStr LetterRating
	SMR         LetterRating
	AccDis      LetterRating
}

type LetterRating pb.LetterGrade

func (r LetterRating) String() string {
	switch pb.LetterGrade(r) {
	case pb.LetterGrade_LETTER_GRADE_E:
		return "E"
	case pb.LetterGrade_LETTER_GRADE_E_PLUS:
		return "E+"
	case pb.LetterGrade_LETTER_GRADE_D_MINUS:
		return "D-"
	case pb.LetterGrade_LETTER_GRADE_D:
		return "D"
	case pb.LetterGrade_LETTER_GRADE_D_PLUS:
		return "D+"
	case pb.LetterGrade_LETTER_GRADE_C_MINUS:
		return "C-"
	case pb.LetterGrade_LETTER_GRADE_C:
		return "C"
	case pb.LetterGrade_LETTER_GRADE_C_PLUS:
		return "C+"
	case pb.LetterGrade_LETTER_GRADE_B_MINUS:
		return "B-"
	case pb.LetterGrade_LETTER_GRADE_B:
		return "B"
	case pb.LetterGrade_LETTER_GRADE_B_PLUS:
		return "B+"
	case pb.LetterGrade_LETTER_GRADE_A_MINUS:
		return "A-"
	case pb.LetterGrade_LETTER_GRADE_A:
		return "A"
	case pb.LetterGrade_LETTER_GRADE_A_PLUS:
		return "A+"
	default:
		return "NA"
	}
}

func LetterRatingFromString(str string) LetterRating {
	switch str {
	case "E":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_E)
	case "E+":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_E_PLUS)
	case "D-":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_D_MINUS)
	case "D":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_D)
	case "D+":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_D_PLUS)
	case "C-":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_C_MINUS)
	case "C":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_C)
	case "C+":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_C_PLUS)
	case "B-":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_B_MINUS)
	case "B":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_B)
	case "B+":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_B_PLUS)
	case "A-":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_A_MINUS)
	case "A":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_A)
	case "A+":
		return LetterRating(pb.LetterGrade_LETTER_GRADE_A_PLUS)
	case "NA":
		fallthrough
	default:
		return LetterRating(pb.LetterGrade_LETTER_GRADE_UNSPECIFIED)
	}
}

func (r LetterRating) Value() (driver.Value, error) {
	return r.String(), nil
}

func (r *LetterRating) Scan(src any) error {
	var source string
	switch v := src.(type) {
	case string:
		source = v
	case []byte:
		source = string(v)
	default:
		return errors.New("incompatible type for LetterRating")
	}
	*r = LetterRatingFromString(source)
	return nil
}