package kdbx

import (
	"errors"
	"sort"
	"strconv"
	"time"
)

// Merge folds source into target, following KeePass's own synchronization
// rules: the union of both databases, with the most recently modified
// version of each object winning, and the version that loses kept in the
// entry's history. Nothing is discarded silently.
//
// Target is modified in place. Source is only read.
//
// The ordering of objects inside a group is not merged; new objects are
// appended. KeePass does not reconcile order either.
func Merge(target, source *Database) (MergeResult, error) {
	var res MergeResult
	if target == nil || source == nil || target.Root == nil || source.Root == nil {
		return res, errors.New("kdbx: merge needs two parsed databases")
	}
	ti := newIndex(target)
	si := newIndex(source)

	mergeMeta(target, source, &res)
	mergeGroups(target, source, ti, si, &res)
	mergeEntries(target, source, ti, si, &res)
	applyDeletions(target, ti, si, &res)
	return res, nil
}

// MergeResult reports what the merge did, for the server's log and for the
// client's status line.
type MergeResult struct {
	GroupsAdded    int
	GroupsUpdated  int
	EntriesAdded   int
	EntriesUpdated int
	HistoryAdded   int
	Relocated      int
	Deleted        int
	Resurrected    int
}

// Changed reports whether the merge altered the target at all.
func (r MergeResult) Changed() bool {
	return r.GroupsAdded+r.GroupsUpdated+r.EntriesAdded+r.EntriesUpdated+
		r.HistoryAdded+r.Relocated+r.Deleted+r.Resurrected > 0
}

func (r MergeResult) String() string {
	return "added " + strconv.Itoa(r.EntriesAdded) + " entries, updated " +
		strconv.Itoa(r.EntriesUpdated) + ", added " + strconv.Itoa(r.GroupsAdded) +
		" groups, moved " + strconv.Itoa(r.Relocated) + ", deleted " +
		strconv.Itoa(r.Deleted)
}

// ---- indexing ----

// objRef is one group or entry together with the group holding it.
type objRef struct {
	node   *Node
	parent *Node // the Group node this object is a child of
}

// index locates every group and entry by UUID. History entries are not
// indexed: they share their parent entry's UUID.
type index struct {
	groups     map[string]*objRef
	entries    map[string]*objRef
	groupOrder []string // document order, parents before children
	entryOrder []string
	deleted    map[string]time.Time
	deletedOf  *Node // the DeletedObjects element
}

func newIndex(db *Database) *index {
	ix := &index{
		groups:    map[string]*objRef{},
		entries:   map[string]*objRef{},
		deleted:   map[string]time.Time{},
		deletedOf: db.Root.Path("Root", "DeletedObjects"),
	}
	root := db.Root.Find("Root")
	if root == nil {
		return ix
	}
	for _, g := range root.FindAll("Group") {
		ix.addGroup(g, nil)
	}
	for _, d := range ix.deletedOf.FindAll("DeletedObject") {
		if t, ok := DecodeTime(d.Value("DeletionTime")); ok {
			uuid := d.Value("UUID")
			if prev, seen := ix.deleted[uuid]; !seen || t.After(prev) {
				ix.deleted[uuid] = t
			}
		}
	}
	return ix
}

func (ix *index) addGroup(g, parent *Node) {
	uuid := g.Value("UUID")
	if uuid != "" {
		if _, dup := ix.groups[uuid]; !dup {
			ix.groups[uuid] = &objRef{node: g, parent: parent}
			ix.groupOrder = append(ix.groupOrder, uuid)
		}
	}
	for _, e := range g.FindAll("Entry") {
		eu := e.Value("UUID")
		if eu == "" {
			continue
		}
		if _, dup := ix.entries[eu]; !dup {
			ix.entries[eu] = &objRef{node: e, parent: g}
			ix.entryOrder = append(ix.entryOrder, eu)
		}
	}
	for _, sub := range g.FindAll("Group") {
		ix.addGroup(sub, g)
	}
}

// lastModified is an object's modification timestamp; a missing one sorts
// oldest, so a well-formed object always wins over a malformed one.
func lastModified(n *Node) time.Time {
	t, _ := DecodeTime(n.Value("Times", "LastModificationTime"))
	return t
}

func locationChanged(n *Node) time.Time {
	t, _ := DecodeTime(n.Value("Times", "LocationChanged"))
	return t
}

// ---- groups ----

func mergeGroups(target, source *Database, ti, si *index, res *MergeResult) {
	for _, uuid := range si.groupOrder {
		src := si.groups[uuid]
		dst, exists := ti.groups[uuid]

		if !exists {
			if deletedLaterThan(ti.deleted[uuid], lastModified(src.node)) {
				continue // target deleted this group after the source touched it
			}
			parent := targetParent(target, ti, si, src)
			shell := groupShell(src.node)
			parent.Nodes = append(parent.Nodes, shell)
			ti.groups[uuid] = &objRef{node: shell, parent: parent}
			res.GroupsAdded++
			continue
		}

		if lastModified(src.node).After(lastModified(dst.node)) {
			copyGroupFields(dst.node, src.node)
			res.GroupsUpdated++
		}
		relocate(target, ti, si, src, dst, res)
	}
}

// groupShell copies a group without its entries or subgroups: those arrive
// through their own iterations, which keeps nesting right and avoids
// duplicating a subtree.
func groupShell(g *Node) *Node {
	out := &Node{Name: g.Name, Text: g.Text, Protected: g.Protected}
	out.Attrs = append([]Attr(nil), g.Attrs...)
	for _, c := range g.Nodes {
		if c.Name == "Group" || c.Name == "Entry" {
			continue
		}
		out.Nodes = append(out.Nodes, c.clone())
	}
	return out
}

// copyGroupFields replaces a group's own properties, keeping its children.
func copyGroupFields(dst, src *Node) {
	kept := []*Node{}
	for _, c := range dst.Nodes {
		if c.Name == "Group" || c.Name == "Entry" {
			kept = append(kept, c)
		}
	}
	fresh := groupShell(src)
	dst.Attrs = fresh.Attrs
	dst.Nodes = append(fresh.Nodes, kept...)
}

// ---- entries ----

func mergeEntries(target, source *Database, ti, si *index, res *MergeResult) {
	maxHistory := historyLimit(target)

	for _, uuid := range si.entryOrder {
		src := si.entries[uuid]
		dst, exists := ti.entries[uuid]

		if !exists {
			if deletedLaterThan(ti.deleted[uuid], lastModified(src.node)) {
				continue
			}
			clone := src.node.clone()
			remapBinaries(clone, source, target)
			parent := targetParent(target, ti, si, src)
			parent.Nodes = append(parent.Nodes, clone)
			ti.entries[uuid] = &objRef{node: clone, parent: parent}
			res.EntriesAdded++
			continue
		}

		tLast, sLast := lastModified(dst.node), lastModified(src.node)
		switch {
		case sLast.After(tLast):
			// The source is newer: its content takes over, and what the
			// target held becomes a history item so it is never lost.
			losing := entryShell(dst.node)
			incoming := src.node.clone()
			remapBinaries(incoming, source, target)
			history := mergeHistory(dst.node, incoming, losing, maxHistory, res)
			replaceEntryContent(dst.node, incoming, history)
			res.EntriesUpdated++
		case tLast.After(sLast):
			// The target is newer: keep it, but file the source's version.
			incoming := src.node.clone()
			remapBinaries(incoming, source, target)
			losing := entryShell(incoming)
			history := mergeHistory(dst.node, dst.node, losing, maxHistory, res)
			setHistory(dst.node, history)
		default:
			history := mergeHistory(dst.node, dst.node, nil, maxHistory, res)
			setHistory(dst.node, history)
		}
		relocate(target, ti, si, src, dst, res)
	}
}

// entryShell copies an entry without its history, which is what a history
// item is: a past state of the entry itself.
func entryShell(e *Node) *Node {
	out := &Node{Name: e.Name, Text: e.Text, Protected: e.Protected}
	out.Attrs = append([]Attr(nil), e.Attrs...)
	for _, c := range e.Nodes {
		if c.Name == "History" {
			continue
		}
		out.Nodes = append(out.Nodes, c.clone())
	}
	return out
}

// replaceEntryContent swaps an entry's fields for another's, then installs
// the merged history.
func replaceEntryContent(dst, src *Node, history []*Node) {
	fresh := entryShell(src)
	dst.Attrs = fresh.Attrs
	dst.Nodes = fresh.Nodes
	setHistory(dst, history)
}

func setHistory(entry *Node, history []*Node) {
	h := entry.Find("History")
	if len(history) == 0 {
		if h != nil {
			entry.Remove(h)
		}
		return
	}
	if h == nil {
		h = &Node{Name: "History"}
		entry.Nodes = append(entry.Nodes, h)
	}
	h.Nodes = history
}

// mergeHistory unions the history of both entries, plus an optional extra
// version, keeping one item per modification time and the newest
// maxHistory of them.
func mergeHistory(targetEntry, sourceEntry, extra *Node, maxHistory int, res *MergeResult) []*Node {
	seen := map[int64]bool{}
	out := []*Node{}
	before := 0

	add := func(n *Node, counts bool) {
		if n == nil {
			return
		}
		key := lastModified(n).Unix()
		if seen[key] {
			return
		}
		seen[key] = true
		out = append(out, n.clone())
		if counts {
			res.HistoryAdded++
		}
	}

	for _, h := range targetEntry.Path("History").FindAll("Entry") {
		add(h, false)
		before++
	}
	for _, h := range sourceEntry.Path("History").FindAll("Entry") {
		add(h, true)
	}
	add(extra, true)

	sort.SliceStable(out, func(i, j int) bool {
		return lastModified(out[i]).Before(lastModified(out[j]))
	})
	if maxHistory > 0 && len(out) > maxHistory {
		out = out[len(out)-maxHistory:]
	}
	return out
}

func historyLimit(db *Database) int {
	v := db.Root.Value("Meta", "HistoryMaxItems")
	n, err := strconv.Atoi(v)
	if err != nil || n < 0 {
		return 10
	}
	return n
}

// ---- relocation ----

// relocate moves an object when the source says it was moved more recently
// than the target did.
func relocate(target *Database, ti, si *index, src, dst *objRef, res *MergeResult) {
	srcParent := si.parentUUID(src)
	if srcParent == "" || dst.parent == nil {
		return
	}
	if dst.parent.Value("UUID") == srcParent {
		return
	}
	if !locationChanged(src.node).After(locationChanged(dst.node)) {
		return
	}
	newParent, ok := ti.groups[srcParent]
	if !ok {
		return
	}
	if dst.parent.Remove(dst.node) {
		newParent.node.Nodes = append(newParent.node.Nodes, dst.node)
		dst.parent = newParent.node
		res.Relocated++
	}
}

func (ix *index) parentUUID(ref *objRef) string {
	if ref.parent == nil {
		return ""
	}
	return ref.parent.Value("UUID")
}

// targetParent finds where an incoming object belongs in the target,
// falling back to the root group when its parent is unknown, so an object is
// never dropped for want of a home.
func targetParent(target *Database, ti, si *index, src *objRef) *Node {
	if uuid := si.parentUUID(src); uuid != "" {
		if ref, ok := ti.groups[uuid]; ok {
			return ref.node
		}
	}
	if g := target.RootGroup(); g != nil {
		return g
	}
	return target.Root.Find("Root")
}

// ---- deletions ----

// applyDeletions removes what either side deleted, unless the other side
// modified the object after the deletion, in which case the modification
// wins and the deletion record is dropped.
func applyDeletions(target *Database, ti, si *index, res *MergeResult) {
	all := map[string]time.Time{}
	for uuid, t := range ti.deleted {
		all[uuid] = t
	}
	for uuid, t := range si.deleted {
		if prev, ok := all[uuid]; !ok || t.After(prev) {
			all[uuid] = t
		}
	}

	keep := map[string]time.Time{}
	uuids := make([]string, 0, len(all))
	for uuid := range all {
		uuids = append(uuids, uuid)
	}
	sort.Strings(uuids)

	for _, uuid := range uuids {
		delTime := all[uuid]
		var ref *objRef
		if r, ok := ti.entries[uuid]; ok {
			ref = r
		} else if r, ok := ti.groups[uuid]; ok {
			ref = r
		}
		if ref == nil {
			keep[uuid] = delTime // nothing here to delete, but remember it
			continue
		}
		if lastModified(ref.node).After(delTime) {
			res.Resurrected++ // edited after it was deleted: the edit wins
			continue
		}
		if ref.parent != nil && ref.parent.Remove(ref.node) {
			res.Deleted++
		}
		delete(ti.entries, uuid)
		delete(ti.groups, uuid)
		keep[uuid] = delTime
	}
	writeDeletedObjects(target, keep)
}

func writeDeletedObjects(target *Database, records map[string]time.Time) {
	root := target.Root.Find("Root")
	if root == nil {
		return
	}
	node := root.Find("DeletedObjects")
	if node == nil {
		node = &Node{Name: "DeletedObjects"}
		root.Nodes = append(root.Nodes, node)
	}
	uuids := make([]string, 0, len(records))
	for uuid := range records {
		uuids = append(uuids, uuid)
	}
	sort.Strings(uuids)
	node.Nodes = nil
	for _, uuid := range uuids {
		node.Nodes = append(node.Nodes, &Node{Name: "DeletedObject", Nodes: []*Node{
			{Name: "UUID", Text: uuid},
			{Name: "DeletionTime", Text: EncodeTime(records[uuid])},
		}})
	}
}

func deletedLaterThan(deletion, modification time.Time) bool {
	return !deletion.IsZero() && !deletion.Before(modification)
}

// ---- attachments ----

// remapBinaries rewrites an incoming entry's attachment references so they
// point into the target's pool, adding any attachment the target lacks.
// Without this, a merged entry would silently gain the wrong attachment.
func remapBinaries(entry *Node, source, target *Database) {
	entry.walk(func(n *Node) {
		if n.Name != "Binary" {
			return
		}
		value := n.Find("Value")
		if value == nil {
			return
		}
		ref := value.Attr("Ref")
		if ref == "" {
			return
		}
		idx, err := strconv.Atoi(ref)
		if err != nil || idx < 0 || idx >= len(source.Binaries) {
			return
		}
		value.SetAttr("Ref", strconv.Itoa(target.addBinary(source.Binaries[idx])))
	})
}

// addBinary returns the index of an attachment in the pool, appending it if
// it is not already there.
func (d *Database) addBinary(b Binary) int {
	for i, have := range d.Binaries {
		if len(have.Data) == len(b.Data) && string(have.Data) == string(b.Data) {
			return i
		}
	}
	d.Binaries = append(d.Binaries, Binary{Flags: b.Flags,
		Data: append([]byte(nil), b.Data...)})
	return len(d.Binaries) - 1
}

// ---- meta ----

// mergeMeta takes the source's database-level settings when they were
// changed more recently, and unions the custom data and icons either way.
func mergeMeta(target, source *Database, res *MergeResult) {
	tm, sm := target.Root.Find("Meta"), source.Root.Find("Meta")
	if tm == nil || sm == nil {
		return
	}
	tChanged, _ := DecodeTime(tm.Value("SettingsChanged"))
	sChanged, _ := DecodeTime(sm.Value("SettingsChanged"))
	if sChanged.After(tChanged) {
		for _, field := range metaSettings {
			src := sm.Find(field)
			if src == nil {
				continue
			}
			if dst := tm.Find(field); dst != nil {
				tm.Nodes[indexOf(tm.Nodes, dst)] = src.clone()
			} else {
				tm.Nodes = append(tm.Nodes, src.clone())
			}
		}
		res.GroupsUpdated += 0 // settings are not objects; nothing to count
	}
	mergeKeyed(tm, sm, "CustomData", "Item", "Key")
	mergeKeyed(tm, sm, "CustomIcons", "Icon", "UUID")
}

// metaSettings is the set of database-level fields that travel together
// under Meta/SettingsChanged.
var metaSettings = []string{
	"DatabaseName", "DatabaseNameChanged",
	"DatabaseDescription", "DatabaseDescriptionChanged",
	"DefaultUserName", "DefaultUserNameChanged",
	"Color", "MemoryProtection",
	"RecycleBinEnabled", "RecycleBinUUID", "RecycleBinChanged",
	"EntryTemplatesGroup", "EntryTemplatesGroupChanged",
	"MaintenanceHistoryDays", "HistoryMaxItems", "HistoryMaxSize",
	"SettingsChanged",
}

// mergeKeyed unions a container of keyed items, newest wins where both sides
// have the same key.
func mergeKeyed(targetMeta, sourceMeta *Node, container, item, keyField string) {
	sc := sourceMeta.Find(container)
	if sc == nil || len(sc.FindAll(item)) == 0 {
		return
	}
	tc := targetMeta.Find(container)
	if tc == nil {
		tc = &Node{Name: container}
		targetMeta.Nodes = append(targetMeta.Nodes, tc)
	}
	have := map[string]*Node{}
	for _, it := range tc.FindAll(item) {
		have[it.Value(keyField)] = it
	}
	for _, it := range sc.FindAll(item) {
		key := it.Value(keyField)
		existing, ok := have[key]
		if !ok {
			tc.Nodes = append(tc.Nodes, it.clone())
			have[key] = it
			continue
		}
		sMod, _ := DecodeTime(it.Value("LastModificationTime"))
		tMod, _ := DecodeTime(existing.Value("LastModificationTime"))
		if sMod.After(tMod) {
			tc.Nodes[indexOf(tc.Nodes, existing)] = it.clone()
		}
	}
}

func indexOf(nodes []*Node, target *Node) int {
	for i, n := range nodes {
		if n == target {
			return i
		}
	}
	return 0
}