package storage

import (
	"fmt"
	"sort"
	"strconv"
	"time"
)

type ExecStatusStg interface {
	OnNewTask(execId string, marshalledExecTask string) error
	SetTaskFinished(execId string) error
	LoadUnfinishedTasks() (marshalledExecTasks map[string]string, err error)
	IsTaskFinished(execId string) (bool, error)
	// CleanFinishedTasks removes finished task records from storage.
	// Time-based: tasks finished longer than retainDuration ago are deleted.
	// Count-based: if maxFinishedTasks > 0, after time-based cleanup the remaining
	// finished tasks are capped at maxFinishedTasks, with oldest deleted first.
	// Returns the total number of cleaned tasks.
	CleanFinishedTasks(retainDuration time.Duration, maxFinishedTasks int64) (int64, error)
}

type KvBasedExecStatusStg struct {
	kvs KvStg
}

func NewKvBasedExecStatusStg(kvs KvStg) *KvBasedExecStatusStg {
	return &KvBasedExecStatusStg{
		kvs: kvs,
	}
}

func (stg *KvBasedExecStatusStg) LoadUnfinishedTasks() (map[string]string, error) {
	chunk, err := stg.kvs.GetChunk(stg.activeTasksKey())
	if err != nil {
		return nil, err
	}

	unfinishedExecTasks := make(map[string]string)

	for _, it := range chunk {
		if len(it.RelativeKeyPath) != 1 {
			return nil, fmt.Errorf("KvBasedExecStatusStg.LoadUnfinished: invalid format `%v`", it.RelativeKeyPath)
		}
		execId := it.RelativeKeyPath[0]
		unfinishedExecTasks[execId] = it.Value
	}
	return unfinishedExecTasks, nil
}

func (stg *KvBasedExecStatusStg) OnNewTask(execId, marshalledExecTask string) error {
	err := stg.kvs.Set(stg.taskKey(execId), marshalledExecTask)
	if err != nil {
		return err
	}
	return stg.kvs.Set(stg.activeTaskKey(execId), marshalledExecTask)
}

func (stg *KvBasedExecStatusStg) SetTaskFinished(execId string) error {
	// Record finished_at first so the timestamp is always present when the
	// active marker is removed. This avoids a partial state where the task
	// looks finished but has no timestamp, which would cause
	// CleanFinishedTasks(retainDuration>0) to skip it indefinitely.
	if err := stg.kvs.Set(stg.finishedAtKey(execId), strconv.FormatInt(time.Now().UnixNano(), 10)); err != nil {
		return err
	}
	return stg.kvs.Delete(stg.activeTaskKey(execId))
}

func (stg *KvBasedExecStatusStg) IsTaskFinished(execId string) (bool, error) {
	chunk, err := stg.kvs.GetChunk(stg.activeTasksKey())
	if err != nil {
		return false, err
	}
	for _, it := range chunk {
		if len(it.RelativeKeyPath) != 1 {
			return false, fmt.Errorf("KvBasedExecStatusStg.LoadUnfinished: invalid format `%v`", it.RelativeKeyPath)
		}
		key := it.RelativeKeyPath[0]
		if key == execId {
			return false, nil
		}
	}
	return true, nil
}

func (stg *KvBasedExecStatusStg) activeTaskKey(execId string) []string {
	return append(stg.activeTasksKey(), execId)
}

func (stg *KvBasedExecStatusStg) activeTasksKey() []string {
	return []string{"running"}
}

func (stg *KvBasedExecStatusStg) taskKey(execId string) []string {
	return []string{"tasks", execId}
}

func (stg *KvBasedExecStatusStg) finishedAtKey(execId string) []string {
	return []string{"finished_at", execId}
}

func (stg *KvBasedExecStatusStg) CleanFinishedTasks(retainDuration time.Duration, maxFinishedTasks int64) (int64, error) {
	// Get all tasks
	allChunk, err := stg.kvs.GetChunk([]string{"tasks"})
	if err != nil {
		return 0, err
	}
	// Get active (unfinished) tasks
	activeChunk, err := stg.kvs.GetChunk(stg.activeTasksKey())
	if err != nil {
		return 0, err
	}
	activeSet := make(map[string]bool)
	for _, it := range activeChunk {
		if len(it.RelativeKeyPath) == 1 {
			activeSet[it.RelativeKeyPath[0]] = true
		}
	}

	// First pass: time-based cleanup. Collect surviving finished tasks for count-based pass.
	type finishedTask struct {
		execId     string
		finishTime time.Time
	}
	cutoff := time.Now().Add(-retainDuration)
	var cleaned int64
	var surviving []finishedTask

	for _, it := range allChunk {
		if len(it.RelativeKeyPath) != 1 {
			continue
		}
		execId := it.RelativeKeyPath[0]
		if activeSet[execId] {
			continue
		}
		// Check when the task was marked finished.
		finishedAtStr, existed, err := stg.kvs.Get(stg.finishedAtKey(execId))
		if err != nil {
			return cleaned, fmt.Errorf("failed to get finished_at for task %s: %w", execId, err)
		}
		if existed && finishedAtStr != "" {
			finishedAtNano, parseErr := strconv.ParseInt(finishedAtStr, 10, 64)
			if parseErr != nil {
				// Unparsable finished_at — skip to avoid violating retention guarantee.
				continue
			}
			finishedAt := time.Unix(0, finishedAtNano)
			if finishedAt.After(cutoff) {
				// Task finished recently, retain it (may still be subject to count limit).
				surviving = append(surviving, finishedTask{execId: execId, finishTime: finishedAt})
				continue
			}
		} else {
			// No finished_at marker: legacy task. Backfill and use it for cleanup decisions.
			// Use cutoff+retainDuration (= the original "now" used to compute cutoff) so that
			// when retainDuration==0 the backfilled time equals cutoff and the task is immediately eligible.
			backfilledAt := cutoff.Add(retainDuration)
			backfilledAtStr := strconv.FormatInt(backfilledAt.UnixNano(), 10)
			if err := stg.kvs.Set(stg.finishedAtKey(execId), backfilledAtStr); err != nil {
				return cleaned, fmt.Errorf("failed to backfill finished_at for legacy task %s: %w", execId, err)
			}
			// Use backfilled timestamp so this task participates in time/count-based cleanup now.
			finishedAt := backfilledAt
			if finishedAt.After(cutoff) {
				surviving = append(surviving, finishedTask{execId: execId, finishTime: finishedAt})
				continue
			}
		}
		// Delete the task record and its finished_at marker.
		if err := stg.kvs.Delete(stg.taskKey(execId)); err != nil {
			return cleaned, fmt.Errorf("failed to delete finished task %s: %w", execId, err)
		}
		if err := stg.kvs.Delete(stg.finishedAtKey(execId)); err != nil {
			return cleaned, fmt.Errorf("failed to delete finished_at for task %s: %w", execId, err)
		}
		cleaned++
	}

	// Second pass: enforce max finished tasks count limit.
	if maxFinishedTasks > 0 && int64(len(surviving)) > maxFinishedTasks {
		sort.Slice(surviving, func(a, b int) bool {
			return surviving[a].finishTime.Before(surviving[b].finishTime)
		})
		excess := int64(len(surviving)) - maxFinishedTasks
		for i := int64(0); i < excess; i++ {
			execId := surviving[i].execId
			if err := stg.kvs.Delete(stg.taskKey(execId)); err != nil {
				return cleaned, fmt.Errorf("failed to delete excess finished task %s: %w", execId, err)
			}
			if err := stg.kvs.Delete(stg.finishedAtKey(execId)); err != nil {
				return cleaned, fmt.Errorf("failed to delete finished_at for excess task %s: %w", execId, err)
			}
			cleaned++
		}
	}

	return cleaned, nil
}