package llmjob

import (
	"github.com/banbox/banbot/config"
	"github.com/banbox/banbot/core"
	"github.com/banbox/banbot/orm/ormo"
	"github.com/banbox/banbot/strat"
	ta "github.com/banbox/banta"
	"math"
)

// SupertrendState 存储每个品种的策略状态
type SupertrendState struct {
	trend     int     // 趋势状态，1表示上升趋势，-1表示下降趋势
	prevTrend int     // 前一期的趋势状态
	up        float64 // 上轨
	dn        float64 // 下轨
	tp        float64 // 止盈价格
}

func J3644(pol *config.RunPolicyConfig) *strat.TradeStrat {
	// 定义可调参数
	atrPeriod := pol.DefInt("atrPeriod", 5, core.PNorm(2, 20))
	atrMultiplier := pol.Def("atrMultiplier", 4.7, core.PNorm(2.0, 8.0))
	profitRatio := pol.Def("profitRatio", 0.382, core.PNorm(0.2, 1.0))
	useSmaATR := pol.DefInt("useSmaATR", 1, core.PUniform(0, 1)) // 1=true, 0=false
	srcType := pol.DefInt("srcType", 0, core.PUniform(0, 7))     // 0=hlc3, 1=high, 2=low, 3=open, 4=hl2, 5=close, 6=hlcc4, 7=ohlc4

	return &strat.TradeStrat{
		WarmupNum:     atrPeriod * 3, // 预热期设置为ATR周期的3倍，确保有足够数据
		HedgeOff:      true,          // 单向持仓模式
		EachMaxLong:   1,             // 最多开1个多单
		EachMaxShort:  1,             // 最多开1个空单
		RunTimeFrames: []string{"8h"},
		OnStartUp: func(s *strat.StratJob) {
			// 初始化策略状态
			s.More = &SupertrendState{
				trend:     1, // 初始趋势状态为上升趋势
				prevTrend: 1, // 前一期趋势状态也为上升趋势
				tp:        0, // 初始止盈价格
			}
		},
		OnBar: func(s *strat.StratJob) {
			e := s.Env
			state, _ := s.More.(*SupertrendState)

			// 获取数据源
			var src *ta.Series
			switch srcType {
			case 0:
				src = ta.HLC3(e.High, e.Low, e.Close)
			case 1:
				src = e.High
			case 2:
				src = e.Low
			case 3:
				src = e.Open
			case 4:
				src = ta.HL2(e.High, e.Low)
			case 5:
				src = e.Close
			case 6:
				// HLCC4: (high + low + close + close) / 4
				src = e.High.Add(e.Low).Add(e.Close).Add(e.Close).Div(4)
			case 7:
				// OHLC4: (open + high + low + close) / 4
				src = e.Open.Add(e.High).Add(e.Low).Add(e.Close).Div(4)
			default:
				src = ta.HLC3(e.High, e.Low, e.Close)
			}

			// 计算ATR
			var atr *ta.Series
			if useSmaATR == 1 {
				// 使用标准ATR计算
				atr = ta.ATR(e.High, e.Low, e.Close, atrPeriod)
			} else {
				// 使用SMA(TR)计算ATR
				trueRange := ta.TR(e.High, e.Low, e.Close)
				atr = ta.SMA(trueRange, atrPeriod)
			}

			// 计算超级趋势通道
			closePrice := e.Close.Get(0)
			prevClose := e.Close.Get(1)
			atrVal := atr.Get(0)
			srcVal := src.Get(0)

			// 计算上下轨
			up := srcVal - (atrMultiplier * atrVal)
			dn := srcVal + (atrMultiplier * atrVal)

			// 获取前一期上下轨
			prevUp := state.up
			prevDn := state.dn

			// 如果是第一期或前一期值无效，使用当前值
			if math.IsNaN(prevUp) {
				prevUp = up
			}
			if math.IsNaN(prevDn) {
				prevDn = dn
			}

			// 趋势跟踪逻辑
			if !math.IsNaN(prevClose) && prevClose > prevUp {
				up = math.Max(up, prevUp)
			}
			if !math.IsNaN(prevClose) && prevClose < prevDn {
				dn = math.Min(dn, prevDn)
			}

			// 更新趋势状态
			state.prevTrend = state.trend
			if state.trend == -1 && closePrice > prevDn {
				state.trend = 1
			} else if state.trend == 1 && closePrice < prevUp {
				state.trend = -1
			}

			// 保存当前上下轨值
			state.up = up
			state.dn = dn

			// 生成交易信号
			buySignal := state.trend == 1 && state.prevTrend == -1
			sellSignal := state.trend == -1 && state.prevTrend == 1

			// 获取当前订单
			longOrders := s.LongOrders
			shortOrders := s.ShortOrders

			// 更新多头订单的止损
			if len(longOrders) > 0 {
				for _, od := range longOrders {
					if od.Status >= ormo.InOutStatusFullEnter {
						// 保持之前的止盈价
						tp := state.tp
						if math.IsNaN(tp) || tp <= 0 {
							tp = od.InitPrice + (od.InitPrice-up)*profitRatio
						}
						// 更新止损价为上轨
						od.SetStopLoss(&ormo.ExitTrigger{
							Price: up,
							Tag:   "super_trend_sl",
						})
						// 更新止盈价
						od.SetTakeProfit(&ormo.ExitTrigger{
							Price: tp,
							Tag:   "super_trend_tp",
						})
						state.tp = tp
					}
				}
			}

			// 更新空头订单的止损
			if len(shortOrders) > 0 {
				for _, od := range shortOrders {
					if od.Status >= ormo.InOutStatusFullEnter {
						// 保持之前的止盈价
						tp := state.tp
						if math.IsNaN(tp) || tp <= 0 {
							tp = od.InitPrice - (dn-od.InitPrice)*profitRatio
						}
						// 更新止损价为下轨
						od.SetStopLoss(&ormo.ExitTrigger{
							Price: dn,
							Tag:   "super_trend_sl",
						})
						// 更新止盈价
						od.SetTakeProfit(&ormo.ExitTrigger{
							Price: tp,
							Tag:   "super_trend_tp",
						})
						state.tp = tp
					}
				}
			}

			// 买入信号
			if buySignal && len(longOrders) == 0 {
				// 计算止盈价
				tp := closePrice + (closePrice-up)*profitRatio
				state.tp = tp
				// 开多单，使用限价单
				s.OpenOrder(&strat.EnterReq{
					Tag:           "super_trend_long",
					Short:         false,
					OrderType:     core.OrderTypeLimit,
					Limit:         tp,
					StopLoss:      up,
					StopLossTag:   "super_trend_sl",
					TakeProfit:    tp,
					TakeProfitTag: "super_trend_tp",
				})
			}

			// 卖出信号
			if sellSignal && len(shortOrders) == 0 {
				// 计算止盈价
				tp := closePrice - (dn-closePrice)*profitRatio
				state.tp = tp
				// 开空单，使用限价单
				s.OpenOrder(&strat.EnterReq{
					Tag:           "super_trend_short",
					Short:         true,
					OrderType:     core.OrderTypeLimit,
					Limit:         tp,
					StopLoss:      dn,
					StopLossTag:   "super_trend_sl",
					TakeProfit:    tp,
					TakeProfitTag: "super_trend_tp",
				})
			}
		},
	}
}