package exchange

import (
	"context"
	"fmt"
	"io"
	"runtime"
	"sig-pub/api/pb"
	"sig-pub/pkg/aside"
	"sig-pub/pkg/data"
	"sig-pub/pkg/storage/kvrocks"
	"sig-pub/pkg/types"
	"sig-pub/pkg/utils/collect"
	"sig-pub/pkg/utils/times"
	"sig-pub/pkg/zlog"
	"sync"
	"sync/atomic"
	"time"

	"google.golang.org/grpc"
)

type ExchangeGrpcServer struct {
	pb.UnimplementedExchangeServiceServer
	exchangeMap         map[types.Exchange]*Exchange
	tradeInstanceAside  *aside.TradeInstanceAside
	exchangeDataService *ExchangeDataService
	kvdb                *kvrocks.KVRocksDB

	klineStreamId  int64
	klinePublisher *Publisher[int64, grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]]
}

// exchanges: 支持的数据源交易所
func NewExchangeGrpcServer(
	tradeInstanceAside *aside.TradeInstanceAside,
	exchangeDataService *ExchangeDataService,
	kvdb *kvrocks.KVRocksDB,
	exchanges ...*Exchange,
) *ExchangeGrpcServer {
	exchangeMap := make(map[types.Exchange]*Exchange)
	for _, exchange := range exchanges {
		exchangeMap[exchange.ExType] = exchange
	}

	return &ExchangeGrpcServer{
		exchangeMap:         exchangeMap,
		tradeInstanceAside:  tradeInstanceAside,
		exchangeDataService: exchangeDataService,
		kvdb:                kvdb,
		klinePublisher:      NewPublisher[int64, grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]](16),
	}
}

func (svc *ExchangeGrpcServer) Init() (err error) {
	svc.subscribeExchanges()
	return
}

// 订阅交易所推送行情
func (svc *ExchangeGrpcServer) subscribeExchanges() {
	// consumerKline
	// 交易所订阅交易产品
	for _, exchange := range svc.exchangeMap {
		go func(exchange *Exchange) {
			// get exchange trade instances
			insts, err := svc.tradeInstanceAside.ListExchangeTradeInstance(context.Background(), exchange.ExType)
			if err != nil {
				zlog.Error(err)
				return
			}

			var exchangeInstIds []string
			var processingInsts []types.TradeInstance
			for _, inst := range insts {
				exchangeInstIds = append(exchangeInstIds, inst.ExchangeInstId)
				tradeInst := &types.TradeInstance{
					InstId:         inst.InstId,
					Status:         inst.Status,
					PriceSz:        0,
					QuantitySz:     0,
					ExchangeInstId: inst.ExchangeInstId,
					Exchange:       exchange.ExType,
				}
				exchange.Insts.Store(inst.ExchangeInstId, &ExchangeTradeInstance{
					Inst:          tradeInst,
					Status:        0,
					LiveMarkTs:    0,
					LiveKMarkTs:   make(map[types.Interval]int64),
					HistoryMarkTs: make(map[types.Interval]int64),
				})
				// 待初始化币种数据
				if inst.Status == data.StatusProcessing {
					processingInsts = append(processingInsts, *tradeInst)
				}
			}

			// instIds := []string{"BTC-USDT", "DOGE-USDT-SWAP"}
			err = exchange.Subscriber.SubscribeKline(exchangeInstIds...)
			if err != nil {
				zlog.Error(err)
				return
			}

			go func() {
				c := exchange.Subscriber.ConsumerKline()
				svc.consumerKline(exchange, c)
				// todo subscribe books 订单簿
				zlog.Infof("unsubscribe exchange: %s", exchange.ExType)
			}()

			// 初始化k线数据
			go svc.initialKlines(exchange, processingInsts)
		}(exchange)
	}
}

// consumerKline 消费交易所k线数据
func (svc *ExchangeGrpcServer) consumerKline(exchange *Exchange, c <-chan *types.ChannelKline) {
	exchangeType := exchange.ExType

	for {
		channelK, ok := <-c
		if !ok {
			return
		}

		// 交易所 instid 转 sig-instid
		var exInst *types.TradeInstance
		if inst, ok := exchange.Insts.Load(channelK.ExgInstId); ok && inst != nil {
			exInst = inst.Inst
			// 标记交易产品开始订阅k线时间
			if len(channelK.Klines) > 0 {
				kline := collect.MustMax(channelK.Klines, func(k *types.Kline) int64 { return k.Ts })
				if inst.GetLiveMarkTs(kline.Interval) == 0 {
					inst.SetLiveMarkTs(kline.Interval, kline.Ts)
				}
			}
			// confirmKlines := collect.Filter(channelK.Klines, func(_ int, k *types.Kline) bool { return k.Confirm })

		} else {
			zlog.Errorf("unknown exchange instId: %v, %s", channelK.Exchange, channelK.ExgInstId)
			continue
		}

		// publish to subscribers
		pubMsgMap := make(map[string]*pb.StreamKline)

		// instId := channelK.InstId
		pbExType, err := channelK.Exchange.Exchange2PB()
		if err != nil {
			zlog.Error(err)
			continue
		}

		var confirmKlines []*types.Kline
		for _, kline := range channelK.Klines {
			// zlog.Infof("recv kline: %#v", kline)
			confirm := 0
			if kline.Confirm {
				confirm = 1
				confirmKlines = append(confirmKlines, kline)
			}
			pubKey := fmt.Sprintf("/kline/%s/%s/%s/%d", exchangeType, exInst.InstId, kline.Interval, confirm)
			// todo 优化没有订阅者就跳过
			msg, ok := pubMsgMap[pubKey]
			if !ok {
				msg = new(pb.StreamKline)
				msg.InstId = exInst.InstId
				msg.Exchange = pbExType
				pubMsgMap[pubKey] = msg
			}

			pbk := kline.ToPBKline()
			msg.Klines = append(msg.Klines, pbk)
		}

		// tsdb storage
		if len(confirmKlines) > 0 {
			// todo 异步处理
			err := svc.exchangeDataService.SaveKlines(*exInst, confirmKlines)
			if err != nil {
				zlog.Errorf("kline save to tsdb error: ", err)
			}
		}

		for pubKey, msg := range pubMsgMap {
			if len(msg.Klines) == 0 {
				continue
			}
			subs := svc.klinePublisher.Publisher(pubKey)
			for _, sub := range subs {
				if err := sub.Send(&pb.RspStreamSubscribeKline{Kline: msg}); err != nil {
					zlog.Error(err)
				}
			}
		}
	}
}

// SubscribeKline 订阅k线stream
func (svc *ExchangeGrpcServer) SubscribeKline(stream grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]) (err0 error) {
	streamId := atomic.AddInt64(&svc.klineStreamId, 1)
	// subKey = /kline/exchange/instId/interval/confirm

	// 接收消息的goroutine
	recvChan := make(chan *pb.ReqStreamSubscribeKline)
	go func() {
		for {
			msg, err := stream.Recv()
			if err == io.EOF {
				zlog.Infof("关闭EOF: %v", err)
				close(recvChan)
				return
			}
			if err != nil {
				zlog.Infof("接收错误: %v", err)
				close(recvChan)
				return
			}
			zlog.Infof("recv stream msg: %#v", msg)
			recvChan <- msg
		}
	}()

	// 发送和处理消息
	for {
		select {
		case <-stream.Context().Done():
			// 客户端断开连接
			svc.klinePublisher.UnsubscribeAll(streamId)
			return stream.Context().Err()
		case msg, ok := <-recvChan:
			if !ok {
				// 接收通道关闭，结束流
				svc.klinePublisher.UnsubscribeAll(streamId)
				return
			}

			if msg.SubType == pb.SubscribeType_UnsubscribeAll {
				svc.klinePublisher.UnsubscribeAll(streamId)
				continue
			}
			for _, exchange := range msg.Exchanges {
				for _, instId := range msg.InstIds {
					for _, interval := range msg.Intervals {
						confirms := []int{1}
						if !msg.OnlyConfirm {
							confirms = append(confirms, 0)
						}
						for _, confirm := range confirms {
							subKey := fmt.Sprintf("/kline/%s/%s/%s/%d", exchange.String(), instId, interval, confirm)
							zlog.Infof("stream: %d sub: %s", streamId, subKey)
							switch msg.SubType {
							case pb.SubscribeType_Subscribe:
								svc.klinePublisher.Subscribe(subKey, streamId, stream)
							case pb.SubscribeType_Unsubscribe:
								svc.klinePublisher.Unsubscribe(subKey, streamId)
							}
						}
					}
				}
			}
		}
	}

	// for {
	// 	select {
	// 	case <-ctx.Done():
	// 		// 客户端断开连接
	// 		// log.Printf("Client disconnected from topic: %s", topic)
	// 		return
	// 	default:
	// 		// 模拟事件生成
	// 		event := &pb.SubscribeKlineStreamRsp{
	// 			Kline: &pb.StreamKline{
	// 				InstId:   "DOGE/USDT",
	// 				Exchange: pb.Exchange_OKX,
	// 				Klines: []*pb.Kline{
	// 					{
	// 						Ts: time.Now().Unix(),
	// 					},
	// 				},
	// 			},
	// 		}

	// 		// 推送事件
	// 		if err := streamRsp.Send(event); err != nil {
	// 			log.Printf("Failed to send event to client: %v", err)
	// 			return err
	// 		}

	// 		// 模拟事件间隔
	// 		time.Sleep(2 * time.Second)
	// 	}
	// }
}

const (
	KlineBefore0      int64  = 1672502400000               // k线开始数据 2023-01-01 00:00:00 GMT+8
	HistoryKlineTsKey string = "history-kline-ts:%s:%s:%s" // exchange:sig-instid:interval
)

type fetchKlineTask struct {
	inst     types.TradeInstance
	interval types.Interval
	afterTs  int64
	beforeTs int64
	times    int32 // 重试次数
}

func (t fetchKlineTask) logKey() string {
	return fmt.Sprintf("%s:%s:%s:%d:%d", t.inst.Exchange, t.inst.InstId, t.interval, t.beforeTs, t.afterTs)
}

// initialKline 初始化交易产品历史k线数据
func (svc *ExchangeGrpcServer) initialKlines(exchange *Exchange, insts []types.TradeInstance) {
	// 记录成功和失败的交易产品
	var success, failed []types.TradeInstance

	for _, inst := range insts {
		err := svc.initialTradeInstanceKlines(exchange, inst)
		if err != nil {
			zlog.Errorf("initial fetch trade instance error: %s(%s), err=%v", inst.InstId, inst.Exchange, err)
			failed = append(failed, inst)
		} else {
			success = append(success, inst)
		}
	}

	zlog.Infof("%d insts initial finished, success %d, failed %d", len(insts), len(success), len(failed))
}

var (
	SingleTaskMaxFailTimes int32 = 10
)

// initTradeInstanceKlines 初始化交易产品历史k线数据
func (svc *ExchangeGrpcServer) initialTradeInstanceKlines(exchange *Exchange, inst types.TradeInstance) (err error) {
	// 并发数
	concurrent := max(8, runtime.NumCPU()*2)
	// 任务 channel
	taskCh := make(chan fetchKlineTask, concurrent)
	retryTaskCh := make(chan fetchKlineTask, concurrent)

	// 发布任务数, 成功任务数, 失败任务次数
	var pubTasks, subTasks, failTasks atomic.Int32
	var pubTaskDone atomic.Bool                        // 所有任务已发布
	var historyMarkTs = make(map[types.Interval]int64) // 已初始化最后k线时间戳
	var historyMarkTsMu sync.Mutex
	ctx, cancel := context.WithCancel(context.Background())

	defer func() {
		if err != nil {
			return
		}
		for interval, ts := range historyMarkTs {
			tsKey := fmt.Sprintf(HistoryKlineTsKey, inst.Exchange, inst.InstId, interval)
			if ex := svc.kvdb.SetI64(context.Background(), tsKey, ts); ex != nil {
				zlog.Errorf("history mark inititaled ts error: key=%s, ts=%d", tsKey, ts)
			}
		}
	}()

	go func() {
		defer func() {
			pubTaskDone.Store(true)
			// 无任务处理
			if subTasks.Load() == 0 {
				cancel()
			}
			zlog.Infof("trade instance initial kline %s(%s), pub %d fetch tasks", inst.InstId, inst.Exchange, pubTasks.Load())
		}()

		// for interval, intervalAdder := range types.SupportedIntervals {
		interval := types.Interval1d
		intervalAdder := types.SupportedIntervals[interval]
		// history 未补全前, history写 kvdb ts mark, 补全后 ws live 写 ts mark
		tsKey := fmt.Sprintf(HistoryKlineTsKey, inst.Exchange, inst.InstId, interval)
		beforeTs, ex := svc.kvdb.GetI64(context.Background(), tsKey)
		if ex != nil {
			err = ex
			zlog.Error(err)
			cancel()
			return
		}
		if beforeTs == 0 {
			beforeTs = intervalAdder(KlineBefore0, -1)
		}
		for {
			// 对比 ws 获取的实时k线
			exchangeInst, ok := exchange.Insts.Load(inst.ExchangeInstId)
			if !ok {
				err = fmt.Errorf("not load exchange trade instance: %s", inst.ExchangeInstId)
				cancel()
				return
			}
			// 订阅完成
			if exchangeInst.LiveMarkTs != 0 && beforeTs >= exchangeInst.LiveMarkTs {
				// history status -> ok
				break
			}
			if beforeTs > time.Now().UnixMilli() {
				if exchangeInst.LiveMarkTs != 0 {
					// history status -> ok
				} else {
					// live status -> not ok
				}
				break
			}

			afterTs := intervalAdder(beforeTs, 101)
			task := fetchKlineTask{
				inst:     inst,
				interval: interval,
				afterTs:  afterTs,
				beforeTs: beforeTs,
				times:    0,
			}

			// 发布任务
			select {
			case taskCh <- task:
				pubTasks.Add(1)
			case <-ctx.Done():
				return
			}

			beforeTs = intervalAdder(afterTs, -1)
		}
	}()

	// 任务消费器 8协程并行
	wg := new(sync.WaitGroup)
	for range concurrent {
		wg.Add(1)
		go func() {
			defer wg.Done()

			var task fetchKlineTask
			for {
				select {
				case <-ctx.Done():
					return
				case task = <-taskCh:
				case task = <-retryTaskCh:
				}

				if task.times > 0 {
					zlog.Infof("retry fetch history kline task %d times: task -> %s", task.times, task.logKey())
				}

				if lastKlineTs, ex := svc.fetchTaskKlines(exchange, task); ex != nil {
					failTasks.Add(1)
					if task.times >= SingleTaskMaxFailTimes {
						err = fmt.Errorf("task failed to many times %d, key: %s, err: %v", task.times, task.logKey(), ex)
						cancel()
						return
					}
					// retry task
					task.times++
					select {
					case retryTaskCh <- task:
					case <-ctx.Done():
						return
					}
				} else {
					// 周期任务最后kline时间
					if lastKlineTs != 0 {
						historyMarkTsMu.Lock()
						if ts, ok := historyMarkTs[task.interval]; ok {
							if lastKlineTs > ts {
								historyMarkTs[task.interval] = lastKlineTs
							}
						} else {
							historyMarkTs[task.interval] = lastKlineTs
						}
						historyMarkTsMu.Unlock()
					}

					zlog.Debugf("trade instance initial kline tasks processing: %s(%s), pub %d, sub %d, fail %d", inst.InstId, inst.Exchange, pubTasks.Load(), subTasks.Load(), failTasks.Load())
					// 任务都已执行成功结束
					subs := subTasks.Add(1)
					if pubTaskDone.Load() && subs >= pubTasks.Load() {
						zlog.Infof("trade instance initial kline tasks success finished, %s(%s), pub %d, sub %d, fail %d", inst.InstId, inst.Exchange, pubTasks.Load(), subTasks.Load(), failTasks.Load())
						cancel()
						return
					}
				}
			}
		}()
	}
	wg.Wait()
	return
}

func (svc *ExchangeGrpcServer) fetchTaskKlines(exchange *Exchange, task fetchKlineTask) (lastKlineTs int64, err error) {
	interval, afterTs, beforeTs := task.interval, task.afterTs, task.beforeTs
	klines, err := exchange.Fetcher.FetchHistoryKlines(context.Background(), task.inst.ExchangeInstId, interval, afterTs, beforeTs)
	if err != nil {
		zlog.Errorf("fetch history kline task error: task -> %s, err -> %v", task.logKey(), err)
		return
	}
	if len(klines) == 0 {
		return
	}

	loc, _ := time.LoadLocation("Asia/Shanghai")

	sts, ets := klines[0].Ts, klines[len(klines)-1].Ts
	lastKlineTs = max(sts, ets)
	ss := time.UnixMilli(sts).In(loc).Format(times.FORMAT_DATE)
	ee := time.UnixMilli(ets).In(loc).Format(times.FORMAT_DATE)
	zlog.Infof("fetch interval %s %d~%d klines: ret=%d~%d, %d klines, %s~%s", interval, beforeTs, afterTs, ets, sts, len(klines), ee, ss)

	// store to tsdb
	err = svc.exchangeDataService.SaveKlines(task.inst, klines)
	if err != nil {
		zlog.Errorf("save history klines to tsdb error: task -> %s, err -> %v", task.logKey(), err)
		return
	}
	return
}