package common

import (
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"sync"

	"github.com/pingcap/errors"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/infra-provider/pkg/k8s"
)

type taskStatus uint8

const (
	taskStatusNotStarted taskStatus = 0
	taskStatusRunning    taskStatus = 1
	taskStatusCompleted  taskStatus = 2
)

type TableProgress struct {
	TotalWritten int64      `json:"w"`
	TotalSize    int64      `json:"z"`
	Status       taskStatus `json:"s"`
	Message      string     `json:"m,omitempty"`
}

type TaskProgress struct {
	Tables  map[string]*TableProgress `json:"t"`
	Status  taskStatus                `json:"s"`
	Message string                    `json:"m,omitempty"`
}

type TasksResponse struct {
	ID    int64  `json:"id"`
	Error string `json:"error"`
}

type EngineCheckpoint struct {
	Status CheckpointStatus
	Chunks []*ChunkCheckpoint // a sorted array
}

type ChunkCheckpoint struct {
	Key               ChunkCheckpointKey
	FileMeta          SourceFileMeta
	ColumnPermutation []int
	Chunk             Chunk
	// Checksum          verify.KVChecksum
	Timestamp int64
}

type SourceFileMeta struct {
	Path string
	// Type        SourceType
	// Compression Compression
	SortKey string
}

// Chunk represents a portion of the data file.
type Chunk struct {
	Offset       int64
	EndOffset    int64
	PrevRowIDMax int64
	RowIDMax     int64
	Columns      []string
}

type ChunkCheckpointKey struct {
	Path   string
	Offset int64
}

type CheckpointStatus uint8

const (
	CheckpointStatusMissing         CheckpointStatus = 0
	CheckpointStatusMaxInvalid      CheckpointStatus = 25
	CheckpointStatusLoaded          CheckpointStatus = 30
	CheckpointStatusAllWritten      CheckpointStatus = 60
	CheckpointStatusClosed          CheckpointStatus = 90
	CheckpointStatusImported        CheckpointStatus = 120
	CheckpointStatusIndexImported   CheckpointStatus = 140
	CheckpointStatusAlteredAutoInc  CheckpointStatus = 150
	CheckpointStatusChecksumSkipped CheckpointStatus = 170
	CheckpointStatusChecksummed     CheckpointStatus = 180
	CheckpointStatusAnalyzeSkipped  CheckpointStatus = 200
	CheckpointStatusAnalyzed        CheckpointStatus = 210
)

type TableCheckpoint struct {
	Status    CheckpointStatus
	AllocBase int64
	Engines   map[int32]*EngineCheckpoint
	TableID   int64
}

type LightningClient struct {
	k8s.HttpLink
}

func newLightningClient(link k8s.HttpLink) LightningClient {
	return LightningClient{HttpLink: link}
}

func (g *LightningClient) CreateTask(cfg string) (int64, error) {
	body, err := g.Post("/tasks", cfg)
	if err != nil {
		return 0, err
	}
	task := TasksResponse{}
	err = json.Unmarshal(body, &task)
	if err != nil {
		return 0, err
	}
	if task.ID == 0 {
		return 0, errors.New(task.Error)
	}
	return task.ID, nil
}

func (g *LightningClient) GetTaskProgress() (*TaskProgress, error) {
	data, err := g.Get("/progress/task")
	if err != nil {
		return nil, errors.AddStack(err)
	}

	task := new(TaskProgress)
	err = json.Unmarshal(data, task)
	if err != nil {
		return nil, errors.Annotatef(err, "failed to unmarshal response %s", data)
	}
	return task, nil
}

func (g *LightningClient) GetTableCheckpoint(t string) (*TableCheckpoint, error) {
	v := url.Values{}
	v.Set("t", t)
	tableURL := fmt.Sprintf("/progress/table?%s", v.Encode())
	data, err := g.Get(tableURL)
	if err != nil {
		return nil, errors.AddStack(err)
	}

	cp := new(TableCheckpoint)
	err = json.Unmarshal(data, cp)
	if err != nil {
		return nil, errors.AddStack(err)
	}
	return cp, nil
}

func (g *LightningClient) GetAllTableCheckpoints(ctx context.Context, task *TaskProgress, fetchWorkerCount int) (map[string]*TableCheckpoint, error) {
	logger := log.FromContext(ctx)
	tblCheckpointMap := map[string]*TableCheckpoint{}

	var wg sync.WaitGroup
	reqCh := make(chan string, fetchWorkerCount)
	var mapLock sync.Mutex

	wg.Add(fetchWorkerCount)
	for i := 0; i < fetchWorkerCount; i++ {
		go func() {
			defer wg.Done()
			for tblName := range reqCh {
				tblCkpt, err := g.GetTableCheckpoint(tblName)
				if err != nil {
					logger.Error(err, "failed to get import table checkpoint, ignore it", "table", tblName)
					continue
				}
				mapLock.Lock()
				tblCheckpointMap[tblName] = tblCkpt
				mapLock.Unlock()
			}
		}()
	}
	var gerr error
sendReqLoop:
	for tblName, tblProgress := range task.Tables {
		if tblProgress == nil {
			continue
		}
		if tblProgress.Status == taskStatusNotStarted {
			continue
		}
		select {
		case reqCh <- tblName:
		case <-ctx.Done():
			logger.Error(ctx.Err(), "context is canceled")
			gerr = ctx.Err()
			break sendReqLoop
		}
	}
	close(reqCh)
	wg.Wait()
	return tblCheckpointMap, gerr
}
