package checkpoint

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
)

type MetricShardState struct {
	Index       int    `json:"index"`
	Status      string `json:"status"`
	Key         string `json:"key,omitempty"`
	FileSize    int64  `json:"fileSize,omitempty"`
	StartedAt   int64  `json:"startedAt,omitempty"`
	CompletedAt int64  `json:"completedAt,omitempty"`
}

type MetricState struct {
	TaskID       string             `json:"taskId"`
	StartTime    int64              `json:"startTime"`
	EndTime      int64              `json:"endTime"`
	TimeInterval string             `json:"timeInterval"`
	Shards       []MetricShardState `json:"shards"`
	LastUpdated  int64              `json:"lastUpdated"`
}

type MetricCheckpoint struct {
	mut        sync.Mutex
	state      *MetricState
	shardCount int
	path       string
}

func NewMetricCheckpoint(path string, shardCount int) *MetricCheckpoint {
	shards := make([]MetricShardState, 0, shardCount)
	for i := 0; i < shardCount; i++ {
		shards = append(shards, MetricShardState{
			Index:  i,
			Status: StatusPending,
		})
	}

	path = filepath.Join(path, "checkpoint.json")
	logrus.Infof("MetricCheckpoint initialized: path=%s", path)
	return &MetricCheckpoint{
		path:       path,
		shardCount: shardCount,
		state: &MetricState{
			Shards: shards,
		},
	}
}

// UpdateState atomically modifies and persists state.
//
// The callback pattern ensures the mutex is held for both modification
// and persistence, preventing lost updates from concurrent goroutines.
func (cp *MetricCheckpoint) UpdateState(f func(*MetricState)) {
	cp.mut.Lock()
	defer cp.mut.Unlock()

	f(cp.state)
	cp.state.LastUpdated = time.Now().Unix()
	if err := saveCheckpoint(cp.state, cp.path); err != nil {
		logrus.Errorf("Failed to save checkpoint: %v", err)
	}
}

// UpdateShardFailed marks a shard as failed.
//
// This is a convenience method for the common failure case where the status and
// lastUpdated timestamp need updating.
func (cp *MetricCheckpoint) UpdateShardFailed(index int) {
	cp.mut.Lock()
	defer cp.mut.Unlock()

	cp.state.Shards[index].Status = StatusFailed
	cp.state.LastUpdated = time.Now().Unix()
	if err := saveCheckpoint(cp.state, cp.path); err != nil {
		logrus.Errorf("Failed to save checkpoint: %v", err)
	}
}

// AccessState provides read-only access to state under lock.
//
// Used by MetricWorker to check shard completion status before enqueuing,
// enabling already-completed shards to be skipped on resume.
func (cp *MetricCheckpoint) AccessState(f func(*MetricState)) {
	cp.mut.Lock()
	defer cp.mut.Unlock()

	f(cp.state)
}

// Load restores state from disk if a checkpoint exists.
//
// Returns nil if no checkpoint exists (first run), allowing the backup to proceed normally with fresh state.
func (cp *MetricCheckpoint) Load() error {
	cp.mut.Lock()
	defer cp.mut.Unlock()

	b, err := os.ReadFile(cp.path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("read checkpoint: %w", err)
	}

	var state MetricState
	if err := json.Unmarshal(b, &state); err != nil {
		return fmt.Errorf("load checkpoint: %w", err)
	}
	if len(state.Shards) != cp.shardCount {
		return fmt.Errorf("checkpoint shard count mismatch, want %d, got %d", cp.shardCount, len(state.Shards))
	}
	cp.state = &state

	logrus.Infof("MetricCheckpoint loaded: path=%s, shards=%d", cp.path, len(state.Shards))
	return nil
}