package client

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"io"

	"github.com/pingcap/log"
	"github.com/tidbcloud/ffs/config"
	"github.com/tidbcloud/ffs/pkg/v1/proto"
	"github.com/tidbcloud/ffs/realm/fserror"
)

const (
	maxBatchCount = 65536
	maxBatchSize  = config.MaxWriteBufferSize
)

type callback struct {
	req    *proto.BatchCommandsRequest_Request
	respCh chan *proto.BatchCommandsResponse_Response
	doneFn func()
}

func (cb *callback) Done(resp *proto.BatchCommandsResponse_Response) {
	if cb != nil {
		cb.respCh <- resp
		if cb.doneFn != nil {
			cb.doneFn()
		}
	}
}

func (cb *callback) Wait() *proto.BatchCommandsResponse_Response {
	return <-cb.respCh
}

func newCallback(req *proto.BatchCommandsRequest_Request) *callback {
	return &callback{
		req:    req,
		respCh: make(chan *proto.BatchCommandsResponse_Response, 2),
	}
}

var batchPool = &sync.Pool{
	New: func() interface{} {
		return &proto.BatchCommandsRequest{}
	},
}

func getBatch() *proto.BatchCommandsRequest {
	return batchPool.Get().(*proto.BatchCommandsRequest)
}

func putBatch(batch *proto.BatchCommandsRequest) {
	if batch != nil {
		for i := 0; i < len(batch.Requests); i++ {
			batch.Requests[i] = nil
		}
		batch.Requests = batch.Requests[:0]
		batch.RequestIds = batch.RequestIds[:0]
		batchPool.Put(batch)
	}
}

const reqChSize = 4096

var reqChPool = &sync.Pool{
	New: func() interface{} {
		return make(chan *callback, reqChSize)
	},
}

func getReqCh() chan *callback {
	return reqChPool.Get().(chan *callback)
}

func putReqCh(reqCh chan *callback) {
	pending := len(reqCh)
	for i := 0; i < pending; i++ {
		select {
		case <-reqCh:
		default:
		}
	}
	if cap(reqCh) == reqChSize && len(reqCh) == 0 {
		reqChPool.Put(reqCh)
	}
}

type batchCommandsClient struct {
	reqCh   chan *callback
	ctx     context.Context
	cancel  context.CancelFunc
	batched sync.Map
	idAlloc uint64
	recv    uint64
	mu      sync.Mutex
	stream  *stream
	closed  int32
}

type batchCommands func(ctx context.Context) (proto.RealmService_BatchCommandsClient, error)

func newBatchCommandsClient(batchCommands batchCommands) (*batchCommandsClient, error) {
	c := &batchCommandsClient{
		reqCh:   getReqCh(),
		batched: sync.Map{},
		idAlloc: 1,
	}
	c.ctx, c.cancel = context.WithCancel(context.Background())
	go c.run(batchCommands)
	return c, nil
}

func (c *batchCommandsClient) run(batchCommands batchCommands) {
	for {
		c.mu.Lock()
		if atomic.LoadInt32(&c.closed) > 0 {
			c.mu.Unlock()
			return
		}
		stream, err := newStream(batchCommands, c)
		if err != nil {
			c.mu.Unlock()
			log.S().Warnf("new stream err %s", err)
			time.Sleep(time.Millisecond * 500)
			continue
		}
		c.stream = stream
		c.mu.Unlock()
		log.S().Infof("new stream ")
		stream.wg.Wait()
	}
}

func (c *batchCommandsClient) Send(cb *callback) error {
	if c == nil || atomic.LoadInt32(&c.closed) > 0 {
		return fserror.ErrClosed
	}
	select {
	case c.reqCh <- cb:
		return nil
	case <-c.ctx.Done():
		return c.ctx.Err()
	}
}

func (c *batchCommandsClient) retryBatchRequest() *proto.BatchCommandsRequest {
	batch := getBatch()
	totalSize := 0
	for i := c.recv + 1; i < c.idAlloc; i++ {
		value, ok := c.batched.Load(i)
		if !ok {
			log.S().Warnf("requestID %d not found", i)
			continue
		}
		totalSize += c.appendRetryRequest(batch, i, value.(*callback))
	}
	return batch
}

func (c *batchCommandsClient) appendRetryRequest(batch *proto.BatchCommandsRequest, id uint64, cb *callback) int {
	size := c.getEstimatedSize(cb.req)
	batch.Requests = append(batch.Requests, cb.req)
	batch.RequestIds = append(batch.RequestIds, id)
	log.S().Infof("append retry request id %d", id)
	return size
}

func (c *batchCommandsClient) getBatchRequest(cb *callback) *proto.BatchCommandsRequest {
	batch := getBatch()
	totalSize := c.appendRequest(batch, cb)
	chLen := len(c.reqCh)
	for i := 0; i < chLen && i < maxBatchCount && totalSize < maxBatchSize; i++ {
		newCb := <-c.reqCh
		totalSize += c.appendRequest(batch, newCb)
	}
	return batch
}

func (c *batchCommandsClient) appendRequest(batch *proto.BatchCommandsRequest, cb *callback) int {
	size := c.getEstimatedSize(cb.req)
	batch.Requests = append(batch.Requests, cb.req)
	batch.RequestIds = append(batch.RequestIds, c.idAlloc)
	c.batched.Store(c.idAlloc, cb)
	log.S().Debugf("append request id %d", c.idAlloc)
	c.idAlloc++
	return size
}

func (c *batchCommandsClient) getEstimatedSize(req *proto.BatchCommandsRequest_Request) int {
	return req.GetCmd().Size()
}

func (c *batchCommandsClient) close() (err error) {
	c.mu.Lock()
	if atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		reqCh := c.reqCh
		c.reqCh = nil
		putReqCh(reqCh)
	}
	stream := c.stream
	c.mu.Unlock()
	c.cancel()
	if stream != nil {
		err = stream.close()
		stream.wg.Wait()
	}
	return err
}

type stream struct {
	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc
	stream proto.RealmService_BatchCommandsClient
	client *batchCommandsClient
}

func newStream(batchCommands batchCommands, client *batchCommandsClient) (*stream, error) {
	ctx, cancel := context.WithCancel(context.Background())
	s := &stream{
		ctx:    ctx,
		cancel: cancel,
		client: client,
	}
	var err error
	s.stream, err = batchCommands(ctx)
	if err != nil {
		return nil, err
	}
	batch := client.retryBatchRequest()
	err = s.Send(batch)
	putBatch(batch)
	if err != nil {
		return nil, err
	}
	s.wg.Add(2)
	go s.batchSendLoop()
	go s.batchRecvLoop()
	return s, nil
}

func (c *stream) batchSendLoop() {
	defer c.wg.Done()
	for {
		select {
		case req := <-c.client.reqCh:
			batch := c.client.getBatchRequest(req)
			if err := c.stream.Send(batch); err != nil {
				log.S().Error(err)
				c.close()
			}
			putBatch(batch)
		case <-c.ctx.Done():
			log.S().Info("stream done")
			return
		}
	}
}

func (c *stream) Send(batch *proto.BatchCommandsRequest) error {
	return c.stream.Send(batch)
}

func (c *stream) batchRecvLoop() {
	defer c.wg.Done()
	for {
		resp, err := c.stream.Recv()
		if err == nil && resp.NotMount {
			err = fserror.ErrNotMount
		}
		if err != nil {
			if fserror.IsCanceled(err) || err == io.EOF {
				log.S().Info(err)
			} else {
				log.S().Error(err)
			}
			c.close()
			return
		}
		ids := resp.GetRequestIds()
		responses := resp.GetResponses()
		for i, requestID := range ids {
			value, ok := c.client.batched.Load(requestID)
			if !ok {
				log.S().Warnf("requestID %d not found", requestID)
				continue
			}
			cb := value.(*callback)
			cb.Done(responses[i])
			c.client.batched.Delete(requestID)
			c.client.recv = requestID
		}
	}
}

func (c *stream) close() error {
	c.cancel()
	err := c.stream.CloseSend()
	return err
}