package userbase

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"time"

	"filippo.io/age"
	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
	"golang.org/x/crypto/ssh"

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

const (
	StageFetch   = "fetch"
	StageVerify  = "verify"
	StageDecrypt = "decrypt"
	StageDecode  = "decode"
	StageImport  = "import"
)

type StageError struct {
	Stage string
	Err   error
}

func (e *StageError) Error() string {
	return e.Stage + ": " + e.Err.Error()
}

func (e *StageError) Unwrap() error {
	return e.Err
}

type RestoreConfig struct {
	Store          *storage.Store
	StateDir       string
	Remote         replica.Remote
	Identity       *age.X25519Identity
	AllowedKeys    []ssh.PublicKey
	Commit         plumbing.Hash
	DryRun         bool
	WithoutDexKeys bool
	DexValidFor    time.Duration
	DexRotation    time.Duration
	Now            func() time.Time
}

type RestoreSummary struct {
	Head        plumbing.Hash
	Commit      plumbing.Hash
	ExportedAt  time.Time
	Users       int
	Groups      int
	Memberships int
	GoogleLinks int
	Invitations int
	DexKeys     bool
}

var verifyErrors = []error{
	replica.ErrChainUnverified,
	replica.ErrMergeCommit,
	replica.ErrNoAllowedKeys,
	replica.ErrNotOnBranch,
	replica.ErrNotExport,
}

func staged(stage string, err error) error {
	var se *StageError
	if errors.As(err, &se) {
		return err
	}
	return &StageError{Stage: stage, Err: err}
}

func replicaStage(err error) error {
	for _, v := range verifyErrors {
		if errors.Is(err, v) {
			return staged(StageVerify, err)
		}
	}
	return staged(StageFetch, err)
}

func openDocument(blob []byte, identity *age.X25519Identity) (Document, error) {
	plaintext, err := Open(blob, identity)
	if err != nil {
		return Document{}, staged(StageDecrypt, err)
	}
	doc, err := Decode(plaintext)
	if err != nil {
		return Document{}, staged(StageDecode, err)
	}
	return doc, nil
}

func summarize(doc Document) RestoreSummary {
	s := doc.Snapshot
	return RestoreSummary{
		ExportedAt:  doc.ExportedAt,
		Users:       len(s.Users),
		Groups:      len(s.Groups),
		Memberships: len(s.Memberships),
		GoogleLinks: len(s.GoogleLinks),
		Invitations: len(s.Invitations),
		DexKeys:     len(doc.DexKeys) > 0,
	}
}

func (cfg RestoreConfig) now() time.Time {
	if cfg.Now != nil {
		return cfg.Now()
	}
	return time.Now()
}

func (cfg RestoreConfig) importDocument(ctx context.Context, doc Document, restored, checkpoint plumbing.Hash) error {
	if len(doc.DexKeys) > 0 && !cfg.WithoutDexKeys {
		dex, err := dexstore.Attach(ctx, cfg.Store)
		if err != nil {
			return staged(StageImport, err)
		}
		now := cfg.now()
		if err := dexstore.ImportPublicKeys(ctx, dex, doc.DexKeys, now, cfg.DexValidFor, now.Add(cfg.DexRotation)); err != nil {
			return staged(StageImport, err)
		}
	}
	if err := cfg.Store.ImportUserbase(ctx, doc.Snapshot, restored.String(), checkpoint.String()); err != nil {
		return staged(StageImport, err)
	}
	if err := os.Remove(filepath.Join(cfg.StateDir, invite.FileName)); err != nil && !errors.Is(err, fs.ErrNotExist) {
		return staged(StageImport, err)
	}
	return nil
}

const headCountersFormat = "the head %s, whose uid and gid counters a restore of an older commit keeps: %w"

func openHead(blob []byte, head plumbing.Hash, identity *age.X25519Identity) (Document, error) {
	plaintext, err := Open(blob, identity)
	if err != nil {
		return Document{}, staged(StageDecrypt, fmt.Errorf(headCountersFormat, head, err))
	}
	current, err := Decode(plaintext)
	if err != nil {
		return Document{}, staged(StageDecode, fmt.Errorf(headCountersFormat, head, err))
	}
	return current, nil
}

func keepLargerCounters(repo *git.Repository, head plumbing.Hash, identity *age.X25519Identity, s storage.UserbaseSnapshot) (storage.UserbaseSnapshot, error) {
	blob, err := replica.ReadFile(repo.Storer, head)
	if err != nil {
		return s, staged(StageImport, err)
	}
	current, err := openHead(blob, head, identity)
	if err != nil {
		return s, err
	}
	s.NextUIDNumber = max(s.NextUIDNumber, current.Snapshot.NextUIDNumber)
	s.NextGIDNumber = max(s.NextGIDNumber, current.Snapshot.NextGIDNumber)
	return s, nil
}

func Restore(ctx context.Context, cfg RestoreConfig) (RestoreSummary, error) {
	if err := cfg.Store.UserbaseImportable(ctx); err != nil {
		return RestoreSummary{}, staged(StageImport, err)
	}
	in := replica.RestoreInput{
		Dir:    cfg.StateDir,
		Remote: cfg.Remote,
		Commit: cfg.Commit,
		Keys:   cfg.AllowedKeys,
		DryRun: cfg.DryRun,
	}
	var summary RestoreSummary
	if !cfg.DryRun {
		repo, err := OpenLocalRepo(cfg.StateDir)
		if err != nil {
			return RestoreSummary{}, staged(StageImport, err)
		}
		in.Repo = repo
		in.Import = func(ctx context.Context, blob []byte, restored, checkpoint plumbing.Hash) error {
			doc, err := openDocument(blob, cfg.Identity)
			if err != nil {
				return err
			}
			if restored != checkpoint {
				if doc.Snapshot, err = keepLargerCounters(repo, checkpoint, cfg.Identity, doc.Snapshot); err != nil {
					return err
				}
			}
			if err := cfg.importDocument(ctx, doc, restored, checkpoint); err != nil {
				return err
			}
			summary = summarize(doc)
			return nil
		}
	}
	res, err := replica.Restore(ctx, in)
	if err != nil {
		return RestoreSummary{}, replicaStage(err)
	}
	if cfg.DryRun {
		doc, err := openDocument(res.Blob, cfg.Identity)
		if err != nil {
			return RestoreSummary{}, err
		}
		if res.Commit != res.Head {
			if _, err := openHead(res.HeadBlob, res.Head, cfg.Identity); err != nil {
				return RestoreSummary{}, err
			}
		}
		summary = summarize(doc)
	}
	summary.Head, summary.Commit = res.Head, res.Commit
	return summary, nil
}

type FirstBootOutcome int

const (
	FirstBootNotEmpty FirstBootOutcome = iota + 1
	FirstBootEmptyRemote
	FirstBootImported
)

func FirstBoot(ctx context.Context, cfg RestoreConfig) (FirstBootOutcome, RestoreSummary, error) {
	err := cfg.Store.UserbaseImportable(ctx)
	if errors.Is(err, storage.ErrUserbaseNotEmpty) {
		return FirstBootNotEmpty, RestoreSummary{}, nil
	}
	if err != nil {
		return 0, RestoreSummary{}, staged(StageImport, err)
	}
	cfg.Commit = plumbing.ZeroHash
	cfg.DryRun = false
	summary, err := Restore(ctx, cfg)
	if errors.Is(err, replica.ErrNothingToRestore) {
		return FirstBootEmptyRemote, RestoreSummary{}, nil
	}
	if err != nil {
		return 0, RestoreSummary{}, err
	}
	return FirstBootImported, summary, nil
}