package picker

import (
	"io/fs"
	"path"
	"sort"
)

// CheckState represents the tri-state checkbox value for a node.
type CheckState int

const (
	// CheckNone means no descendants are selected.
	CheckNone CheckState = iota
	// CheckPartial means some but not all descendants are selected.
	CheckPartial
	// CheckAll means all selectable descendants are selected.
	CheckAll
)

// nodeKind distinguishes leaf-selectable nodes from structural ones.
type nodeKind int

const (
	kindFile     nodeKind = iota // regular file — always selectable
	kindEmptyDir                 // directory with no children — selectable leaf
	kindDir                      // directory with children — state is derived
)

// Selection tracks multi-select state over an fs.FS tree. It maintains
// a flat set of selected leaves (files and empty dirs) and cached
// subtree counts for O(1) tri-state lookups.
//
// The tree index is built once from the FS and is immutable; only the
// selected set and subtree counts change on toggle.
type Selection struct {
	selected map[string]bool // selected leaves (files + empty dirs)

	// Immutable tree index, built once from the FS.
	children        map[string][]string // parent → sorted child paths
	parent          map[string]string   // child → parent path
	kind            map[string]nodeKind
	totalSelectable map[string]int // total selectable leaves in subtree

	// Mutable counter, updated incrementally on toggle.
	selectedCount map[string]int // selected leaves in subtree
}

// NewSelection builds a Selection from the given fs.FS by walking the
// entire tree once. The FS must implement fs.ReadDirFS.
func NewSelection(fsys fs.FS) *Selection {
	s := &Selection{
		selected:        make(map[string]bool),
		children:        make(map[string][]string),
		parent:          make(map[string]string),
		kind:            make(map[string]nodeKind),
		totalSelectable: make(map[string]int),
		selectedCount:   make(map[string]int),
	}

	s.buildIndex(fsys, ".")
	return s
}

// buildIndex recursively walks the FS to populate the tree index.
// Returns the total selectable count for the subtree rooted at dir.
func (s *Selection) buildIndex(fsys fs.FS, dir string) int {
	entries, err := fs.ReadDir(fsys, dir)
	if err != nil {
		return 0
	}

	total := 0
	for _, e := range entries {
		child := path.Join(dir, e.Name())

		s.parent[child] = dir
		if e.IsDir() {
			childTotal := s.buildIndex(fsys, child)
			if childTotal == 0 {
				// Empty directory — treat as a selectable leaf.
				s.kind[child] = kindEmptyDir
				s.totalSelectable[child] = 1
				total++
			} else {
				s.kind[child] = kindDir
				s.totalSelectable[child] = childTotal
				total += childTotal
			}
		} else {
			s.kind[child] = kindFile
			s.totalSelectable[child] = 1
			total++
		}
		s.children[dir] = append(s.children[dir], child)
	}

	s.totalSelectable[dir] = total
	s.kind[dir] = kindDir

	return total
}

// State returns the tri-state checkbox value for the given path.
func (s *Selection) State(p string) CheckState {
	k := s.kind[p]
	switch k {
	case kindFile, kindEmptyDir:
		if s.selected[p] {
			return CheckAll
		}
		return CheckNone
	default:
		total := s.totalSelectable[p]
		count := s.selectedCount[p]
		switch count {
		case 0:
			return CheckNone
		case total:
			return CheckAll
		default:
			return CheckPartial
		}
	}
}

// Toggle flips the selection of the given path. For files and empty
// dirs, this is a simple toggle. For directories, it selects all
// descendants if the current state is none or partial, or deselects
// all if fully selected.
func (s *Selection) Toggle(p string) {
	state := s.State(p)

	k := s.kind[p]
	switch k {
	case kindFile, kindEmptyDir:
		if s.selected[p] {
			delete(s.selected, p)
			s.adjustCounts(p, -1)
		} else {
			s.selected[p] = true
			s.adjustCounts(p, +1)
		}
	default:
		if state == CheckAll {
			s.setSubtree(p, false)
		} else {
			s.setSubtree(p, true)
		}
	}
}

// setSubtree selects or deselects all selectable leaves under p.
// It walks the subtree once, flipping leaves and returning the count
// of changes, then adjusts ancestor counts in a single upward pass.
func (s *Selection) setSubtree(p string, sel bool) {
	changed := s.applySubtree(p, sel)
	if changed == 0 {
		return
	}

	// Adjust ancestors of p (not p itself — applySubtree already
	// updated interior counts within the subtree).
	delta := changed
	if !sel {
		delta = -changed
	}
	cur := p
	for {
		par, ok := s.parent[cur]
		if !ok {
			break
		}
		s.selectedCount[par] += delta
		cur = par
	}
}

// applySubtree recursively sets leaves under p to the desired state.
// For interior nodes within the subtree, it updates selectedCount
// directly. Returns the number of leaves that actually changed.
func (s *Selection) applySubtree(p string, sel bool) int {
	k := s.kind[p]
	switch k {
	case kindFile, kindEmptyDir:
		was := s.selected[p]
		if sel == was {
			return 0
		}
		if sel {
			s.selected[p] = true
		} else {
			delete(s.selected, p)
		}
		return 1
	default:
		changed := 0
		for _, child := range s.children[p] {
			changed += s.applySubtree(child, sel)
		}
		// Update this interior node's count directly.
		if sel {
			s.selectedCount[p] += changed
		} else {
			s.selectedCount[p] -= changed
		}
		return changed
	}
}

// adjustCounts increments or decrements selectedCount for all ancestors
// of p (not p itself, since leaves don't have subtree counts).
func (s *Selection) adjustCounts(p string, delta int) {
	cur := p
	for {
		par, ok := s.parent[cur]
		if !ok {
			break
		}
		s.selectedCount[par] += delta
		cur = par
	}
}

// AllSelected reports whether every selectable leaf in the tree is
// selected. When true, the caller should perform a full restore
// rather than emitting --include flags.
func (s *Selection) AllSelected() bool {
	return s.State(".") == CheckAll
}

// SelectedPaths returns the minimal set of FS-relative paths covering
// the current selection. Fully-selected subtrees are compressed to
// their root directory. Returns nil if nothing is selected or if
// everything is selected (check AllSelected() to distinguish).
func (s *Selection) SelectedPaths() []string {
	if s.State(".") == CheckNone || s.State(".") == CheckAll {
		return nil
	}
	var paths []string
	s.collectPaths(".", &paths)
	sort.Strings(paths)
	return paths
}

// collectPaths walks the tree top-down, emitting fully-selected
// subtrees as single paths and recursing into partial ones.
func (s *Selection) collectPaths(p string, out *[]string) {
	for _, child := range s.children[p] {
		state := s.State(child)
		switch state {
		case CheckNone:
			continue
		case CheckAll:
			// Fully selected — emit just this path, don't recurse.
			*out = append(*out, child)
		case CheckPartial:
			s.collectPaths(child, out)
		}
	}
}