package dirtree

import (
	"os"
	"path/filepath"
	"strings"
)

type Tree struct {
	Root       *Node
	ShowHidden bool
}

type Node struct {
	tree     *Tree
	parent   *Node
	children []*Node
	depth    int
	Path     string
	Name     string
	Expanded bool
	isDir    bool
	selected bool
}

func NewTree(path string) *Tree {
	tree := &Tree{
		ShowHidden: false,
	}
	tree.Root = &Node{
		tree:     tree,
		parent:   nil,
		Path:     path,
		Name:     filepath.Base(path),
		isDir:    true,
		Expanded: true,
	}
	return tree
}

func (node *Node) IsRoot() bool {
	return node.tree.Root == node
}

func (node *Node) Parent() *Node {
	return node.parent
}

func (node *Node) IsSelected() bool {
	return node.selected
}

func (node *Node) IsDir() bool {
	return node.isDir
}

func (node *Node) Select(recurse bool, maxDepth int) {
	node.selected = true
	if node.isDir && recurse {
		node.LoadChildren(recurse, maxDepth, func (n *Node) {
			n.selected = true
		})
	}
}

func (node *Node) Deselect(recurse bool) {
	node.selected = false
	if node.isDir && recurse {
		for _, child := range node.children {
			child.Deselect(recurse)
		}
	}
}

func (node *Node) LoadChildren(recurse bool, maxDepth int, visitFunc func (node *Node)) error {
	startDepth := node.depth
	var loadChildren func(node *Node) error
	loadChildren = func(node *Node) error {
		// Load new entries
		entries, err := os.ReadDir(node.Path)
		if err != nil {
			return err
		}

		// Save current children before loading
		oldChildren := make(map[string]*Node)
		for _, child := range node.children {
			oldChildren[child.Path] = child
		}

		node.children = make([]*Node, 0)
		for _, entry := range entries {
			if !node.tree.ShowHidden && strings.HasPrefix(entry.Name(), ".") {
				continue
			}

			child := &Node{
				tree:   node.tree,
				parent: node,
				depth:  node.depth + 1,
				Name:   entry.Name(),
				isDir:  entry.IsDir(),
				Path:   filepath.Join(node.Path, entry.Name()),
			}

			oldChild, exists := oldChildren[child.Path]
			if child.isDir {
				// Preserve child's grandchildren and expanded/selected state
				if exists && (oldChild.Expanded || len(oldChild.children) > 0) {
					child.Expanded = oldChild.Expanded
					child.selected = oldChild.selected
					child.children = oldChild.children
					for i := range child.children {
						child.children[i].parent = child
					}
					loadChildren(child)
				} else if recurse && (maxDepth <= 0 || node.depth - startDepth < maxDepth) {
					loadChildren(child)
				}
			} else if exists {
				child.selected = oldChild.selected
			}

			node.children = append(node.children, child)
			if visitFunc != nil {
				visitFunc(child)
			}
		}
		return nil
	}
	return loadChildren(node)
}