package common

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sort"
	"time"

	"github.com/pingcap/kvproto/pkg/keyspacepb"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type KeyspaceRecord struct {
	KeyspaceId             uint32
	KeyspaceName           string
	KeyspaceCreatedAt      int64
	KeyspaceStateChangedAt int64

	KeyspaceMeta *keyspacepb.KeyspaceMeta

	ArchivedKeyspace    bool
	ArchivedPdRule      bool
	ArchivedRegionLabel bool
	ArchivedEtcdPath    bool
}

type KeyspaceRecords []*KeyspaceRecord

func (r *KeyspaceRecords) WriteTo(f *os.File) error {
	recordBytes, err := json.Marshal(r)
	if err != nil {
		return err
	}
	return WriteFile(f, recordBytes)
}

func (r *KeyspaceRecords) ReadFrom(f *os.File) error {
	recordBytes, err := io.ReadAll(f)
	if err != nil {
		return err
	}

	return json.Unmarshal(recordBytes, r)
}

func CacheArchiveKeyspaceRecords(f *os.File) (map[string]*KeyspaceRecord, error) {
	var records KeyspaceRecords

	err := records.ReadFrom(f)
	if err != nil {
		return nil, err
	}

	recordsMap := make(map[string]*KeyspaceRecord)
	for _, record := range records {
		idStr := fmt.Sprintf("%d", record.KeyspaceId)
		recordsMap[idStr] = record
	}

	return recordsMap, nil
}

func UpdateArchiveKeyspaceRecords(recordsMap map[string]*KeyspaceRecord, f *os.File) error {
	err := f.Truncate(0)
	if err != nil {
		return err
	}
	_, err = f.Seek(0, 0)
	if err != nil {
		return err
	}

	var records KeyspaceRecords
	for _, record := range recordsMap {
		records = append(records, record)
	}
	sort.Slice(records, func(i, j int) bool {
		return records[i].KeyspaceId < records[j].KeyspaceId
	})

	return records.WriteTo(f)
}

func GetKeyspace(ctx context.Context, metaServerAddr string, delayDuration time.Duration,
	expectedState keyspacepb.KeyspaceState, limit int) (KeyspaceRecords, error) {
	delayDate := time.Now().Add(-delayDuration)
	ksMeta, err := GetKeyspaceByState(ctx, metaServerAddr, expectedState)
	if err != nil {
		return nil, err
	}

	var records KeyspaceRecords

	for i := range ksMeta {
		keyspace := ksMeta[i]
		if keyspace.State != expectedState {
			log.Warn("keyspace state does not match expected state, skip keyspace",
				zap.String("keyspace", keyspace.Name),
				zap.Uint32("keyspaceId", keyspace.Id),
				zap.String("state", keyspace.State.String()),
				zap.String("expectedState", expectedState.String()))
			continue
		}
		if keyspace.StateChangedAt == 0 {
			log.Warn("keyspace missing StateChangedAt from meta server, skip keyspace",
				zap.String("keyspace", keyspace.Name),
				zap.Uint32("keyspaceId", keyspace.Id))
			continue
		}
		stateChangedTime := time.Unix(keyspace.StateChangedAt, 0)

		if stateChangedTime.Before(delayDate) {
			record := &KeyspaceRecord{}
			record.KeyspaceId = keyspace.Id
			record.KeyspaceName = keyspace.Name
			record.KeyspaceCreatedAt = keyspace.CreatedAt
			record.KeyspaceStateChangedAt = keyspace.StateChangedAt
			record.KeyspaceMeta = keyspace

			records = append(records, record)
		} else {
			log.Info("keyspace archive time did not exceed delay time",
				zap.String("keyspace", keyspace.Name),
				zap.Time("stateChangeTime", stateChangedTime),
				zap.Time("delayTime", delayDate))
		}
	}
	log.Info("get total keyspace",
		zap.Int("count", len(ksMeta)),
		zap.Int("eligible", len(records)),
		zap.String("state", expectedState.String()))
	return LimitKeyspaces(records, limit), nil
}

// LimitKeyspaces caps the number of keyspaces processed in a single run.
// A non-positive limit means no cap. Eligible keyspaces are processed oldest-first
// by StateChangedAt so leftover work continues in later runs.
func LimitKeyspaces(records KeyspaceRecords, limit int) KeyspaceRecords {
	if limit <= 0 || len(records) <= limit {
		return records
	}
	sorted := make(KeyspaceRecords, len(records))
	copy(sorted, records)
	sort.Slice(sorted, func(i, j int) bool {
		if sorted[i].KeyspaceStateChangedAt == sorted[j].KeyspaceStateChangedAt {
			return sorted[i].KeyspaceId < sorted[j].KeyspaceId
		}
		return sorted[i].KeyspaceStateChangedAt < sorted[j].KeyspaceStateChangedAt
	})
	log.Info("limit keyspaces for this run",
		zap.Int("eligible", len(sorted)),
		zap.Int("limit", limit),
		zap.Int("skipped", len(sorted)-limit))
	return sorted[:limit]
}
