package demo

import (
	"fmt"
	"github.com/banbox/banbot/btime"
	"github.com/banbox/banbot/config"
	"github.com/banbox/banbot/strat"
	"github.com/banbox/banbot/utils"
	"github.com/banbox/banexg"
	"github.com/banbox/banexg/errs"
	"github.com/banbox/banexg/log"
	ta "github.com/banbox/banta"
	"go.uber.org/zap"
	"path/filepath"
	"strings"
)

type AggMore struct {
	bars    []*banexg.Kline
	barsBig []*banexg.Kline
}

func getBar(e *ta.BarEnv) *banexg.Kline {
	return &banexg.Kline{
		Time:   e.TimeStart,
		Open:   e.Open.Get(0),
		High:   e.High.Get(0),
		Low:    e.Low.Get(0),
		Close:  e.Close.Get(0),
		Volume: e.Volume.Get(0),
	}
}

func dumpKline(bars []*banexg.Kline, name string) *errs.Error {
	if len(bars) == 0 {
		return nil
	}
	outDir := config.GetDataDir()
	rows := utils.KlineToStr(bars, btime.UTCLocale)
	path := filepath.Join(outDir, "tmp", fmt.Sprintf("%s.csv", name))
	return utils.WriteCsvFile(path, rows, false)
}

func BarAgg(pol *config.RunPolicyConfig) *strat.TradeStrat {
	return &strat.TradeStrat{
		WarmupNum: 1,
		OnPairInfos: func(s *strat.StratJob) []*strat.PairSub {
			return []*strat.PairSub{
				{"_cur_", "1h", 1},
			}
		},
		OnStartUp: func(s *strat.StratJob) {
			s.More = &AggMore{}
		},
		OnBar: func(s *strat.StratJob) {
			b := getBar(s.Env)
			m, _ := s.More.(*AggMore)
			m.bars = append(m.bars, b)
		},
		OnInfoBar: func(s *strat.StratJob, e *ta.BarEnv, pair, tf string) {
			b := getBar(e)
			m, _ := s.More.(*AggMore)
			m.barsBig = append(m.barsBig, b)
		},
		OnShutDown: func(s *strat.StratJob) {
			m, _ := s.More.(*AggMore)
			code := strings.ReplaceAll(s.Symbol.Symbol, "/", "_")
			code = strings.ReplaceAll(code, ":", "_")
			err := dumpKline(m.bars, fmt.Sprintf("%s_%s", code, s.TimeFrame))
			if err != nil {
				log.Error("dump klines fail", zap.Error(err))
			}
			err = dumpKline(m.barsBig, fmt.Sprintf("%s_1h", code))
			if err != nil {
				log.Error("dump big klines fail", zap.Error(err))
			}
		},
	}
}