probo/cmd/filter/filter.go

package filter

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"git.andreafazzi.eu/andrea/probo/pkg/store/file"
	"github.com/alecthomas/chroma/quick"
	"github.com/charmbracelet/bubbles/key"
	"github.com/charmbracelet/bubbles/spinner"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/itchyny/gojq"
	foam "github.com/remogatto/sugarfoam"
	"github.com/remogatto/sugarfoam/components/group"
	"github.com/remogatto/sugarfoam/components/header"
	"github.com/remogatto/sugarfoam/components/help"
	"github.com/remogatto/sugarfoam/components/statusbar"
	"github.com/remogatto/sugarfoam/components/textinput"
	"github.com/remogatto/sugarfoam/components/viewport"
	"github.com/remogatto/sugarfoam/layout"
)

type FilterModel struct {
	// UI
	textInput *textinput.Model
	viewport  *viewport.Model
	group     *group.Model
	help      *help.Model
	statusBar *statusbar.Model
	spinner   spinner.Model

	// Layout
	document *layout.Layout

	// Key bindings
	bindings *keyBindings

	// file store
	store  []any
	result []any

	// json
	lastQuery    string
	FilteredJson string
	InputJson    string
	Result       string

	// filter file
	filterFilePath string

	state int

	filterType string
}

func New(path string, filterType string, stdin string) *FilterModel {
	textInput := textinput.New(
		textinput.WithPlaceholder("Write your jq filter here..."),
	)

	viewport := viewport.New()

	group := group.New(
		group.WithItems(textInput, viewport),
		group.WithLayout(
			layout.New(
				layout.WithStyles(&layout.Styles{Container: lipgloss.NewStyle().Padding(1, 0, 1, 0)}),
				layout.WithItem(textInput),
				layout.WithItem(viewport),
			),
		),
	)

	bindings := newBindings(group)
	statusBar := statusbar.New(bindings)

	s := spinner.New(
		spinner.WithStyle(
			lipgloss.NewStyle().Foreground(lipgloss.Color("265"))),
	)
	s.Spinner = spinner.Dot

	header := header.New(
		header.WithContent(
			lipgloss.NewStyle().
				Bold(true).
				Border(lipgloss.NormalBorder(), false, false, true, false).
				Render(filterTypeFormats[filterType]),
		),
	)

	help := help.New(
		bindings,
		help.WithStyles(&foam.Styles{NoBorder: lipgloss.NewStyle().Padding(1, 1)}))

	document := layout.New(
		layout.WithStyles(&layout.Styles{Container: lipgloss.NewStyle().Margin(1)}),
		layout.WithItem(header),
		layout.WithItem(group),
		layout.WithItem(help),
		layout.WithItem(statusBar),
	)

	return &FilterModel{
		textInput:      textInput,
		viewport:       viewport,
		group:          group,
		statusBar:      statusBar,
		spinner:        s,
		document:       document,
		bindings:       bindings,
		help:           help,
		filterType:     filterType,
		filterFilePath: path,
		InputJson:      stdin,
	}

}

func (m *FilterModel) Init() tea.Cmd {
	var cmds []tea.Cmd

	cmds = append(cmds, m.group.Init(), m.loadStore(), m.spinner.Tick)

	m.group.Focus()

	return tea.Batch(cmds...)
}

func (m *FilterModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd

	switch msg := msg.(type) {

	case tea.WindowSizeMsg:
		m.document.SetSize(msg.Width, msg.Height)

	case tea.KeyMsg:
		switch {
		case key.Matches(msg, m.bindings.quit):
			m.FilteredJson = ""
			return m, tea.Quit

		case key.Matches(msg, m.bindings.enter):
			m.marshalJSON()
			return m, tea.Quit
		}

	case storeLoadedMsg:
		cmds = append(cmds, m.handleStoreLoaded(msg))

	case resultMsg:
		cmds = append(cmds, m.handleFiltered(msg))
		m.state = FilterState

	case errorMsg:
		m.handleError(msg)
		m.state = ErrorState
	}

	cmds = m.handleState(msg, cmds)

	return m, tea.Batch(cmds...)
}

func (m *FilterModel) View() string {
	return m.document.View()
}

func (m *FilterModel) marshalJSON() {
	if m.FilteredJson == "" {
		return
	}
	if m.InputJson != "" {
		m.Result = fmt.Sprintf("{%s, \"%s\": %s}", strings.Trim(m.InputJson, "{}"), m.filterType, m.FilteredJson)
	} else {
		var result interface{}

		filtered := fmt.Sprintf("{\"%s\": %s}", m.filterType, m.FilteredJson)
		err := json.Unmarshal([]byte(filtered), &result)
		if err != nil {
			panic(err)
		}

		resultJson, err := json.Marshal(result)
		if err != nil {
			panic(err)
		}
		m.Result = string(resultJson)
	}
}

func (m *FilterModel) showErrorOnStatusBar(err error) {
	m.statusBar.SetContent(
		stateFormats[ErrorState][0],
		fmt.Sprintf(stateFormats[ErrorState][1], err),
		stateFormats[ErrorState][2],
	)
}

func (m *FilterModel) handleError(msg tea.Msg) {
	err := msg.(errorMsg)

	m.statusBar.SetContent(
		stateFormats[ErrorState][0],
		fmt.Sprintf(stateFormats[ErrorState][1], err.error),
		stateFormats[ErrorState][2],
	)
}

func (m *FilterModel) handleStoreLoaded(msg tea.Msg) tea.Cmd {
	return func() tea.Msg {
		storeMsg := msg.(storeLoadedMsg)
		m.store = storeMsg.store
		m.state = FilterState

		if m.filterFilePath != "" {
			jq, err := os.ReadFile(m.filterFilePath)
			if err != nil {
				panic(err)
			}
			m.textInput.SetValue(strings.TrimSpace(string(jq)))
			return m.query(strings.TrimSpace(string(jq)))
		}

		coloredJson, err := toColoredJson(m.store)
		if err != nil {
			return errorMsg{err}
		}

		m.viewport.SetContent(coloredJson)

		return m.query(".")
	}
}

func (m *FilterModel) handleFiltered(msg tea.Msg) tea.Cmd {
	return func() tea.Msg {
		m.result = msg.(resultMsg).result

		coloredJson, err := toColoredJson(m.result)
		if err != nil {
			return errorMsg{err}
		}

		json, err := toJson(m.result)
		if err != nil {
			return errorMsg{err}
		}

		m.FilteredJson = json

		m.viewport.SetContent(coloredJson)

		return nil
	}
}

func (m *FilterModel) handleState(msg tea.Msg, cmds []tea.Cmd) []tea.Cmd {
	_, cmd := m.group.Update(msg)

	if m.state == LoadingStoreState {
		return m.updateSpinner(msg, cmd, cmds)
	}

	cmds = append(cmds, cmd, m.query(m.textInput.Value()))

	if m.state != ErrorState {
		m.statusBar.SetContent(stateFormats[FilterState]...)
	}

	return cmds
}

func (m *FilterModel) updateSpinner(msg tea.Msg, cmd tea.Cmd, cmds []tea.Cmd) []tea.Cmd {
	m.spinner, cmd = m.spinner.Update(msg)

	m.statusBar.SetContent(fmt.Sprintf(stateFormats[m.state][0], m.spinner.View()), stateFormats[m.state][1], stateFormats[m.state][2])

	cmds = append(cmds, cmd)

	return cmds
}

func (m *FilterModel) loadStore() tea.Cmd {
	return func() tea.Msg {
		var jsonStore []byte

		switch m.filterType {
		case "participants":
			pStore, err := file.NewDefaultParticipantFileStore()
			if err != nil {
				panic(err)
			}

			jsonStore, err = pStore.Storer.Json()
			if err != nil {
				panic(err)
			}
		case "quizzes":
			qStore, err := file.NewDefaultQuizFileStore()
			if err != nil {
				panic(err)
			}

			jsonStore, err = qStore.Storer.Json()
			if err != nil {
				panic(err)
			}

		case "responses":
			qStore, err := file.NewDefaultResponseFileStore()
			if err != nil {
				panic(err)
			}

			jsonStore, err = qStore.Storer.Json()
			if err != nil {
				panic(err)
			}
		default:
			panic("Unknown filter type!")
		}

		v := make([]any, 0)

		err := json.Unmarshal(jsonStore, &v)
		if err != nil {
			return errorMsg{err}
		}

		return storeLoadedMsg{v}
	}
}

func (m *FilterModel) query(input string) tea.Cmd {
	return func() tea.Msg {
		if input == m.lastQuery {
			return nil
		}

		if m.state == LoadingStoreState {
			return nil
		}

		m.lastQuery = input

		query, err := gojq.Parse(input)
		if err != nil {
			return errorMsg{fmt.Errorf("jq query parse error: %v", err)}
		}

		var result []string

		iter := query.Run(m.store)
		for {
			v, ok := iter.Next()
			if !ok {
				break
			}

			if err, ok := v.(error); ok {
				return errorMsg{fmt.Errorf("jq query run error: %v", err)}
			}

			b, err := json.MarshalIndent(v, "", "  ")
			if err != nil {
				return errorMsg{err}
			}

			result = append(result, string(b))
		}

		v := make([]any, 0)

		err = json.Unmarshal([]byte(strings.Join(result, "\n")), &v)
		if err != nil {
			return errorMsg{err}
		}

		return resultMsg{v}
	}
}

func toColoredJson(data []any) (string, error) {
	result, err := json.MarshalIndent(data, "", " ")
	if err != nil {
		return "", err
	}

	coloredBytes := make([]byte, 0)
	buffer := bytes.NewBuffer(coloredBytes)

	err = quick.Highlight(buffer, string(result), "json", "terminal16m", "dracula")
	if err != nil {
		panic(err)
	}

	return sanitize(buffer.String()), nil
}

func toJson(data []any) (string, error) {
	result, err := json.Marshal(data)
	if err != nil {
		return "", err
	}
	return string(result), nil
}
func sanitize(text string) string {
	// FIXME: The use of a standard '-' character causes rendering
	// issues within the viewport. Further investigation is
	// required to resolve this problem.
	return strings.Replace(text, "-", "–", -1)
}