package worker

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/sirupsen/logrus"

	"github.com/tidbcloud/observability/backup-worker/pkg/client"
	"github.com/tidbcloud/observability/backup-worker/pkg/configs"
	"github.com/tidbcloud/observability/backup-worker/pkg/source"
	"github.com/tidbcloud/observability/backup-worker/pkg/worker/checkpoint"
)

type MetricWorker struct {
	cfg     *configs.Config
	cp      *checkpoint.MetricCheckpoint
	shards  []configs.MetricShard
	sourcer source.Sourcer
	sinker  client.Client
}

func NewMetricWorker(ctx context.Context, cfg *configs.Config) (*MetricWorker, error) {
	sinker, err := client.New(ctx, *cfg)
	if err != nil {
		return nil, err
	}
	shards := cfg.Sources.Shards()
	return &MetricWorker{
		cfg:     cfg,
		cp:      checkpoint.NewMetricCheckpoint(filepath.Join(cfg.Checkpoint.Path, cfg.Mode, cfg.ID()), len(shards)),
		shards:  shards,
		sourcer: source.New(ctx, *cfg),
		sinker:  sinker,
	}, nil
}

func (w *MetricWorker) Close() error {
	return w.sinker.Close()
}

func (w *MetricWorker) Run(ctx context.Context) error {
	if err := w.cp.Load(); err != nil {
		return err
	}

	logrus.Infof("Starting metrics backup: sink=%s, shards=%d, concurrency=%d",
		w.sinker.Type(), len(w.shards), w.cfg.Execution.Concurrency)

	return runWorker(ctx, w.shards, w.cfg.Execution.Concurrency,
		func(i int, shard configs.MetricShard) bool {
			var skip bool
			w.cp.AccessState(func(state *checkpoint.MetricState) {
				skip = state.Shards[i].Status == checkpoint.StatusCompleted
			})
			return skip
		},
		w.process,
	)
}

func (w *MetricWorker) process(ctx context.Context, shard configs.MetricShard) error {
	idx := shard.Index
	key := fmt.Sprintf("%s/%s.object", w.cfg.ID(), shard.ID())

	logrus.Infof("Start process shard: %+v", shard)
	w.cp.UpdateState(func(state *checkpoint.MetricState) {
		state.Shards[idx].Status = checkpoint.StatusRunning
		state.Shards[idx].Key = key
		state.Shards[idx].StartedAt = time.Now().Unix()
	})

	ok, err := w.sinker.CheckObject(ctx, key)
	if err != nil {
		w.cp.UpdateShardFailed(idx)
		return err
	}
	if ok {
		logrus.Infof("Object already exists, skipping: shard=%+v", shard)
		w.cp.UpdateState(func(state *checkpoint.MetricState) {
			state.Shards[idx].Status = checkpoint.StatusCompleted
			state.Shards[idx].Key = key
			state.Shards[idx].CompletedAt = time.Now().Unix()
		})
		return nil
	}

	backupFile, err := w.pull(ctx, shard)
	defer os.Remove(backupFile)
	if err != nil {
		w.cp.UpdateShardFailed(idx)
		return err
	}

	if err := w.sinker.SinkObject(ctx, key, backupFile); err != nil {
		w.cp.UpdateShardFailed(idx)
		return err
	}

	fileInfo, err := os.Stat(backupFile)
	if err != nil {
		w.cp.UpdateShardFailed(idx)
		return err
	}

	w.cp.UpdateState(func(state *checkpoint.MetricState) {
		state.Shards[idx].Status = checkpoint.StatusCompleted
		state.Shards[idx].Key = key
		state.Shards[idx].FileSize = fileInfo.Size()
		state.Shards[idx].CompletedAt = time.Now().Unix()
	})

	logrus.Infof("MetricShard completed: %+v", shard)
	return nil
}

func (w *MetricWorker) pull(ctx context.Context, shard configs.MetricShard) (string, error) {
	backupFile, err := os.CreateTemp("", "backup-*.bin")
	if err != nil {
		return "", err
	}
	defer backupFile.Close()

	readCloser, err := w.sourcer.Pull(ctx, shard.Start, shard.End)
	if err != nil {
		return "", err
	}
	defer readCloser.Close()

	if _, err := io.Copy(backupFile, readCloser); err != nil {
		return "", err
	}
	return backupFile.Name(), nil
}