package main

import (
	"git.beejay.kim/Gshopper/sentio"
	"git.beejay.kim/Gshopper/sentio/indicator"
	"git.beejay.kim/Gshopper/sentio/util"
)

const (
	ATR_MULTIPLIER = 3
	ATR_PERIOD     = 21
	ATR_HHV        = 10
)

var Strategy = qqq{}

type qqq struct{}

func (strategy qqq) Name() string {
	return "Alpaca: QQQ [TQQQ : SQQQ]"
}

func (strategy qqq) Model() string {
	return "qqq200"
}

func (strategy qqq) MarketId() string {
	return "alpaca"
}

func (strategy qqq) Interval() uint8 {
	return 1
}

func (strategy qqq) PositionSymbols() map[sentio.Side]string {
	return map[sentio.Side]string{
		sentio.BASE:  "QQQ",
		sentio.LONG:  "TQQQ",
		sentio.SHORT: "SQQQ",
	}
}

func (strategy qqq) PositionProbabilities() map[sentio.Side]float32 {
	return map[sentio.Side]float32{
		sentio.BASE:  -1,
		sentio.LONG:  1.000097,
		sentio.SHORT: .999867,
	}
}

func (strategy qqq) Handle(market sentio.Market, probability *sentio.Probability) error {
	var (
		now        = market.Clock().Now()
		stamp      *sentio.Probability_Stamp
		symbols    = util.Symbols(strategy)
		stoplosses map[string]float64
		quotes     map[string]sentio.Quote
		orders     []sentio.Order
		ok         bool
		err        error
	)

	// skip too early orders
	if now.Hour() == 9 && now.Minute() < 50 {
		return nil
	}

	// close all orders before market close
	if now.Hour() == 15 && now.Minute() > 50 {
		return util.CloseAllOrders(market, strategy)
	}

	if stamp, ok = probability.First(); !ok {
		return nil
	}

	// retrieve running orders
	if orders, err = market.Orders(sentio.OrderListCriteria{
		Status:  "open",
		Symbols: symbols,
	}); err != nil {
		return err
	}

	if stoplosses, err = indicator.AtrTrailingStopLoss(market, 1, ATR_PERIOD, ATR_MULTIPLIER, ATR_HHV, symbols...); err != nil {
		return err
	}

	// update stoplosses or close running orders
	hasClosedOrders := false
	for i := range orders {
		if strategy.PositionSymbols()[sentio.LONG] == orders[i].GetSymbol() &&
			stamp.Value < strategy.PositionProbabilities()[sentio.LONG] {
			if _, err = market.CloseOrder(orders[i]); err != nil {
				return err
			}

			hasClosedOrders = true
			continue
		}

		if strategy.PositionSymbols()[sentio.SHORT] == orders[i].GetSymbol() &&
			stamp.Value > strategy.PositionProbabilities()[sentio.SHORT] {
			if _, err = market.CloseOrder(orders[i]); err != nil {
				return err
			}

			hasClosedOrders = true
			continue
		}

		if f, ok := stoplosses[orders[i].GetSymbol()]; ok && f > 0 {
			if err = market.UpdateOrder(orders[i].GetId(), f); err != nil {
				return err
			}

			continue
		}
	}

	// Prevent new orders if we just closed one
	if hasClosedOrders || len(orders) > 0 {
		return nil
	}

	// Prevent BUYs on closing market
	if now.Hour() == 15 && now.Minute() > 46 {
		return nil
	}

	// Prevent Market.CreateOrder while probability is not clear
	if stamp.Value == 1 || stamp.Value < 0 {
		return nil
	}

	if quotes, err = market.Quotes(symbols...); err != nil {
		return err
	}

	var t = sentio.SHORT
	if stamp.Value > 1 {
		t = sentio.LONG
	}

	return util.CreateOrder(market, strategy, t, probability, quotes, stoplosses)
}