//! Stealing: bringing another machine's saves in, and combining when both
//! machines moved.
//!
//! The order is always the same, and it is what makes an interrupted steal
//! harmless: read everything from the remote first, then land it in one
//! transaction, then write the working tree. Divergence is resolved per file —
//! a clash is never blended line by line, it is chosen.

module exchange

import compress.zlib
import src.apply
import src.events
import src.fmt
import src.history
import src.remote
import src.snapshot
import src.storage
import strconv

pub enum Side {
	mine
	theirs
}

// Clash is one file both machines changed differently. `stat` is the short
// diff stat the terminal prompt shows, so the choice is informed.
pub struct Clash {
pub mut:
	path         string
	mine         snapshot.FileEntry
	theirs       snapshot.FileEntry
	base         snapshot.FileEntry
	mine_stat    string
	theirs_stat  string
	mine_patch   string
	theirs_patch string
	mine_save    history.Save
	theirs_save  history.Save
}

pub struct StealPreview {
pub:
	repo      string
	base      history.Save
	local     []history.Save
	remote    []SaveRecord
	changed   int
	automatic int
	identical int
	clashes   []Clash
	yes       bool
	all_side  string
}

pub struct StealOptions {
pub mut:
	mine   bool // resolve every clash as mine
	theirs bool // resolve every clash as theirs
	// Shown the whole combine and asked about every clash. Nil means "nobody to ask".
	guide fn (preview StealPreview) !map[string]Side = unsafe { nil }
	yes   bool
}

pub struct StealReport {
pub mut:
	repo            string
	forward         bool
	fetched_saves   int
	fetched_objects int
	fetched_events  int
	clashes         []string
	kept_mine       int
	replayed        int
	removed         int
	head            string
	flagged         []events.Flag // incoming saves that are flagged
	// Set when a combine took this machine's saves off the line: loading it
	// brings back the tree from before the steal, less anything stripped.
	previous_head string
	kept_versions []string
	unsafe_paths  []string
}

pub const dirty_tree = 'save or checkpoint first — steal needs a clean tree'

// ClashError is a clash nobody was there to decide. It carries the paths because the advice
// depends on the command: `koh steal` takes --mine or --theirs, `koh offer` does not.
pub struct ClashError {
	Error
pub:
	paths []string
}

pub fn (e ClashError) msg() string {
	return '${e.paths.join(', ')} changed on both machines — run on a terminal to choose, or pass --mine or --theirs'
}

pub const no_saves_here = 'this folder has files but no saves — steal a repo into an empty folder with "koh steal <name>"'

// steal brings the remote's new saves in. Inside a repo with its own work this
// may have to combine; a repo with nothing of its own simply fast-forwards.
pub fn steal(store &storage.Store, mut backend remote.Backend, opts StealOptions) !StealReport {
	name := repo_name(store, mut backend)!
	if snapshot.restore_pending(store) {
		return error(snapshot.restore_unfinished)
	}
	head_before := history.head(store)
	// Combining rewrites the working tree, so unsaved work would be lost.
	if dirty(store, head_before)! {
		if head_before == none {
			return error(no_saves_here)
		}
		return error(dirty_tree)
	}
	fetched := fetch_back_to_known(store, mut backend, name)!
	mut report := StealReport{
		repo:            name
		fetched_saves:   fetched.saves.len
		fetched_objects: fetched.repo_objects.len
		fetched_events:  fetched.events_out.len
		head:            fetched.head
		unsafe_paths:    fetched.unsafe_paths
	}
	mut local_only := []history.Save{}
	if head := head_before {
		// The remote builds on `base`. If that save is here but no longer on
		// the line, neither line contains the other; jumping to the remote
		// would silently leave this line behind.
		local_only = saves_after(store, head.id, fetched.base) or {
			return error('the remote builds on save ${fetched.base}, which is not on this line — load the save you want, then offer')
		}
	}
	if fetched.base == '' && local_only.len > 0 {
		return error('the remote keeps a different history — it shares no save with this repo')
	}
	if fetched.saves.len > 0 {
		store.begin()!
		write_fetched_records(store, fetched) or {
			store.rollback()
			return err
		}
		store.end()!
	}

	strips_before := events.strip_events(store)!.len
	mut clashes := []Clash{}
	mut side_of := map[string]Side{}
	if fetched.saves.len > 0 && local_only.len > 0 {
		clashes = combine_clashes(store, fetched, local_only)!
		report.clashes = clashes.map(it.path)
		if opts.guide != unsafe { nil } {
			side_of = opts.guide(combine_preview(store, name, fetched, local_only, clashes,
				opts.yes, if opts.mine {
					'mine'
				} else if opts.theirs { 'theirs' } else { '' })!)!
		} else {
			// Asked before the write lock: a person thinking must not stall every
			// agent's save on this repo.
			side_of = choose_sides(clashes, opts)!
		}
	}
	for clash in clashes {
		if side_of[clash.path] == .theirs {
			report.kept_versions << '${clash.path}:${local_only.last().id}'
		}
	}

	store.begin() or { return err }
	// The questions ran without the lock. If a save or an edit landed
	// meanwhile, the answers were about a line that no longer exists.
	if moved := changed_since(store, head_before, strips_before) {
		if moved {
			store.rollback()
			return error('this repo changed while steal was asking — run koh steal again')
		}
	} else {
		store.rollback()
		return err
	}
	final := (write_and_combine(store, fetched, local_only, clashes, side_of, mut report) or {
		store.rollback()
		return err
	})
	files := history.files_of(store, final.id) or {
		store.rollback()
		return err
	}
	snapshot.check_links(store.root, files) or {
		store.rollback()
		return err
	}
	// The files follow the transaction; the mark lands with it, in one step.
	// No head before means a fresh steal, whose whole tree is new.
	moved := if head := head_before { final.id != head.id } else { true }
	if moved {
		snapshot.mark_restore(store) or {
			store.rollback()
			return err
		}
	}
	store.end() or { return err }
	report.head = final.id
	if local_only.len > 0 && report.fetched_saves > 0 {
		report.previous_head = local_only.last().id
	}
	flags := events.latest_flags(store)!
	for record in fetched.saves {
		if flag := flags[record.id] {
			if flag.flagged {
				report.flagged << flag
			}
		}
	}

	if head := head_before {
		if final.id == head.id {
			// The line did not move, so neither does the working tree.
			return report
		}
		mut wanted := map[string]bool{}
		for file in files {
			wanted[file.path] = true
		}
		mut remove := []string{}
		for file in history.files_of(store, head.id)! {
			if file.path !in wanted {
				remove << file.path
			}
		}
		result := snapshot.restore_files(store, store.root, snapshot.load(store.root), files,
			remove)!
		report.removed = result.removed
	} else {
		snapshot.restore_files(store, store.root, snapshot.load(store.root), files, [])!
	}
	history.refresh_log_md(store)
	return report
}

// write_and_combine lands a fetch and, when the local line has saves the remote
// lacks, replays them on top of the remote head. It runs inside the caller's
// transaction and returns the new head.
fn write_and_combine(store &storage.Store, fetched Fetched, local_only []history.Save, clashes []Clash, side_of map[string]Side, mut report StealReport) !history.Save {
	write_fetched(store, fetched)!
	if fetched.saves.len == 0 {
		// The remote head is already on this line: level, or this machine is
		// ahead and the offer that follows a steal sends the rest.
		return history.head(store) or { return error('this repo has no head') }
	}
	if local_only.len == 0 {
		// Nothing of ours is missing from the remote: the line simply moves on.
		history.set_head(store, fetched.head)!
		report.forward = true
		return history.head(store) or { return error('the remote head did not land') }
	}
	remote_files := history.files_of(store, fetched.head)!
	mut mine_files := []snapshot.FileEntry{}
	if our_head := history.head(store) {
		mine_files = history.files_of(store, our_head.id)!
	}

	// A file the remote changed is a clash, decided below, or one this machine
	// left as it was or changed the same way, so the remote's version stands.
	// Replaying this machine's steps on it could only conflict: an edit it
	// later undid would no longer apply.
	base_files := history.files_of(store, fetched.base)!
	settled := changed_paths(base_files, remote_files)

	// Take the remote's line first, so the replayed saves chain onto it.
	history.set_head(store, fetched.head)!
	mut current := remote_files.clone()
	for save in local_only {
		parent_files := if save.parent != '' {
			history.files_of(store, save.parent)!
		} else {
			[]snapshot.FileEntry{}
		}
		target_files := history.files_of(store, save.id)!
		plan := apply.three_way(current, parent_files, target_files, [], settled)
		if apply.conflicts(plan) {
			return error('combine could not apply ${plan.conflict} — it is not a clash file')
		}
		if plan.applied.len == 0 {
			// Everything this save touched was a file the remote changed.
			continue
		}
		history.record_save(store, '', plan.files, history.SaveInput{
			message:       save.message
			created_at:    save.created_at
			machine:       save.machine
			agent:         save.agent
			session:       save.session
			is_checkpoint: save.is_checkpoint
			replayed_from: save.id
		})!
		current = plan.files.clone()
		report.replayed++
	}

	mut mine_paths := []string{}
	for clash in clashes {
		if side_of[clash.path] == .mine {
			mine_paths << clash.path
		}
	}
	if mine_paths.len > 0 {
		// One last save puts my versions of the chosen files back.
		final_plan := apply.three_way(current, current, mine_files, mine_paths, [])
		if final_plan.applied.len > 0 {
			history.record_save(store, '', final_plan.files, history.SaveInput{
				message: 'combine: kept mine for ${mine_paths.len} ${if mine_paths.len == 1 {
					'file'
				} else {
					'files'
				}}'
			})!
			report.kept_mine = mine_paths.len
		}
	}
	return history.head(store) or { return error('the combined line has no head') }
}

// find_clashes lists the files both sides changed, differently. Files changed
// identically on both sides are not clashes: there is nothing to decide.
pub fn find_clashes(base []snapshot.FileEntry, mine []snapshot.FileEntry, theirs []snapshot.FileEntry) []Clash {
	base_map := snapshot.index_files(base)
	mine_map := snapshot.index_files(mine)
	theirs_map := snapshot.index_files(theirs)
	mut paths := mine_map.keys()
	for path in theirs_map.keys() {
		if path !in mine_map {
			paths << path
		}
	}
	paths.sort()
	mut clashes := []Clash{}
	for path in paths {
		here := mine_map[path] or {
			snapshot.FileEntry{
				path: path
			}
		}
		there := theirs_map[path] or {
			snapshot.FileEntry{
				path: path
			}
		}
		was := base_map[path] or {
			snapshot.FileEntry{
				path: path
			}
		}
		if same(here, there) {
			continue
		}
		if same(here, was) || same(there, was) {
			continue
		}
		clashes << Clash{
			path:   path
			mine:   here
			theirs: there
			base:   was
		}
	}
	return clashes
}

// changed_paths lists the paths whose entry differs between two trees.
fn changed_paths(before []snapshot.FileEntry, after []snapshot.FileEntry) []string {
	was := snapshot.index_files(before)
	now := snapshot.index_files(after)
	mut paths := []string{}
	for path, entry in now {
		if !same(entry, was[path] or { snapshot.FileEntry{} }) {
			paths << path
		}
	}
	for path, _ in was {
		if path !in now {
			paths << path
		}
	}
	return paths
}

// combine_clashes finds the files both lines changed, with the diff stat each
// prompt shows. It reads the remote head from the fetch, not the store, so it
// can run before anything is written.
fn combine_clashes(store &storage.Store, fetched Fetched, local_only []history.Save) ![]Clash {
	base_files := if fetched.base != '' {
		history.files_of(store, fetched.base)!
	} else {
		[]snapshot.FileEntry{}
	}
	head := history.head(store) or { return []Clash{} }
	// The transaction applies every known strip before it replays, so the
	// clashes are found on the trees as they will be then, or the answers
	// would be about files that are no longer there.
	strips := known_strips(store, fetched)!
	remote_head := fetched.saves.last()
	mut clashes := find_clashes(after_strips(base_files, fetched.base, strips), after_strips(history.files_of(store,
		head.id)!, head.id, strips), after_strips(to_files(remote_head), remote_head.id, strips))
	for mut clash in clashes {
		base := fetched_bytes(store, fetched, clash.base)!
		mine := fetched_bytes(store, fetched, clash.mine)!
		theirs := fetched_bytes(store, fetched, clash.theirs)!
		mine_delta := snapshot.diff_bytes(clash.path, 'modified', base, mine)!
		theirs_delta := snapshot.diff_bytes(clash.path, 'modified', base, theirs)!
		clash.mine_stat = delta_stat(mine_delta, clash.mine.size)
		clash.theirs_stat = delta_stat(theirs_delta, clash.theirs.size)
		clash.mine_patch = mine_delta.patch
		clash.theirs_patch = theirs_delta.patch
		clash.mine_save = last_touch(store, local_only, clash.path)!
		clash.theirs_save = last_touch(store, fetched.saves.map(to_save(it)), clash.path)!
	}
	return clashes
}

fn delta_stat(delta snapshot.Delta, size i64) string {
	if delta.binary {
		return fmt.human_bytes(size)
	}
	return '+${delta.added} -${delta.removed}'
}

fn last_touch(store &storage.Store, saves []history.Save, path string) !history.Save {
	for index := saves.len - 1; index >= 0; index-- {
		save := saves[index]
		current := snapshot.index_files(history.files_of(store, save.id)!)[path] or {
			snapshot.FileEntry{}
		}
		parent := if save.parent != '' {
			snapshot.index_files(history.files_of(store, save.parent)!)[path] or {
				snapshot.FileEntry{}
			}
		} else {
			snapshot.FileEntry{}
		}
		if !same(current, parent) {
			return save
		}
	}
	return saves.last()
}

// combine_preview counts what a combine would do to the paths either side changed since `base`:
// `identical` were changed the same way on both sides, `clashes` need a choice, the rest go in by themselves.
fn combine_preview(store &storage.Store, name string, fetched Fetched, local []history.Save, clashes []Clash, yes bool, all_side string) !StealPreview {
	base := history.files_of(store, fetched.base)!
	mine := history.files_of(store, local.last().id)!
	theirs := history.files_of(store, fetched.head)!
	mine_changed := changed_paths(base, mine)
	theirs_changed := changed_paths(base, theirs)
	mut changed := map[string]bool{}
	for path in mine_changed {
		changed[path] = true
	}
	for path in theirs_changed {
		changed[path] = true
	}
	mut identical := 0
	base_map := snapshot.index_files(base)
	mine_map := snapshot.index_files(mine)
	theirs_map := snapshot.index_files(theirs)
	for path, _ in changed {
		if !same(mine_map[path] or { snapshot.FileEntry{} }, base_map[path] or { snapshot.FileEntry{} })
			&& same(mine_map[path] or { snapshot.FileEntry{} }, theirs_map[path] or { snapshot.FileEntry{} }) {
			identical++
		}
	}
	return StealPreview{
		repo:      name
		base:      history.by_id(store, fetched.base) or { history.Save{} }
		local:     local
		remote:    fetched.saves
		changed:   changed.len
		automatic: changed.len - identical - clashes.len
		identical: identical
		clashes:   clashes
		yes:       yes
		all_side:  all_side
	}
}

// after_strips is a save's file list once the given strip events are applied,
// by the same rule write_fetched uses.
fn after_strips(files []snapshot.FileEntry, save_id string, strips []events.Event) []snapshot.FileEntry {
	return files.filter(fn [strips, save_id] (file snapshot.FileEntry) bool {
		return !strips.any(reaches(it, save_id) && covers(it.path, file.path))
	})
}

// fetched_bytes reads one side of a clash: from the fetch when it has just
// arrived, else from the store. An absent side is empty.
fn fetched_bytes(store &storage.Store, fetched Fetched, entry snapshot.FileEntry) ![]u8 {
	if entry.hash == '' {
		return []u8{}
	}
	if blob := fetched.repo_objects[entry.hash] {
		return zlib.decompress(blob)!
	}
	return store.get_object(entry.hash)!
}

fn choose_sides(clashes []Clash, opts StealOptions) !map[string]Side {
	mut sides := map[string]Side{}
	if clashes.len == 0 {
		return sides
	}
	for clash in clashes {
		if opts.mine {
			sides[clash.path] = .mine
			continue
		}
		if opts.theirs {
			sides[clash.path] = .theirs
			continue
		}
		return ClashError{
			paths: clashes.map(it.path)
		}
	}
	return sides
}

fn same(a snapshot.FileEntry, b snapshot.FileEntry) bool {
	if a.hash == '' && b.hash == '' {
		return true
	}
	return a.hash == b.hash && a.executable == b.executable
}

fn changed_since(store &storage.Store, before ?history.Save, strips_before int) !bool {
	if events.strip_events(store)!.len != strips_before {
		return true
	}
	now := history.head(store)
	if now or { history.Save{} }.id != before or { history.Save{} }.id {
		return true
	}
	return dirty(store, now)!
}

// dirty reports whether the working tree holds anything a steal would
// overwrite. With no saves yet, that is every file but the untouched ignore
// list `koh init` wrote.
pub fn dirty(store &storage.Store, head ?history.Save) !bool {
	files := snapshot.walk(store, store.root, snapshot.load(store.root))!
	if saved := head {
		return !snapshot.same_tree(history.files_of(store, saved.id)!, files)
	}
	seed := storage.hash_hex(snapshot.default_ignore.bytes())
	return files.any(it.path != snapshot.ignore_file || it.hash != seed)
}

// fetch_repo takes a whole repo — head, saves, objects and events — for
// `koh steal <name>`, which lands a copy in a new folder.
pub fn fetch_repo(mut backend remote.Backend, name string) !Fetched {
	check_format(mut backend, name)!
	remote_head := backend.get(head_key(name))!
	if !remote_head.found {
		return error('the remote has no repo called "${name}"')
	}
	mut fetched := Fetched{
		head: remote_head.body.bytestr().trim_space()
	}
	bundle := backend.get(bundle_key(name))!
	if bundle.found {
		take_bundle(mut fetched, bundle.body)!
	} else {
		for kind in ['objects', 'saves', 'events'] {
			for key in backend.list('${name}/${kind}')! {
				body := backend.get(key)!
				if body.found {
					take(mut fetched, key.all_after('${name}/'), body.body)!
				}
			}
		}
	}
	fetched.saves = parent_first(fetched.saves)!
	fetched.events_out.sort(a.id < b.id)
	return fetched
}

// take_bundle splits §22.7's bundle — `<key> <length>\n` and then that many bytes, over and over,
// until `end 0` — into what one GET of each key would have returned. A connection that drops
// between two entries still looks like a whole bundle, so only the end line proves it arrived.
fn take_bundle(mut fetched Fetched, body []u8) ! {
	mut at := 0
	for at < body.len {
		// Scanned in place: a slice of `body` copies the rest of it, and that per entry is quadratic.
		mut newline := at
		for newline < body.len && body[newline] != `\n` {
			newline++
		}
		if newline == body.len {
			return error('the bundle ends part way through an entry')
		}
		header := body[at..newline].bytestr()
		if header == 'end 0' && newline + 1 == body.len {
			return
		}
		length := strconv.atoi(header.all_after(' '))!
		start := newline + 1
		if length < 0 || length > body.len - start {
			return error('the bundle ends part way through ${header.all_before(' ')}')
		}
		take(mut fetched, header.all_before(' '), body[start..start + length])!
		at = start + length
	}
	return error('the bundle was cut short — try the steal again')
}

// take files one fetched key, relative to the repo: an object, a save or an event.
fn take(mut fetched Fetched, key string, body []u8) ! {
	match key.all_before('/') {
		'objects' {
			fetched.repo_objects[key.all_after('/')] = body
		}
		'saves' {
			mut record := decode_save(body)!
			fetched.unsafe_paths << drop_unsafe_paths(mut record)
			fetched.saves << record
		}
		'events' {
			event := decode_event(body)!
			if !ignored_strip(event) {
				fetched.events_out << event
			}
		}
		else {
			return error('the remote sent "${key}", which is not part of a repo')
		}
	}
}

// parent_first orders saves so every parent lands before its children, as the
// foreign key on `parent` demands. Time cannot do it: a replayed save keeps its
// original, older time but sits on a newer parent.
fn parent_first(records []SaveRecord) ![]SaveRecord {
	mut by_id := map[string]SaveRecord{}
	for record in records {
		by_id[record.id] = record
	}
	mut placed := map[string]bool{}
	mut ordered := []SaveRecord{}
	for record in records {
		mut chain := []SaveRecord{}
		mut cursor := record.id
		for cursor != '' && cursor !in placed {
			if chain.len > records.len {
				return error('the save records on the remote form a loop')
			}
			current := by_id[cursor] or { return error('the remote is missing save ${cursor}') }
			chain << current
			cursor = current.parent
		}
		chain.reverse_in_place()
		for current in chain {
			placed[current.id] = true
			ordered << current
		}
	}
	return ordered
}