gguf-editor/main.go

package main

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

	"github.com/charmbracelet/bubbles/table"
	"github.com/charmbracelet/bubbles/textinput"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

// ── Styles ──────────────────────────────────────────────────────────

var (
	titleStyle = lipgloss.NewStyle().
			Bold(true).
			Foreground(lipgloss.Color("#FAFAFA")).
			Background(lipgloss.Color("#7C3AED")).
			Padding(0, 2)

	infoStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#A78BFA")).
			Bold(true)

	statusOkStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#10B981"))

	statusErrStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#EF4444"))

	helpStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#6B7280"))

	overlayStyle = lipgloss.NewStyle().
			Border(lipgloss.RoundedBorder()).
			BorderForeground(lipgloss.Color("#7C3AED")).
			Padding(1, 2).
			Width(60)

	overlayTitleStyle = lipgloss.NewStyle().
				Bold(true).
				Foreground(lipgloss.Color("#C4B5FD")).
				MarginBottom(1)

	dimStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#4B5563"))

	typeActiveStyle = lipgloss.NewStyle().
			Bold(true).
			Foreground(lipgloss.Color("#FAFAFA")).
			Background(lipgloss.Color("#7C3AED")).
			Padding(0, 1)

	typeInactiveStyle = lipgloss.NewStyle().
				Foreground(lipgloss.Color("#9CA3AF")).
				Padding(0, 1)

	modifiedStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#F59E0B")).
			Bold(true)
)

// ── App State ───────────────────────────────────────────────────────

type appState int

const (
	stateTable   appState = iota
	stateEdit             // editing a value
	stateAdd              // adding new KV (multi-step)
	stateConfirm          // confirm dialog (save/delete)
)

type confirmAction int

const (
	confirmSave   confirmAction = iota
	confirmDelete
)

// ── Model ───────────────────────────────────────────────────────────

type model struct {
	state    appState
	gguf     *GGUFFile
	filePath string
	dirty    bool

	// table
	table  table.Model
	width  int
	height int

	// edit
	editInput textinput.Model
	editIdx   int

	// add (step 0: key, step 1: type picker, step 2: value)
	addStep  int
	addKey   string
	addType  int // index into ScalarTypes
	addInput textinput.Model

	// confirm
	confirmMsg    string
	confirmAction confirmAction

	// status
	statusMsg   string
	statusIsErr bool
}

func initialModel(g *GGUFFile, path string) model {
	ti := textinput.New()
	ti.CharLimit = 4096

	ai := textinput.New()
	ai.CharLimit = 4096

	m := model{
		state:     stateTable,
		gguf:      g,
		filePath:  path,
		editInput: ti,
		addInput:  ai,
		statusMsg: fmt.Sprintf("Loaded %d metadata entries, %d tensors (%s)",
			len(g.Metadata), g.TensorCount, humanSize(g.fileSize)),
	}

	m.table = m.buildTable(80, 30)
	return m
}

func (m model) buildTable(width, height int) table.Model {
	keyW := max(20, (width-8)*45/100)
	typeW := 14
	valW := max(10, width-keyW-typeW-8)

	cols := []table.Column{
		{Title: "#", Width: 4},
		{Title: "Key", Width: keyW},
		{Title: "Type", Width: typeW},
		{Title: "Value", Width: valW},
	}

	rows := make([]table.Row, len(m.gguf.Metadata))
	for i, kv := range m.gguf.Metadata {
		val := kv.Value.DisplayString()
		if len(val) > valW-2 {
			val = val[:valW-5] + "..."
		}
		rows[i] = table.Row{
			fmt.Sprintf("%d", i),
			kv.Key,
			kv.Value.TypeName(),
			val,
		}
	}

	tblHeight := height - 7
	if tblHeight < 5 {
		tblHeight = 5
	}

	t := table.New(
		table.WithColumns(cols),
		table.WithRows(rows),
		table.WithFocused(true),
		table.WithHeight(tblHeight),
	)

	s := table.DefaultStyles()
	s.Header = s.Header.
		BorderStyle(lipgloss.NormalBorder()).
		BorderForeground(lipgloss.Color("#7C3AED")).
		BorderBottom(true).
		Bold(true).
		Foreground(lipgloss.Color("#C4B5FD"))
	s.Selected = s.Selected.
		Foreground(lipgloss.Color("#FAFAFA")).
		Background(lipgloss.Color("#4C1D95")).
		Bold(false)
	t.SetStyles(s)

	return t
}

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

// ── Update ──────────────────────────────────────────────────────────

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		cur := m.table.Cursor()
		m.table = m.buildTable(m.width, m.height)
		m.table.SetCursor(cur)
		return m, nil
	}

	switch m.state {
	case stateTable:
		return m.updateTable(msg)
	case stateEdit:
		return m.updateEdit(msg)
	case stateAdd:
		return m.updateAdd(msg)
	case stateConfirm:
		return m.updateConfirm(msg)
	}
	return m, nil
}

func (m model) updateTable(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "q", "ctrl+c":
			return m, tea.Quit
		case "e", "enter":
			if len(m.gguf.Metadata) == 0 {
				return m, nil
			}
			idx := m.table.Cursor()
			kv := m.gguf.Metadata[idx]
			if kv.Value.Type == TypeArray {
				m.statusMsg = "Array values cannot be edited directly"
				m.statusIsErr = true
				return m, nil
			}
			m.state = stateEdit
			m.editIdx = idx
			m.editInput.SetValue(kv.Value.DisplayString())
			m.editInput.Focus()
			m.editInput.CursorEnd()
			return m, textinput.Blink
		case "a":
			m.state = stateAdd
			m.addStep = 0
			m.addType = 0
			m.addKey = ""
			m.addInput.SetValue("")
			m.addInput.Placeholder = "metadata.key.name"
			m.addInput.Focus()
			return m, textinput.Blink
		case "d":
			if len(m.gguf.Metadata) == 0 {
				return m, nil
			}
			idx := m.table.Cursor()
			m.state = stateConfirm
			m.confirmAction = confirmDelete
			m.confirmMsg = fmt.Sprintf("Delete key %q?", m.gguf.Metadata[idx].Key)
			return m, nil
		case "s":
			outPath := m.outputPath()
			m.state = stateConfirm
			m.confirmAction = confirmSave
			m.confirmMsg = fmt.Sprintf("Save as %s?", filepath.Base(outPath))
			return m, nil
		}
	}

	var cmd tea.Cmd
	m.table, cmd = m.table.Update(msg)
	return m, cmd
}

func (m model) updateEdit(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "enter":
			val, err := ParseValueFromString(m.editInput.Value(), m.gguf.Metadata[m.editIdx].Value.Type)
			if err != nil {
				m.statusMsg = fmt.Sprintf("Invalid value: %v", err)
				m.statusIsErr = true
				return m, nil
			}
			m.gguf.Metadata[m.editIdx].Value = val
			m.dirty = true
			m.statusMsg = fmt.Sprintf("Updated %q", m.gguf.Metadata[m.editIdx].Key)
			m.statusIsErr = false
			m.state = stateTable
			m.rebuildRows()
			return m, nil
		case "esc":
			m.state = stateTable
			return m, nil
		}
	}

	var cmd tea.Cmd
	m.editInput, cmd = m.editInput.Update(msg)
	return m, cmd
}

func (m model) updateAdd(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "esc":
			m.state = stateTable
			return m, nil
		case "enter":
			switch m.addStep {
			case 0: // key entered
				key := strings.TrimSpace(m.addInput.Value())
				if key == "" {
					m.statusMsg = "Key cannot be empty"
					m.statusIsErr = true
					return m, nil
				}
				m.addKey = key
				m.addStep = 1
				return m, nil
			case 1: // type selected
				m.addStep = 2
				m.addInput.SetValue("")
				m.addInput.Placeholder = fmt.Sprintf("Enter %s value", ScalarTypes[m.addType])
				m.addInput.Focus()
				return m, textinput.Blink
			case 2: // value entered
				t := ScalarTypes[m.addType]
				val, err := ParseValueFromString(m.addInput.Value(), t)
				if err != nil {
					m.statusMsg = fmt.Sprintf("Invalid %s: %v", t, err)
					m.statusIsErr = true
					return m, nil
				}
				m.gguf.Metadata = append(m.gguf.Metadata, KVPair{Key: m.addKey, Value: val})
				m.dirty = true
				m.statusMsg = fmt.Sprintf("Added %q (%s)", m.addKey, t)
				m.statusIsErr = false
				m.state = stateTable
				m.rebuildRows()
				return m, nil
			}
		case "left", "shift+tab":
			if m.addStep == 1 {
				m.addType = (m.addType - 1 + len(ScalarTypes)) % len(ScalarTypes)
				return m, nil
			}
		case "right", "tab":
			if m.addStep == 1 {
				m.addType = (m.addType + 1) % len(ScalarTypes)
				return m, nil
			}
		}
	}

	if m.addStep != 1 {
		var cmd tea.Cmd
		m.addInput, cmd = m.addInput.Update(msg)
		return m, cmd
	}
	return m, nil
}

func (m model) updateConfirm(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "y", "enter":
			switch m.confirmAction {
			case confirmSave:
				outPath := m.outputPath()
				if err := m.gguf.Save(outPath); err != nil {
					m.statusMsg = fmt.Sprintf("Save failed: %v", err)
					m.statusIsErr = true
				} else {
					m.statusMsg = fmt.Sprintf("Saved to %s", filepath.Base(outPath))
					m.statusIsErr = false
					m.dirty = false
				}
			case confirmDelete:
				idx := m.table.Cursor()
				key := m.gguf.Metadata[idx].Key
				m.gguf.Metadata = append(m.gguf.Metadata[:idx], m.gguf.Metadata[idx+1:]...)
				m.dirty = true
				m.statusMsg = fmt.Sprintf("Deleted %q", key)
				m.statusIsErr = false
				m.rebuildRows()
				if idx >= len(m.gguf.Metadata) && idx > 0 {
					m.table.SetCursor(idx - 1)
				}
			}
			m.state = stateTable
			return m, nil
		case "n", "esc":
			m.state = stateTable
			return m, nil
		}
	}
	return m, nil
}

// ── View ────────────────────────────────────────────────────────────

func (m model) View() string {
	var b strings.Builder

	// Title bar
	title := titleStyle.Render(" GGUF Editor ")
	fname := infoStyle.Render(filepath.Base(m.filePath))
	ver := infoStyle.Render(fmt.Sprintf("v%d", m.gguf.Version))
	dirtyMark := ""
	if m.dirty {
		dirtyMark = modifiedStyle.Render(" [modified]")
	}
	b.WriteString(title + " " + fname + " " + ver + dirtyMark + "\n\n")

	// Table
	b.WriteString(m.table.View())
	b.WriteString("\n\n")

	// Status line
	if m.statusIsErr {
		b.WriteString(statusErrStyle.Render("  " + m.statusMsg))
	} else {
		b.WriteString(statusOkStyle.Render("  " + m.statusMsg))
	}
	b.WriteString("\n")

	// Help bar
	help := "  e/Enter:edit  a:add  d:delete  s:save  q:quit"
	b.WriteString(helpStyle.Render(help))

	view := b.String()

	// Overlay rendering
	switch m.state {
	case stateEdit:
		kv := m.gguf.Metadata[m.editIdx]
		overlay := overlayTitleStyle.Render(fmt.Sprintf("Edit: %s (%s)", kv.Key, kv.Value.TypeName())) +
			"\n" + m.editInput.View() +
			"\n\n" + dimStyle.Render("Enter: save  Esc: cancel")
		view = m.renderOverlay(view, overlay)
	case stateAdd:
		var overlay string
		switch m.addStep {
		case 0:
			overlay = overlayTitleStyle.Render("Add Metadata: Enter Key") +
				"\n" + m.addInput.View() +
				"\n\n" + dimStyle.Render("Enter: next  Esc: cancel")
		case 1:
			overlay = overlayTitleStyle.Render(fmt.Sprintf("Add: %s - Select Type", m.addKey)) +
				"\n" + m.renderTypePicker() +
				"\n\n" + dimStyle.Render("Left/Right: select  Enter: next  Esc: cancel")
		case 2:
			overlay = overlayTitleStyle.Render(fmt.Sprintf("Add: %s (%s) - Enter Value", m.addKey, ScalarTypes[m.addType])) +
				"\n" + m.addInput.View() +
				"\n\n" + dimStyle.Render("Enter: add  Esc: cancel")
		}
		view = m.renderOverlay(view, overlay)
	case stateConfirm:
		overlay := overlayTitleStyle.Render(m.confirmMsg) +
			"\n\n" + dimStyle.Render("y/Enter: yes  n/Esc: no")
		view = m.renderOverlay(view, overlay)
	}

	return view
}

func (m model) renderTypePicker() string {
	var parts []string
	for i, t := range ScalarTypes {
		if i == m.addType {
			parts = append(parts, typeActiveStyle.Render(t.String()))
		} else {
			parts = append(parts, typeInactiveStyle.Render(t.String()))
		}
	}
	// Wrap into rows of ~6
	var rows []string
	for i := 0; i < len(parts); i += 6 {
		end := i + 6
		if end > len(parts) {
			end = len(parts)
		}
		rows = append(rows, strings.Join(parts[i:end], " "))
	}
	return strings.Join(rows, "\n")
}

func (m model) renderOverlay(base, content string) string {
	box := overlayStyle.Render(content)
	boxLines := strings.Split(box, "\n")
	baseLines := strings.Split(base, "\n")

	// Center the overlay vertically and horizontally
	startY := (m.height - len(boxLines)) / 2
	if startY < 0 {
		startY = 0
	}
	startX := (m.width - lipgloss.Width(box)) / 2
	if startX < 0 {
		startX = 0
	}

	// Dim the base
	for i := range baseLines {
		baseLines[i] = dimStyle.Render(stripAnsi(baseLines[i]))
	}

	// Overlay the box
	for i, line := range boxLines {
		row := startY + i
		if row >= 0 && row < len(baseLines) {
			padded := baseLines[row]
			if startX > lipgloss.Width(padded) {
				padded += strings.Repeat(" ", startX-lipgloss.Width(padded))
			}
			// Replace the section
			baseLines[row] = truncateToWidth(padded, startX) + line
		}
	}

	return strings.Join(baseLines, "\n")
}

// ── Helpers ─────────────────────────────────────────────────────────

func (m *model) rebuildRows() {
	cur := m.table.Cursor()
	m.table = m.buildTable(m.width, m.height)
	if cur >= len(m.gguf.Metadata) {
		cur = len(m.gguf.Metadata) - 1
	}
	if cur < 0 {
		cur = 0
	}
	m.table.SetCursor(cur)
}

func (m model) outputPath() string {
	dir := filepath.Dir(m.filePath)
	base := filepath.Base(m.filePath)
	return filepath.Join(dir, "modified_"+base)
}

// stripAnsi removes ANSI escape sequences for width calculations.
func stripAnsi(s string) string {
	var b strings.Builder
	inEsc := false
	for _, r := range s {
		if r == '\033' {
			inEsc = true
			continue
		}
		if inEsc {
			if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') {
				inEsc = false
			}
			continue
		}
		b.WriteRune(r)
	}
	return b.String()
}

// truncateToWidth takes the first n visible characters of a string,
// preserving any ANSI escapes that precede them.
func truncateToWidth(s string, n int) string {
	var b strings.Builder
	visible := 0
	inEsc := false
	for _, r := range s {
		if visible >= n && !inEsc {
			break
		}
		b.WriteRune(r)
		if r == '\033' {
			inEsc = true
			continue
		}
		if inEsc {
			if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') {
				inEsc = false
			}
			continue
		}
		visible++
	}
	// Pad if needed
	for visible < n {
		b.WriteByte(' ')
		visible++
	}
	return b.String()
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// ── Main ────────────────────────────────────────────────────────────

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s <file.gguf>\n", os.Args[0])
		os.Exit(1)
	}

	path := os.Args[1]
	g, err := ParseGGUF(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	absPath, err := filepath.Abs(path)
	if err != nil {
		absPath = path
	}

	m := initialModel(g, absPath)
	p := tea.NewProgram(m, tea.WithAltScreen())
	if _, err := p.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}