package remembering

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"
)



type envT struct {
	allArgs []string
	in      io.Reader
	out     io.Writer
	err     io.Writer
}

// entryT is one line of a profile file.  Its on-disk form is
// "COUNT profile TEXT", and TEXT runs to the end of the line, so
// anything a menu can show --- spaces, UTF-8, leading blanks ---
// round-trips byte-for-byte.
type entryT struct {
	count int
	text  string
}



// splitLines turns raw input into menu lines: the final newline
// doesn't open a last empty line, but blank lines in the middle
// are kept, since a menu tool would show them.
func splitLines(data string) []string {
	if data == "" {
		return []string{}
	}
	return strings.Split(strings.TrimSuffix(data, "\n"), "\n")
}

// parseProfile reads the on-disk "COUNT profile TEXT" lines.
// TEXT is recovered positionally, like the original awk substr()
// did, so blanks inside it survive.  Lines without the "profile"
// tag are dropped; a non-numeric count reads as 0.
func parseProfile(content string) []entryT {
	entries := []entryT{}
	for _, line := range splitLines(content) {
		fields := strings.Fields(line)
		if len(fields) < 2 || fields[1] != "profile" {
			continue
		}
		count, _ := strconv.Atoi(fields[0])
		idx := len(fields[0]) + len(fields[1]) + 2
		text := ""
		if idx <= len(line) {
			text = line[idx:]
		}
		entries = append(entries, entryT{count, text})
	}
	return entries
}

func serializeProfile(entries []entryT) string {
	str := strings.Builder{}
	for _, entry := range entries {
		fmt.Fprintf(
			&str, "%d profile %s\n",
			entry.count, entry.text,
		)
	}
	return str.String()
}

// menuFor ranks the stdin lines: most-picked first, ties in byte
// order, never-picked entries at the bottom.  Counts of duplicate
// profile lines for the same text are summed.  Only stdin defines
// what is offered --- profile entries absent from stdin are
// remembered, not shown.
func menuFor(profile []entryT, lines []string) []entryT {
	counts := map[string]int{}
	for _, entry := range profile {
		counts[entry.text] += entry.count
	}

	menu := make([]entryT, 0, len(lines))
	for _, line := range lines {
		menu = append(menu, entryT{counts[line], line})
	}
	sort.SliceStable(menu, func(i int, j int) bool {
		if menu[i].count != menu[j].count {
			return menu[i].count > menu[j].count
		}
		return menu[i].text < menu[j].text
	})
	return menu
}

func menuText(menu []entryT) string {
	str := strings.Builder{}
	for _, entry := range menu {
		str.WriteString(entry.text)
		str.WriteByte('\n')
	}
	return str.String()
}

// nextProfile is the profile after choice was picked: the union
// of what the profile knew and what stdin offered, in byte order,
// with the pick's count bumped.  A pick the profile has never
// seen is appended at the end, starting at 1; texts offered but
// never picked enter at 0, so the profile accumulates the
// universe of choices it has seen.  Duplicate profile lines for
// the same text collapse to the highest count.
func nextProfile(
	profile []entryT,
	lines []string,
	choice string,
) []entryT {
	counts := map[string]int{}
	for _, entry := range profile {
		prev, seen := counts[entry.text]
		if !seen || entry.count > prev {
			counts[entry.text] = entry.count
		}
	}
	for _, line := range lines {
		_, seen := counts[line]
		if !seen {
			counts[line] = 0
		}
	}

	texts := make([]string, 0, len(counts))
	for text := range counts {
		texts = append(texts, text)
	}
	sort.Strings(texts)

	next := make([]entryT, 0, len(texts)+1)
	found := false
	for _, text := range texts {
		count := counts[text]
		if text == choice {
			count++
			found = true
		}
		next = append(next, entryT{count, text})
	}
	if !found {
		next = append(next, entryT{1, choice})
	}
	return next
}



// profilePath is where a profile lives:
// $XDG_DATA_HOME/remembering/NAME, honouring the XDG default of
// ~/.local/share.  The default profile name is the current
// directory with slashes turned into "!", so each directory gets
// its own ranking without any setup.
func profilePath(name string) (string, error) {
	if name == "" {
		cwd, err := os.Getwd()
		if err != nil {
			return "", err
		}
		name = strings.ReplaceAll(cwd, "/", "!")
	}
	data := os.Getenv("XDG_DATA_HOME")
	if data == "" {
		data = filepath.Join(
			os.Getenv("HOME"), ".local", "share",
		)
	}
	return filepath.Join(data, Name, name), nil
}

func exists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// loadProfile reads the profile, creating an empty one (and its
// directory) on first use, so the file is in place even when this
// run ends up picking nothing.
func loadProfile(path string) ([]entryT, error) {
	if !exists(path) {
		err := os.MkdirAll(filepath.Dir(path), 0755)
		if err != nil {
			return nil, err
		}
		err = os.WriteFile(path, []byte{}, 0644)
		if err != nil {
			return nil, err
		}
	}
	content, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return parseProfile(string(content)), nil
}

// saveProfile writes the new ranking atomically --- a sibling
// .tmp plus rename --- so a crash never leaves a half-written
// profile behind.
func saveProfile(path string, entries []entryT) error {
	tmp := path + ".tmp"
	err := os.WriteFile(
		tmp, []byte(serializeProfile(entries)), 0644,
	)
	if err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

// forwardedStatus mirrors how sh(1) reports a child: its exit
// status as-is, or 128 plus the signal number when it was killed.
func forwardedStatus(err *exec.ExitError) int {
	status, ok := err.Sys().(syscall.WaitStatus)
	if ok && status.Signaled() {
		return 128 + int(status.Signal())
	}
	return err.ExitCode()
}



func usage(w io.Writer) {
	fmt.Fprintf(
		w,
		"Usage:\n  %s [-p PROFILE] -- COMMAND...\n",
		Name,
	)
}

func run(env envT) int {
	flags := flag.NewFlagSet("", flag.ContinueOnError)
	flags.Usage = func() {}
	flags.SetOutput(env.err)
	profileName := flags.String("p", "", "")

	if flags.Parse(env.allArgs[1:]) != nil {
		usage(env.err)
		return 2
	}

	command := flags.Args()
	if len(command) == 0 {
		fmt.Fprintf(env.err, "Missing \"-- COMMAND\"\n")
		usage(env.err)
		return 2
	}

	path, err := profilePath(*profileName)
	if err != nil {
		fmt.Fprintf(env.err, "%s: %v\n", Name, err)
		return 1
	}
	profile, err := loadProfile(path)
	if err != nil {
		fmt.Fprintf(env.err, "%s: %v\n", Name, err)
		return 1
	}

	stdin, err := io.ReadAll(env.in)
	if err != nil {
		fmt.Fprintf(env.err, "%s: %v\n", Name, err)
		return 1
	}
	lines := splitLines(string(stdin))

	// the command is transparent in pipelines: its stderr is
	// passed through (that's where dmenu and fzf draw), its
	// stdout is the pick, and a nonzero exit --- a cancelled
	// menu --- aborts the run with the same status, leaving the
	// profile untouched
	cmd := exec.Command(command[0], command[1:]...)
	cmd.Stdin = strings.NewReader(
		menuText(menuFor(profile, lines)),
	)
	picked := bytes.Buffer{}
	cmd.Stdout = &picked
	cmd.Stderr = env.err
	err = cmd.Run()
	if err != nil {
		exitErr := (*exec.ExitError)(nil)
		if errors.As(err, &exitErr) {
			return forwardedStatus(exitErr)
		}
		fmt.Fprintf(env.err, "%s: %v\n", Name, err)
		if errors.Is(err, fs.ErrPermission) {
			return 126
		}
		return 127
	}

	choice := strings.TrimRight(picked.String(), "\n")
	if choice == "" {
		return 0
	}

	err = saveProfile(path, nextProfile(profile, lines, choice))
	if err != nil {
		fmt.Fprintf(env.err, "%s: %v\n", Name, err)
		return 1
	}
	fmt.Fprintf(env.out, "%s\n", choice)
	return 0
}



func Main() {
	os.Exit(run(envT{
		allArgs: os.Args,
		in:      os.Stdin,
		out:     os.Stdout,
		err:     os.Stderr,
	}))
}