package status

import (
	"context"
	"errors"
	"fmt"
	"os"
	"sync"

	"github.com/tidbcloud/observability/vm-migration/internal/schema"
)

// Store is the in-memory + on-disk consistency manager for a status JSON file.
// All read/write access is serialized through an internal mutex, so concurrent
// Mutate / View calls from multiple goroutines are safe.
//
// Each successful Mutate runs the caller's mutation function under the write
// lock, refreshes doc.Summary from doc.Shards, and invokes writeFn. Callers do
// NOT call WriteJSONAtomic / RefreshSummary manually; the store owns that
// side-effect.
//
// Mutate is the unit of atomicity and is copy-on-write: the fn mutates a
// clone of the current doc, and the clone replaces the in-memory doc only
// after writeFn succeeds. On fn error or write failure the in-memory doc is
// unchanged. A Mutate fn must perform ONE coherent state change.
//
// Mutate checks ctx before acquiring the lock and again after — a cancelled
// context never triggers fn execution.
type Store struct {
	writeFn func(*schema.Document) error
	mu      sync.Mutex
	doc     *schema.Document
}

// Open loads an existing status file, or initializes a new one via init when the
// file does not exist. init may be nil — in that case, a missing file is an
// error. The freshly-initialized doc is written to disk immediately.
func Open(path string, init func() (*schema.Document, error)) (*Store, error) {
	doc, err := LoadDocument(path)
	if err == nil {
		return &Store{writeFn: diskWriter(path), doc: doc}, nil
	}
	if !errors.Is(err, os.ErrNotExist) || init == nil {
		return nil, err
	}
	doc, err = init()
	if err != nil {
		return nil, fmt.Errorf("init status: %w", err)
	}
	RefreshSummary(doc)
	w := diskWriter(path)
	if err := w(doc); err != nil {
		return nil, fmt.Errorf("write initial status: %w", err)
	}
	return &Store{writeFn: w, doc: doc}, nil
}

// NewInMemory wraps an existing doc with no disk persistence. Useful for
// callers that need Store's API surface but must not touch disk.
func NewInMemory(doc *schema.Document) *Store {
	return &Store{writeFn: func(*schema.Document) error { return nil }, doc: doc}
}

// View runs fn under the read lock with the current doc. fn must not mutate the
// doc — if it does, the in-memory copy and the on-disk copy will diverge
// silently. Use Mutate for any state change.
func (s *Store) View(fn func(*schema.Document)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	fn(s.doc)
}

// Mutate runs fn under the write lock with a clone of the current doc. On fn
// returning nil, RefreshSummary + writeFn run automatically and the clone
// becomes the store's doc. On fn or writeFn returning a non-nil error,
// nothing is written and the in-memory doc is left unchanged.
//
// ctx is checked twice (before and after acquiring the lock) so a cancelled
// context short-circuits even when the lock is contested.
func (s *Store) Mutate(ctx context.Context, fn func(*schema.Document) error) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	next := schema.CloneDocument(s.doc)
	if err := fn(next); err != nil {
		return err
	}
	RefreshSummary(next)
	if err := s.writeFn(next); err != nil {
		return err
	}
	s.doc = next
	return nil
}

// Document returns the current in-memory doc WITHOUT taking the lock. Provided
// exclusively for tests that need to assert post-state after a Mutate;
// production callers should use View or Mutate.
func (s *Store) Document() *schema.Document { return s.doc }

// diskWriter returns a writeFn that atomic-writes to path.
func diskWriter(path string) func(*schema.Document) error {
	return func(doc *schema.Document) error { return DumpJSONAtomic(path, doc) }
}