// Copyright 2021 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package actions

import (
	"context"
	"errors"

	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
	"github.com/dolthub/dolt/go/libraries/doltcore/env"
	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
	"github.com/dolthub/dolt/go/libraries/utils/set"
	"github.com/dolthub/dolt/go/store/hash"
)

func CheckoutAllTables(ctx context.Context, roots doltdb.Roots) (doltdb.Roots, error) {
	tbls, err := doltdb.UnionTableNames(ctx, roots.Working, roots.Staged, roots.Head)
	if err != nil {
		return doltdb.Roots{}, err
	}

	return MoveTablesFromHeadToWorking(ctx, roots, tbls)
}

// CheckoutTables takes in a set of tables and docs and checks them out to another branch.
func CheckoutTables(ctx context.Context, roots doltdb.Roots, tables []string) (doltdb.Roots, error) {
	return MoveTablesFromHeadToWorking(ctx, roots, tables)
}

// MoveTablesFromHeadToWorking replaces the tables named from the given head to the given working root, overwriting any
// working changes, and returns the new resulting roots
func MoveTablesFromHeadToWorking(ctx context.Context, roots doltdb.Roots, tbls []string) (doltdb.Roots, error) {
	var unknownTbls []string
	for _, tblName := range tbls {
		tbl, ok, err := roots.Staged.GetTable(ctx, tblName)
		if err != nil {
			return doltdb.Roots{}, err
		}
		fkc, err := roots.Staged.GetForeignKeyCollection(ctx)
		if err != nil {
			return doltdb.Roots{}, err
		}

		if !ok {
			tbl, ok, err = roots.Head.GetTable(ctx, tblName)
			if err != nil {
				return doltdb.Roots{}, err
			}

			fkc, err = roots.Head.GetForeignKeyCollection(ctx)
			if err != nil {
				return doltdb.Roots{}, err
			}

			if !ok {
				unknownTbls = append(unknownTbls, tblName)
				continue
			}
		}

		roots.Working, err = roots.Working.PutTable(ctx, tblName, tbl)
		if err != nil {
			return doltdb.Roots{}, err
		}

		roots.Working, err = roots.Working.PutForeignKeyCollection(ctx, fkc)
		if err != nil {
			return doltdb.Roots{}, err
		}
	}

	if len(unknownTbls) > 0 {
		// Return table not exist error before RemoveTables, which fails silently if the table is not on the root.
		err := validateTablesExist(ctx, roots.Working, unknownTbls)
		if err != nil {
			return doltdb.Roots{}, err
		}

		roots.Working, err = roots.Working.RemoveTables(ctx, false, false, unknownTbls...)

		if err != nil {
			return doltdb.Roots{}, err
		}
	}

	return roots, nil
}

// rootsForBranch returns the roots needed for a branch checkout. |roots.Head| should be the pre-checkout head. The
// returned roots struct has |Head| set to |branchRoot|.
func rootsForBranch(ctx context.Context, roots doltdb.Roots, branchRoot *doltdb.RootValue, force bool) (doltdb.Roots, error) {
	conflicts := set.NewStrSet([]string{})
	if roots.Head == nil {
		roots.Working = branchRoot
		roots.Staged = branchRoot
		roots.Head = branchRoot
		return roots, nil
	}

	wrkTblHashes, err := moveModifiedTables(ctx, roots.Head, branchRoot, roots.Working, conflicts, force)
	if err != nil {
		return doltdb.Roots{}, err
	}

	stgTblHashes, err := moveModifiedTables(ctx, roots.Head, branchRoot, roots.Staged, conflicts, force)
	if err != nil {
		return doltdb.Roots{}, err
	}

	if conflicts.Size() > 0 {
		return doltdb.Roots{}, ErrCheckoutWouldOverwrite{conflicts.AsSlice()}
	}

	workingForeignKeys, err := moveForeignKeys(ctx, roots.Head, branchRoot, roots.Working, force)
	if err != nil {
		return doltdb.Roots{}, err
	}

	stagedForeignKeys, err := moveForeignKeys(ctx, roots.Head, branchRoot, roots.Staged, force)
	if err != nil {
		return doltdb.Roots{}, err
	}

	roots.Working, err = writeTableHashes(ctx, branchRoot, wrkTblHashes)
	if err != nil {
		return doltdb.Roots{}, err
	}

	roots.Staged, err = writeTableHashes(ctx, branchRoot, stgTblHashes)
	if err != nil {
		return doltdb.Roots{}, err
	}

	roots.Working, err = roots.Working.PutForeignKeyCollection(ctx, workingForeignKeys)
	if err != nil {
		return doltdb.Roots{}, err
	}

	roots.Staged, err = roots.Staged.PutForeignKeyCollection(ctx, stagedForeignKeys)
	if err != nil {
		return doltdb.Roots{}, err
	}

	roots.Head = branchRoot
	return roots, nil
}

func CheckoutBranch(ctx context.Context, dEnv *env.DoltEnv, brName string, force bool) error {
	branchRef := ref.NewBranchRef(brName)
	initialHeadRef := dEnv.RepoStateReader().CWBHeadRef()

	db := dEnv.DoltDB
	hasRef, err := db.HasRef(ctx, branchRef)
	if err != nil {
		return err
	}
	if !hasRef {
		return doltdb.ErrBranchNotFound
	}

	if ref.Equals(dEnv.RepoStateReader().CWBHeadRef(), branchRef) {
		return doltdb.ErrAlreadyOnBranch
	}

	branchHead, err := branchHeadRoot(ctx, db, brName)
	if err != nil {
		return err
	}

	workingSetExists := true
	initialWs, err := dEnv.WorkingSet(ctx)
	if err == doltdb.ErrWorkingSetNotFound {
		// ignore, but don't reset the working set
		workingSetExists = false
	} else if err != nil {
		return err
	}

	if !force {
		if checkoutWouldStompWorkingSetChanges(ctx, dEnv, branchRef) {
			return ErrWorkingSetsOnBothBranches
		}
	}

	initialRoots, err := dEnv.Roots(ctx)

	// roots will be empty/nil if the working set is not set (working set is not set if the current branch was deleted)
	if errors.Is(err, doltdb.ErrBranchNotFound) || errors.Is(err, doltdb.ErrWorkingSetNotFound) {
		workingSetExists = false
	} else if err != nil {
		return err
	}

	hasChanges := false
	if workingSetExists {
		hasChanges, _, _, err = rootHasUncommittedChanges(initialRoots)
		if err != nil {
			return err
		}
	}

	// Only if the current working set has uncommitted changes do we carry them forward to the branch being checked out.
	// If this is the case, then the destination branch must *not* have any uncommitted changes, as checked by
	// checkoutWouldStompWorkingSetChanges
	if hasChanges {
		err = transferWorkingChanges(ctx, dEnv, initialRoots, branchHead, branchRef, force)
		if err != nil {
			return err
		}
	} else {
		err = dEnv.RepoStateWriter().SetCWBHeadRef(ctx, ref.MarshalableRef{Ref: branchRef})
		if err != nil {
			return err
		}
	}

	if workingSetExists && hasChanges {
		err = cleanOldWorkingSet(ctx, dEnv, initialRoots, initialHeadRef, initialWs)
		if err != nil {
			return err
		}
	}

	return nil
}

func transferWorkingChanges(
	ctx context.Context,
	dEnv *env.DoltEnv,
	initialRoots doltdb.Roots,
	branchHead *doltdb.RootValue,
	branchRef ref.BranchRef,
	force bool,
) error {
	newRoots, err := rootsForBranch(ctx, initialRoots, branchHead, force)
	if err != nil {
		return err
	}

	// important to not update the checked out branch until after we have done the error checking above, otherwise we
	// potentially leave the client in a bad state
	err = dEnv.RepoStateWriter().SetCWBHeadRef(ctx, ref.MarshalableRef{Ref: branchRef})
	if err != nil {
		return err
	}

	ws, err := dEnv.WorkingSet(ctx)

	// For backwards compatibility we support the branch not having a working set, but generally speaking it already
	// should have one
	if err == doltdb.ErrWorkingSetNotFound {
		wsRef, err := ref.WorkingSetRefForHead(branchRef)
		if err != nil {
			return err
		}
		ws = doltdb.EmptyWorkingSet(wsRef)
	} else if err != nil {
		return err
	}

	err = dEnv.UpdateWorkingSet(ctx, ws.WithWorkingRoot(newRoots.Working).WithStagedRoot(newRoots.Staged))
	if err != nil {
		return err
	}

	return nil
}

// cleanOldWorkingSet resets the source branch's working set to the branch head, leaving the source branch unchanged
func cleanOldWorkingSet(
	ctx context.Context,
	dEnv *env.DoltEnv,
	initialRoots doltdb.Roots,
	initialHeadRef ref.DoltRef,
	initialWs *doltdb.WorkingSet,
) error {
	// reset the source branch's working set to the branch head, leaving the source branch unchanged
	err := ResetHard(ctx, dEnv, "", initialRoots, initialHeadRef, initialWs)
	if err != nil {
		return err
	}

	// Annoyingly, after the ResetHard above we need to get all the roots again, because the working set has changed
	cm, err := dEnv.DoltDB.ResolveCommitRef(ctx, initialHeadRef)
	if err != nil {
		return err
	}

	headRoot, err := cm.ResolveRootValue(ctx)
	if err != nil {
		return err
	}

	workingSet, err := dEnv.DoltDB.ResolveWorkingSet(ctx, initialWs.Ref())
	if err != nil {
		return err
	}

	resetRoots := doltdb.Roots{
		Head:    headRoot,
		Working: workingSet.WorkingRoot(),
		Staged:  workingSet.StagedRoot(),
	}

	// we also have to do a clean, because we the ResetHard won't touch any new tables (tables only in the working set)
	newRoots, err := CleanUntracked(ctx, resetRoots, []string{}, false, true)
	if err != nil {
		return err
	}

	h, err := workingSet.HashOf()
	if err != nil {
		return err
	}

	err = dEnv.DoltDB.UpdateWorkingSet(
		ctx,
		initialWs.Ref(),
		initialWs.WithWorkingRoot(newRoots.Working).WithStagedRoot(newRoots.Staged).ClearMerge(),
		h,
		dEnv.NewWorkingSetMeta("reset hard"),
	)
	if err != nil {
		return err
	}
	return nil
}

// branchHeadRoot returns the root value at the branch head with the name given
func branchHeadRoot(ctx context.Context, db *doltdb.DoltDB, brName string) (*doltdb.RootValue, error) {
	cs, err := doltdb.NewCommitSpec(brName)
	if err != nil {
		return nil, doltdb.RootValueUnreadable{RootType: doltdb.HeadRoot, Cause: err}
	}

	cm, err := db.Resolve(ctx, cs, nil)
	if err != nil {
		return nil, doltdb.RootValueUnreadable{RootType: doltdb.HeadRoot, Cause: err}
	}

	branchRoot, err := cm.GetRootValue(ctx)
	if err != nil {
		return nil, err
	}
	return branchRoot, nil
}

// moveModifiedTables handles working set changes during a branch change.
// When moving between branches, changes in the working set should travel with you.
// Working set changes cannot be moved if the table differs between the old and new head,
// in this case, we throw a conflict and error (as per Git).
func moveModifiedTables(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.RootValue, conflicts *set.StrSet, force bool) (map[string]hash.Hash, error) {
	resultMap := make(map[string]hash.Hash)
	tblNames, err := newRoot.GetTableNames(ctx)
	if err != nil {
		return nil, err
	}

	for _, tblName := range tblNames {
		oldHash, _, err := oldRoot.GetTableHash(ctx, tblName)
		if err != nil {
			return nil, err
		}

		newHash, _, err := newRoot.GetTableHash(ctx, tblName)
		if err != nil {
			return nil, err
		}

		changedHash, _, err := changedRoot.GetTableHash(ctx, tblName)
		if err != nil {
			return nil, err
		}

		if oldHash == changedHash {
			resultMap[tblName] = newHash
		} else if oldHash == newHash {
			resultMap[tblName] = changedHash
		} else if force {
			resultMap[tblName] = newHash
		} else {
			conflicts.Add(tblName)
		}
	}

	tblNames, err = changedRoot.GetTableNames(ctx)
	if err != nil {
		return nil, err
	}

	for _, tblName := range tblNames {
		if _, exists := resultMap[tblName]; !exists {
			oldHash, _, err := oldRoot.GetTableHash(ctx, tblName)
			if err != nil {
				return nil, err
			}

			changedHash, _, err := changedRoot.GetTableHash(ctx, tblName)
			if err != nil {
				return nil, err
			}

			if oldHash == emptyHash {
				resultMap[tblName] = changedHash
			} else if force {
				resultMap[tblName] = oldHash
			} else if oldHash != changedHash {
				conflicts.Add(tblName)
			}
		}
	}

	return resultMap, nil
}

// moveForeignKeys returns the foreign key collection that should be used for the new working set.
func moveForeignKeys(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.RootValue, force bool) (*doltdb.ForeignKeyCollection, error) {
	oldFks, err := oldRoot.GetForeignKeyCollection(ctx)
	if err != nil {
		return nil, err
	}

	newFks, err := newRoot.GetForeignKeyCollection(ctx)
	if err != nil {
		return nil, err
	}

	changedFks, err := changedRoot.GetForeignKeyCollection(ctx)
	if err != nil {
		return nil, err
	}

	oldHash, err := oldFks.HashOf(ctx, oldRoot.VRW())
	if err != nil {
		return nil, err
	}

	newHash, err := newFks.HashOf(ctx, newRoot.VRW())
	if err != nil {
		return nil, err
	}

	changedHash, err := changedFks.HashOf(ctx, changedRoot.VRW())
	if err != nil {
		return nil, err
	}

	if oldHash == changedHash {
		return newFks, nil
	} else if oldHash == newHash {
		return changedFks, nil
	} else {
		// Both roots have modified the foreign keys. We need to do more work to merge them together into a new foreign
		// key collection.
		return mergeForeignKeyChanges(ctx, oldFks, newRoot, newFks, changedRoot, changedFks, force)
	}
}

// mergeForeignKeyChanges merges the foreign key changes from the old and changed roots into a new foreign key
// collection, or returns an error if the changes are incompatible. Changes are incompatible if the changed root
// and new root both altered foreign keys on the same table.
func mergeForeignKeyChanges(
	ctx context.Context,
	oldFks *doltdb.ForeignKeyCollection,
	newRoot *doltdb.RootValue,
	newFks *doltdb.ForeignKeyCollection,
	changedRoot *doltdb.RootValue,
	changedFks *doltdb.ForeignKeyCollection,
	force bool,
) (*doltdb.ForeignKeyCollection, error) {
	fksByTable := make(map[string][]doltdb.ForeignKey)

	conflicts := set.NewEmptyStrSet()
	tblNames, err := newRoot.GetTableNames(ctx)
	if err != nil {
		return nil, err
	}

	for _, tblName := range tblNames {
		oldFksForTable, _ := oldFks.KeysForTable(tblName)
		newFksForTable, _ := newFks.KeysForTable(tblName)
		changedFksForTable, _ := changedFks.KeysForTable(tblName)

		oldHash := doltdb.CombinedHash(oldFksForTable)
		newHash := doltdb.CombinedHash(newFksForTable)
		changedHash := doltdb.CombinedHash(changedFksForTable)

		if oldHash == changedHash {
			fksByTable[tblName] = append(fksByTable[tblName], newFksForTable...)
		} else if oldHash == newHash {
			fksByTable[tblName] = append(fksByTable[tblName], changedFksForTable...)
		} else if force {
			fksByTable[tblName] = append(fksByTable[tblName], newFksForTable...)
		} else {
			conflicts.Add(tblName)
		}
	}

	tblNames, err = changedRoot.GetTableNames(ctx)
	if err != nil {
		return nil, err
	}

	for _, tblName := range tblNames {
		if _, exists := fksByTable[tblName]; !exists {
			oldKeys, _ := oldFks.KeysForTable(tblName)
			oldHash := doltdb.CombinedHash(oldKeys)

			changedKeys, _ := changedFks.KeysForTable(tblName)
			changedHash := doltdb.CombinedHash(changedKeys)

			if oldHash == emptyHash {
				fksByTable[tblName] = append(fksByTable[tblName], changedKeys...)
			} else if force {
				fksByTable[tblName] = append(fksByTable[tblName], oldKeys...)
			} else if oldHash != changedHash {
				conflicts.Add(tblName)
			}
		}
	}

	if conflicts.Size() > 0 {
		return nil, ErrCheckoutWouldOverwrite{conflicts.AsSlice()}
	}

	fks := make([]doltdb.ForeignKey, 0)
	for _, v := range fksByTable {
		fks = append(fks, v...)
	}

	return doltdb.NewForeignKeyCollection(fks...)
}

// writeTableHashes writes new table hash values for the root given and returns it.
// This is an inexpensive and convenient way of replacing all the tables at once.
func writeTableHashes(ctx context.Context, head *doltdb.RootValue, tblHashes map[string]hash.Hash) (*doltdb.RootValue, error) {
	names, err := head.GetTableNames(ctx)
	if err != nil {
		return nil, err
	}

	var toDrop []string
	for _, name := range names {
		if _, ok := tblHashes[name]; !ok {
			toDrop = append(toDrop, name)
		}
	}

	head, err = head.RemoveTables(ctx, false, false, toDrop...)
	if err != nil {
		return nil, err
	}

	for k, v := range tblHashes {
		if v == emptyHash {
			continue
		}

		head, err = head.SetTableHash(ctx, k, v)
		if err != nil {
			return nil, err
		}
	}

	return head, nil
}

// checkoutWouldStompWorkingSetChanges checks that the current working set is "compatible" with the dest working set.
// This means that if both working sets are present (ie there are changes on both source and dest branches),
// we check if the changes are identical before allowing a clobbering checkout.
// Working set errors are ignored by this function, because they are properly handled elsewhere.
func checkoutWouldStompWorkingSetChanges(ctx context.Context, dEnv *env.DoltEnv, branchRef ref.BranchRef) bool {
	sourceRoots, err := dEnv.Roots(ctx)
	if err != nil {
		return false
	}

	destRoots, err := dEnv.DoltDB.ResolveBranchRoots(ctx, branchRef)
	if err != nil {
		return false
	}

	sourceHasChanges, sourceWorkingHash, sourceStagedHash, err := rootHasUncommittedChanges(sourceRoots)
	if err != nil {
		return false
	}

	destHasChanges, destWorkingHash, destStagedHash, err := rootHasUncommittedChanges(destRoots)
	if err != nil {
		return false
	}

	// This is a stomping checkout operation if both the source and dest have uncommitted changes, and they're not the
	// same uncommitted changes
	return sourceHasChanges && destHasChanges && (sourceWorkingHash != destWorkingHash || sourceStagedHash != destStagedHash)
}

// rootHasUncommittedChanges returns whether the roots given have uncommitted changes, and the hashes of the working and staged roots
func rootHasUncommittedChanges(roots doltdb.Roots) (hasChanges bool, workingHash hash.Hash, stagedHash hash.Hash, err error) {
	headHash, err := roots.Head.HashOf()
	if err != nil {
		return false, hash.Hash{}, hash.Hash{}, err
	}

	workingHash, err = roots.Working.HashOf()
	if err != nil {
		return false, hash.Hash{}, hash.Hash{}, err
	}

	stagedHash, err = roots.Staged.HashOf()
	if err != nil {
		return false, hash.Hash{}, hash.Hash{}, err
	}

	hasChanges = workingHash != stagedHash || stagedHash != headHash
	return hasChanges, workingHash, stagedHash, nil
}
