package userbase

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"math/rand/v2"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"filippo.io/age"
	dexstorage "github.com/dexidp/dex/storage"
	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
	"github.com/go-git/go-git/v5/plumbing/object"
	"github.com/go-git/go-git/v5/plumbing/transport"

	"tug/internal/platform/replica"
	"tug/internal/storage"
	"tug/internal/storage/dexstore"
)

const RepoDir = "userbase.git"

var errReplicatorConfig = errors.New("userbase: the replicator needs a store, a state directory, a repo and the three keys")

type Keys struct {
	Identity   *age.X25519Identity
	BreakGlass *age.X25519Recipient
	Signer     *replica.Signer
	DeployKey  []byte
	KnownHosts string
}

func (k Keys) Remote(url string) (replica.Remote, error) {
	if !strings.HasPrefix(url, "ssh://") {
		return replica.Remote{URL: url}, nil
	}
	auth, err := replica.SSHAuth(url, k.DeployKey, k.KnownHosts)
	if err != nil {
		return replica.Remote{URL: url}, err
	}
	return replica.Remote{URL: url, Auth: auth}, nil
}

func (k Keys) Digest(blob []byte) ([32]byte, error) {
	plaintext, err := Open(blob, k.Identity)
	if err != nil {
		return [32]byte{}, err
	}
	doc, err := Decode(plaintext)
	if err != nil {
		return [32]byte{}, err
	}
	return ContentDigest(doc.Snapshot, doc.DexKeys)
}

func OpenLocalRepo(stateDir string) (*git.Repository, error) {
	path := filepath.Join(stateDir, RepoDir)
	repo, err := git.PlainOpen(path)
	if err == nil {
		return repo, nil
	}
	if !errors.Is(err, git.ErrRepositoryNotExists) {
		return nil, err
	}
	repo, err = git.PlainInit(path, true)
	if err != nil {
		return nil, err
	}
	if err := repo.Storer.SetReference(plumbing.NewSymbolicReference(plumbing.HEAD, replica.Branch)); err != nil {
		return nil, err
	}
	return repo, nil
}

type checkpoints struct{ store *storage.Store }

func (c checkpoints) ReadCheckpoint(ctx context.Context) (replica.Checkpoint, error) {
	cp, err := c.store.UserbaseCommit(ctx)
	return replica.Checkpoint{Value: cp.Value, Set: cp.Set}, err
}

func (c checkpoints) SwapCheckpoint(ctx context.Context, expected replica.Checkpoint, next string) error {
	return c.store.CompareAndSwapUserbaseCommit(ctx, storage.Checkpoint{Value: expected.Value, Set: expected.Set}, next)
}

type ReplicatorConfig struct {
	Store             *storage.Store
	StateDir          string
	Repo              string
	Mirror            string
	Keys              Keys
	ExportSigningKeys bool
	Logger            *slog.Logger
	ObservePush       func(at time.Time)
}

type Replicator struct {
	store    *storage.Store
	stateDir string
	repoDir  string
	primary  string
	mirror   string
	keys     Keys
	dex      dexstorage.Storage
	logger   *slog.Logger
	repo     *git.Repository
	sched    *schedule
	lastPush time.Time
	pushed   func(at time.Time)
	inflight sync.WaitGroup

	now        func() time.Time
	after      func(time.Duration) <-chan time.Time
	remoteHead func(ctx context.Context, url string, auth transport.AuthMethod, branch plumbing.ReferenceName) (plumbing.Hash, error)
	planStart  func(ctx context.Context, in replica.StartInput) (replica.StartPlan, error)
	step       func(ctx context.Context, in replica.StepInput) (replica.Result, error)
	prune      func(dir string, keep ...plumbing.Hash) (*git.Repository, error)
}

func NewReplicator(ctx context.Context, cfg ReplicatorConfig) (*Replicator, error) {
	if cfg.Store == nil || cfg.StateDir == "" || cfg.Repo == "" || cfg.Keys.Identity == nil || cfg.Keys.BreakGlass == nil || cfg.Keys.Signer == nil {
		return nil, errReplicatorConfig
	}
	repo, err := OpenLocalRepo(cfg.StateDir)
	if err != nil {
		return nil, err
	}
	var dex dexstorage.Storage
	if cfg.ExportSigningKeys {
		if dex, err = dexstore.Attach(ctx, cfg.Store); err != nil {
			return nil, err
		}
	}
	logger := cfg.Logger
	if logger == nil {
		logger = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	pushed := cfg.ObservePush
	if pushed == nil {
		pushed = func(time.Time) {}
	}
	if cp, err := cfg.Store.UserbaseCommit(ctx); err == nil && cp.Set {
		if c, err := object.GetCommit(repo.Storer, plumbing.NewHash(cp.Value)); err == nil {
			pushed(c.Committer.When)
		}
	}
	return &Replicator{
		store:      cfg.Store,
		stateDir:   cfg.StateDir,
		repoDir:    filepath.Join(cfg.StateDir, RepoDir),
		primary:    cfg.Repo,
		mirror:     cfg.Mirror,
		keys:       cfg.Keys,
		dex:        dex,
		logger:     logger,
		pushed:     pushed,
		repo:       repo,
		sched:      newSchedule(func() float64 { return 2*rand.Float64() - 1 }),
		now:        time.Now,
		after:      time.After,
		remoteHead: replica.RemoteHead,
		planStart:  replica.PlanStart,
		step:       replica.Step,
		prune:      replica.Prune,
	}, nil
}

const maxActionsPerTick = 4

var refusalKinds = []struct {
	err  error
	kind string
}{
	{replica.ErrRewound, "rewound"},
	{replica.ErrForeignChild, "foreign-child"},
	{replica.ErrNotLocalHead, "not-local-head"},
	{replica.ErrDiverged, "diverged"},
	{replica.ErrRemoteEmpty, "remote-empty"},
	{replica.ErrRemoteHasHistory, "remote-has-history"},
	{replica.ErrOwnRootDiffers, "own-root-differs"},
	{replica.ErrCheckpointMissing, "checkpoint-missing"},
}

func refusalKind(err error) string {
	for _, k := range refusalKinds {
		if errors.Is(err, k.err) {
			return k.kind
		}
	}
	return "refused"
}

func shortID(id plumbing.Hash) string {
	if id.IsZero() {
		return ""
	}
	return id.String()[:12]
}

func abandonable[T any](ctx context.Context, wg *sync.WaitGroup, f func(context.Context) (T, error)) (T, error) {
	type outcome struct {
		v   T
		err error
	}
	out := make(chan outcome, 1)
	wg.Add(1)
	go func() {
		defer wg.Done()
		v, err := f(ctx)
		out <- outcome{v, err}
	}()
	select {
	case o := <-out:
		return o.v, o.err
	case <-ctx.Done():
		var zero T
		return zero, ctx.Err()
	}
}

type exportRead struct {
	snapshot storage.UserbaseSnapshot
	dexKeys  json.RawMessage
	digest   [32]byte
}

func (r *Replicator) Run(ctx context.Context) {
	for ctx.Err() == nil {
		r.tick(ctx)
		select {
		case <-ctx.Done():
			return
		case <-r.after(pollInterval):
		}
	}
}

func (r *Replicator) tick(ctx context.Context) {
	read, err := r.read(ctx)
	if err != nil {
		if ctx.Err() == nil {
			r.logger.Warn("userbase export read failed", "err", err)
		}
		return
	}
	r.sched.observe(r.now(), read.digest)
	for i := 0; i < maxActionsPerTick && ctx.Err() == nil; i++ {
		now := r.now()
		if r.sched.next(now) == actWait {
			return
		}
		r.act(ctx, now, read)
	}
}

func (r *Replicator) read(ctx context.Context) (exportRead, error) {
	snapshot, err := r.store.ExportUserbase(ctx)
	if err != nil {
		return exportRead{}, err
	}
	var keys json.RawMessage
	if r.dex != nil {
		keys, err = dexstore.ExportPublicKeys(ctx, r.dex)
		if errors.Is(err, dexstore.ErrNoKeys) {
			keys, err = nil, nil
		}
		if err != nil {
			return exportRead{}, err
		}
	}
	digest, err := ContentDigest(snapshot, keys)
	if err != nil {
		return exportRead{}, err
	}
	return exportRead{snapshot: snapshot, dexKeys: keys, digest: digest}, nil
}

func (r *Replicator) act(ctx context.Context, now time.Time, read exportRead) {
	primary, err := r.keys.Remote(r.primary)
	if err != nil {
		r.fail(now, err)
		return
	}
	head, headErr := abandonable(ctx, &r.inflight, func(ctx context.Context) (plumbing.Hash, error) {
		return r.remoteHead(ctx, primary.URL, primary.Auth, replica.Branch)
	})
	if ctx.Err() != nil {
		return
	}
	if headErr != nil {
		r.fail(now, headErr)
		return
	}
	if refused, refusing := r.sched.currentRefusal(); refusing && head == refused.head {
		r.sched.unchanged(now)
		r.logger.Warn("userbase replication still refused", "kind", refused.kind, "remote", r.primary, "head", shortID(head), "last_push", r.lastPushText())
		return
	}
	if err := replica.DropEmptyRefs(r.repoDir); err != nil {
		r.fail(now, err)
		return
	}
	if !r.sched.isPlanned {
		r.plan(ctx, now, primary, head, read)
		return
	}
	r.stepOnce(ctx, now, primary, head, read)
}

func (r *Replicator) plan(ctx context.Context, now time.Time, primary replica.Remote, head plumbing.Hash, read exportRead) {
	plan, err := abandonable(ctx, &r.inflight, func(ctx context.Context) (replica.StartPlan, error) {
		return r.planStart(ctx, replica.StartInput{
			Repo:        r.repo,
			Dir:         r.stateDir,
			Checkpoints: checkpoints{r.store},
			Primary:     primary,
			Digest:      r.keys.Digest,
			Current:     read.digest,
		})
	})
	if ctx.Err() != nil {
		return
	}
	if err != nil {
		r.settleError(ctx, now, head, err)
		return
	}
	r.sched.planned(now, plan)
	r.logger.Info("userbase replication planned", "remote", r.primary, "step_now", plan.StepNow, "pending", plan.Pending, "daily_due", plan.DailyDue, "checkpoint_unreadable", plan.CheckpointUnreadable)
	if plan.CheckpointUnreadable {
		r.logger.Warn("userbase checkpoint export cannot be read with the configured keys; committing a new export", "commit", r.checkpointID(ctx))
	}
}

func (r *Replicator) checkpointID(ctx context.Context) string {
	cp, err := r.store.UserbaseCommit(ctx)
	if err != nil || !cp.Set {
		return ""
	}
	return shortID(plumbing.NewHash(cp.Value))
}

func (r *Replicator) stepOnce(ctx context.Context, now time.Time, primary replica.Remote, head plumbing.Hash, read exportRead) {
	plaintext, err := Encode(read.snapshot, now, read.dexKeys)
	if err != nil {
		r.fail(now, err)
		return
	}
	blob, err := Seal(plaintext, r.keys.Identity.Recipient(), r.keys.BreakGlass)
	if err != nil {
		r.fail(now, err)
		return
	}
	in := replica.StepInput{
		Repo:        r.repo,
		Dir:         r.stateDir,
		Checkpoints: checkpoints{r.store},
		Primary:     primary,
		Signer:      r.keys.Signer,
		Blob:        blob,
		When:        now,
		Digest:      r.keys.Digest,
		Current:     read.digest,
	}
	var mirrorErr error
	if r.mirror != "" {
		mirror, err := r.keys.Remote(r.mirror)
		if err != nil {
			mirrorErr = err
		} else {
			in.Mirror = &mirror
		}
	}
	res, err := abandonable(ctx, &r.inflight, func(ctx context.Context) (replica.Result, error) {
		return r.step(ctx, in)
	})
	if ctx.Err() != nil {
		return
	}
	if err != nil {
		r.settleError(ctx, now, head, err)
		return
	}
	if mirrorErr != nil && res.Outcome == replica.Pushed && res.MirrorErr == nil {
		res.MirrorErr = mirrorErr
	}
	r.succeeded(now, res, read)
}

func (r *Replicator) settleError(ctx context.Context, now time.Time, head plumbing.Hash, err error) {
	if !errors.Is(err, replica.ErrRefused) {
		r.fail(now, err)
		return
	}
	kind := refusalKind(err)
	audit := r.sched.refused(now, refusal{kind: kind, head: head})
	r.logger.Error("userbase replication refused", "kind", kind, "remote", r.primary, "head", shortID(head), "err", err, "retry_after", recheckInterval, "last_push", r.lastPushText())
	if !audit {
		return
	}
	detail, err := json.Marshal(map[string]string{"kind": kind, "remote": r.primary, "head": shortID(head)})
	if err == nil {
		_, err = r.store.AppendAudit(ctx, storage.AppendAuditInput{Actor: "system", Action: "userbase.refused", Subject: "userbase", Detail: string(detail)})
	}
	if err != nil {
		r.logger.Error("userbase refusal audit failed", "kind", kind, "err", err)
	}
}

func (r *Replicator) fail(now time.Time, err error) {
	r.sched.failed(now, err)
	r.logger.Warn("userbase replication attempt failed", "remote", r.primary, "err", err, "retry_after", r.sched.retryAt().Sub(now).Round(time.Second), "last_push", r.lastPushText())
}

func (r *Replicator) succeeded(now time.Time, res replica.Result, read exportRead) {
	committed := read.digest
	if res.Outcome != replica.Pushed {
		committed = r.committedDigest(res.Commit)
	}
	commitTime := now
	if c, err := object.GetCommit(r.repo.Storer, res.Commit); err == nil {
		commitTime = c.Committer.When
	}
	r.sched.succeeded(now, read.digest, committed, commitTime)
	if res.Outcome != replica.Pushed {
		r.logger.Info("userbase push", "remote", r.primary, "role", "repo", "result", "ok", "outcome", "adopted", "commit", res.Commit.String(), "last_push", r.lastPushText())
		return
	}
	r.lastPush = now
	r.pushed(now)
	r.logger.Info("userbase export committed", "commit", res.Commit.String(), "remote", r.primary)
	r.logger.Info("userbase push", "remote", r.primary, "role", "repo", "result", "ok", "outcome", "pushed", "commit", res.Commit.String(), "last_push", r.lastPushText())
	if r.mirror != "" {
		if res.MirrorErr != nil {
			r.logger.Warn("userbase push", "remote", r.mirror, "role", "mirror", "result", "failed", "reseeded", res.MirrorReseeded, "err", res.MirrorErr, "last_push", r.lastPushText())
		} else {
			r.logger.Info("userbase push", "remote", r.mirror, "role", "mirror", "result", "ok", "reseeded", res.MirrorReseeded, "commit", res.Commit.String(), "last_push", r.lastPushText())
		}
	}
	if r.sched.pruneDue(now, res) {
		r.pruneLocal(now, res.Commit)
	}
}

func (r *Replicator) committedDigest(id plumbing.Hash) [32]byte {
	blob, err := replica.ReadFile(r.repo.Storer, id)
	if err != nil {
		return [32]byte{}
	}
	digest, err := r.keys.Digest(blob)
	if err != nil {
		return [32]byte{}
	}
	return digest
}

func (r *Replicator) pruneLocal(now time.Time, keep plumbing.Hash) {
	r.sched.pruned(now)
	repo, err := r.prune(r.repoDir, keep)
	if err == nil {
		r.repo = repo
		r.logger.Info("userbase local repository pruned", "kept", keep.String())
		return
	}
	r.logger.Warn("userbase local repository prune failed", "err", err)
	if repo, err = git.PlainOpen(r.repoDir); err != nil {
		r.logger.Error("userbase local repository reopen failed", "err", err)
		return
	}
	r.repo = repo
}

func (r *Replicator) lastPushText() string {
	if r.lastPush.IsZero() {
		return "never"
	}
	return r.lastPush.UTC().Format(time.RFC3339)
}