lmcli/pkg/tui/bubbles/filepicker/model.go

package filepicker

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

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

type Node struct {
    path     string
    name     string
    isDir    bool
    expanded bool
    children []*Node
}

type Model struct {
    root         *Node
    cursor       int
    showHidden   bool
    displayLines []*Node // Flattened view of visible nodes
    rootPath     string

    SelectedPaths []string
}

func New(path string) Model {
    root := &Node{
        path:     path,
        name:     filepath.Base(path),
        isDir:    true,
        expanded: true,
    }

    m := Model{
        root:          root,
        rootPath:      path,
        SelectedPaths: make([]string, 0),
        showHidden:    false,
    }

    m.loadDirectory(root)
    m.updateDisplayLines()
    return m
}

func (m *Model) loadDirectory(node *Node) error {
    entries, err := os.ReadDir(node.path)
    if err != nil {
        return err
    }

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

        child := &Node{
            path:     filepath.Join(node.path, entry.Name()),
            name:     entry.Name(),
            isDir:    entry.IsDir(),
            expanded: false,
        }
        node.children = append(node.children, child)
    }
    return nil
}

func (m *Model) updateDisplayLines() {
    m.displayLines = make([]*Node, 0)
    m.flattenTree(m.root, 0)
}

func (m *Model) flattenTree(node *Node, depth int) {
    m.displayLines = append(m.displayLines, node)
    if node.isDir && node.expanded {
        for _, child := range node.children {
            m.flattenTree(child, depth+1)
        }
    }
}

func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
    switch msg := msg.(type) {
    case tea.KeyMsg:
        switch msg.String() {
        case "j":
            if m.cursor < len(m.displayLines)-1 {
                m.cursor++
            }
        case "k":
            if m.cursor > 0 {
                m.cursor--
            }
        case "l":
            current := m.displayLines[m.cursor]
            if current.isDir {
                current.expanded = true
                m.loadDirectory(current)
                m.updateDisplayLines()
            }
        case "h":
            current := m.displayLines[m.cursor]
            if current.isDir && current.expanded {
                current.expanded = false
                m.updateDisplayLines()
            } else {
                // Find and collapse parent
                parent := m.findParent(current)
                if parent != nil && parent.path != m.rootPath {
                    parent.expanded = false
                    m.updateDisplayLines()
                    // Move cursor to parent
                    for i, node := range m.displayLines {
                        if node == parent {
                            m.cursor = i
                            break
                        }
                    }
                }
            }
        case " ":
            current := m.displayLines[m.cursor]
            if m.isSelected(current.path) {
                m.SelectedPaths = removeString(m.SelectedPaths, current.path)
            } else {
                m.SelectedPaths = append(m.SelectedPaths, current.path)
            }
        case "H":
            m.showHidden = !m.showHidden
            m.reloadTree()
        }
    }
    return m, nil
}

func (m Model) Init() tea.Cmd {
    return nil
}

func (m Model) View() string {
    var sb strings.Builder

    // Define styles
    dirStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("12"))    // Blue for directories
    fileStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("7"))    // White for files
    cursorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("yellow"))
    selectedStyle := lipgloss.NewStyle().Bold(true)

    // Track the parent-child relationships for proper indentation
    levelIsLast := make(map[int]bool)

    for i, node := range m.displayLines {
        // Calculate depth and prefix
        depth := 0
        if node != m.root {
            depth = strings.Count(node.path[len(m.rootPath):], string(os.PathSeparator))
        }

        // Build the prefix (tree structure)
        var prefix strings.Builder

        if node != m.root {
            // For each level up to current depth, add the appropriate prefix
            for d := 0; d < depth; d++ {
                if levelIsLast[d] {
                    prefix.WriteString("    ")
                } else {
                    prefix.WriteString("│   ")
                }
            }

            // Find if this node is the last among its siblings
            parent := m.findParent(node)
            if parent != nil {
                isLast := parent.children[len(parent.children)-1] == node
                levelIsLast[depth] = isLast

                if isLast {
                    prefix.WriteString("└── ")
                } else {
                    prefix.WriteString("├── ")
                }
            }
        }

        // Cursor indicator
        cursor := " "
        if i == m.cursor {
            cursor = cursorStyle.Render(">")
        }

        // Node name with optional caret for directories
        var nodeName string
        if node.isDir {
            caret := "▶ "
            if node.expanded {
                caret = "▼ "
            }
            nodeName = caret + node.name
            if node != m.root {
                nodeName += "/"
            }
            nodeName = dirStyle.Render(nodeName)
        } else {
            nodeName = fileStyle.Render(node.name)
        }

        // Apply bold style if selected
        if m.isSelected(node.path) {
            nodeName = selectedStyle.Render(nodeName)
        }

        // Build the line
        line := fmt.Sprintf("%s %s%s\n",
            cursor,
            prefix.String(),
            nodeName)

        sb.WriteString(line)
    }

    return sb.String()
}

func (m *Model) reloadTree() {
    m.root.children = nil
    m.loadDirectory(m.root)
    m.updateDisplayLines()
}

func (m *Model) findParent(node *Node) *Node {
    var findParentRec func(*Node) *Node

    findParentRec = func(current *Node) *Node {
        if current == nil {
            return nil
        }
        for _, child := range current.children {
            if child == node {
                return current
            }
            if result := findParentRec(child); result != nil {
                return result
            }
        }
        return nil
    }

    return findParentRec(m.root)
}

func (m Model) isSelected(path string) bool {
    for _, selected := range m.SelectedPaths {
        if selected == path {
            return true
        }
    }
    return false
}

func removeString(slice []string, s string) []string {
    for i, v := range slice {
        if v == s {
            return append(slice[:i], slice[i+1:]...)
        }
    }
    return slice
}