package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path/filepath"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/chzyer/readline"
)

// ErrNoPath is returned when 'cd' was called without a second argument.
var ErrNoPath = errors.New("path required")

// version is set at build time via -ldflags "-X main.version=..."
var version = "dev"

// previousDir tracks the previous working directory for 'cd -' support
var previousDir string

// getHomeDir returns the user's home directory
func getHomeDir() string {
	usr, err := user.Current()
	if err != nil {
		// Fallback to HOME environment variable
		return os.Getenv("HOME")
	}
	return usr.HomeDir
}

// expandTilde expands ~ to the home directory and ~user to user's home
func expandTilde(path string) string {
	if len(path) == 0 {
		return path
	}
	
	// Handle ~ at the start
	if strings.HasPrefix(path, "~") {
		home := getHomeDir()
		if home == "" {
			return path // Can't expand, return as-is
		}
		
		// Handle ~ (just home)
		if len(path) == 1 {
			return home
		}
		
		// Handle ~/path or ~user/path
		// For now, just handle ~/path
		if path[1] == '/' {
			return filepath.Join(home, path[2:])
		}
	}
	
	return path
}

// splitWords splits input into words like a (very) simple POSIX shell:
//   - words are separated by unquoted blanks, an unquoted # starts a comment
//   - '...' is taken literally, nothing is expanded inside
//   - "..." expands $VAR and ${VAR}; \ only escapes $ " \ and `
//   - outside quotes \ escapes any character
//   - $VAR, ${VAR} and ${VAR:-default} are expanded outside single quotes;
//     an unset variable is an error, so that e.g. "rm -rf $DIR/x" never runs
//     with an empty $DIR
//   - ~ and ~/... are expanded at the start of an unquoted word
//
// Expanded values are never split into several words.
func splitWords(input string) ([]string, error) {
	var (
		words   []string
		b       strings.Builder
		inWord  bool // a word has been started (even if it is still empty)
		quote   byte
		escaped bool
	)

	endWord := func() {
		if inWord {
			words = append(words, b.String())
		}
		b.Reset()
		inWord = false
	}

	for i := 0; i < len(input); i++ {
		ch := input[i]
		switch {
		case escaped:
			if quote == '"' && !strings.ContainsRune("$\"\\`", rune(ch)) {
				b.WriteByte('\\')
			}
			b.WriteByte(ch)
			inWord = true
			escaped = false
		case quote == '\'':
			if ch == '\'' {
				quote = 0
			} else {
				b.WriteByte(ch)
			}
		case ch == '\\':
			escaped = true
		case ch == '$':
			value, n, err := expandVar(input[i:])
			if err != nil {
				return nil, err
			}
			b.WriteString(value)
			// A literal "$" counts as content, an expansion that is empty does
			// not (unless quoted), like in sh
			if n == 1 || value != "" {
				inWord = true
			}
			i += n - 1
		case quote == '"':
			if ch == '"' {
				quote = 0
			} else {
				b.WriteByte(ch)
			}
		case ch == '"' || ch == '\'':
			quote = ch
			inWord = true
		case ch == ' ' || ch == '\t' || ch == '\n':
			endWord()
		case ch == '#' && !inWord && b.Len() == 0:
			// Comment until the end of the line
			i = len(input)
		case ch == '~' && !inWord && b.Len() == 0 &&
			(i+1 == len(input) || strings.IndexByte("/ \t\n", input[i+1]) >= 0):
			home := getHomeDir()
			if home == "" {
				home = "~"
			}
			b.WriteString(home)
			inWord = true
		default:
			b.WriteByte(ch)
			inWord = true
		}
	}

	switch {
	case escaped:
		return nil, errors.New("unexpected end of input after \\")
	case quote != 0:
		return nil, fmt.Errorf("unterminated %c quote", quote)
	}
	endWord()
	return words, nil
}

// expandVar expands the variable reference at the start of s (which starts
// with "$"). It returns the value and the number of bytes consumed. A "$" that
// is not followed by a variable name is returned literally.
func expandVar(s string) (string, int, error) {
	if len(s) > 1 && s[1] == '{' {
		end := strings.IndexByte(s, '}')
		if end < 0 {
			return "", 0, errors.New("missing } in variable reference")
		}
		name, def, hasDef := strings.Cut(s[2:end], ":-")
		if !isVarName(name) {
			return "", 0, fmt.Errorf("bad substitution: %s", s[:end+1])
		}
		value, ok := os.LookupEnv(name)
		if hasDef && value == "" {
			return def, end + 1, nil
		}
		if !ok {
			return "", 0, fmt.Errorf("%s: unbound variable", name)
		}
		return value, end + 1, nil
	}

	n := 1
	for n < len(s) && isVarNameByte(s[n], n == 1) {
		n++
	}
	if n == 1 {
		return "$", 1, nil
	}
	value, ok := os.LookupEnv(s[1:n])
	if !ok {
		return "", 0, fmt.Errorf("%s: unbound variable", s[1:n])
	}
	return value, n, nil
}

func isVarName(name string) bool {
	if name == "" {
		return false
	}
	for i := 0; i < len(name); i++ {
		if !isVarNameByte(name[i], i == 0) {
			return false
		}
	}
	return true
}

func isVarNameByte(c byte, first bool) bool {
	return c == '_' || c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' ||
		!first && c >= '0' && c <= '9'
}

// sanitize replaces non-printable characters (e.g. terminal escape sequences
// in directory names) with a visible escape so they cannot manipulate the
// terminal.
func sanitize(s string) string {
	var b strings.Builder
	for i := 0; i < len(s); {
		r, size := utf8.DecodeRuneInString(s[i:])
		switch {
		case r == utf8.RuneError && size == 1:
			fmt.Fprintf(&b, "\\x%02x", s[i])
		case !unicode.IsPrint(r):
			if r < 0x100 {
				fmt.Fprintf(&b, "\\x%02x", r)
			} else {
				fmt.Fprintf(&b, "\\u%04x", r)
			}
		default:
			b.WriteRune(r)
		}
		i += size
	}
	return b.String()
}

// isPrintable reports whether s only consists of printable characters
func isPrintable(s string) bool {
	return sanitize(s) == s
}

// historyFile returns the path of the history file: $GOSH_HISTFILE if set,
// otherwise $XDG_STATE_HOME/gosh/history (default ~/.local/state/gosh/history).
// An empty string disables the history.
func historyFile() string {
	if path, ok := os.LookupEnv("GOSH_HISTFILE"); ok {
		return path
	}
	state := os.Getenv("XDG_STATE_HOME")
	if state == "" {
		home := getHomeDir()
		if home == "" {
			return ""
		}
		state = filepath.Join(home, ".local", "state")
	}
	return filepath.Join(state, "gosh", "history")
}

// prepareHistoryFile makes sure that the history file and its directory are
// only accessible by the current user. readline itself creates the file (and
// a temporary file next to it) with mode 0666 & ~umask.
func prepareHistoryFile(path string) error {
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0700); err != nil {
		return err
	}
	// Only restrict the directory if it is ours (and not e.g. $HOME itself)
	if filepath.Base(dir) == "gosh" {
		if err := os.Chmod(dir, 0700); err != nil {
			return err
		}
	}
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	return f.Chmod(0600)
}

// getPrompt returns a dynamic prompt string like "user@host:dir $"
func getPrompt() string {
	// Get current working directory
	wd, err := os.Getwd()
	if err != nil {
		wd = "?"
	}
	
	// Get username
	usr, err := user.Current()
	username := "?"
	if err == nil && usr.Username != "" {
		username = usr.Username
	}
	
	// Get hostname
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "unknown"
	}
	
	// Replace the home directory with ~
	home := getHomeDir()
	if home != "" && (wd == home || strings.HasPrefix(wd, home+"/")) {
		wd = "~" + strings.TrimPrefix(wd, home)
	}
	
	var colors *promptColors
	if useColor() {
		c := parseColors(os.Getenv("GOSH_COLORS"))
		colors = &c
	}
	return buildPrompt(username, hostname, wd, os.Geteuid() == 0, colors)
}

// promptColors holds the SGR parameters (e.g. "1;32") for the parts of the
// prompt. An empty value means no color for that part.
type promptColors struct {
	user string // user@host
	root string // user@host when running as root
	path string // working directory
}

// defaultColors: bold green user, bold red root, bold blue directory
var defaultColors = promptColors{user: "1;32", root: "1;31", path: "1;34"}

// parseColors reads a GOSH_COLORS value like "user=1;33:root=1;31:path=36"
// and overrides the corresponding defaults. Unknown keys and values that are
// not plain SGR parameters (digits and ';') are ignored, so the variable
// cannot be used to write arbitrary escape sequences to the terminal.
func parseColors(spec string) promptColors {
	colors := defaultColors
	for _, entry := range strings.Split(spec, ":") {
		key, value, ok := strings.Cut(entry, "=")
		if !ok || !isSGR(value) {
			continue
		}
		switch key {
		case "user":
			colors.user = value
		case "root":
			colors.root = value
		case "path":
			colors.path = value
		}
	}
	return colors
}

// isSGR reports whether s consists only of digits and ';'. The empty string
// is allowed and disables the color.
func isSGR(s string) bool {
	for _, r := range s {
		if (r < '0' || r > '9') && r != ';' {
			return false
		}
	}
	return true
}

// colorize wraps text in the SGR sequence for sgr, or returns it unchanged if
// sgr is empty.
func colorize(text, sgr string) string {
	if sgr == "" {
		return text
	}
	return "\x1b[" + sgr + "m" + text + "\x1b[0m"
}

// buildPrompt formats the prompt as user@host:path $. The parts are sanitized
// before the color codes are added, so only our own escape sequences end up
// in the prompt. colors == nil disables colors.
func buildPrompt(username, hostname, path string, root bool, colors *promptColors) string {
	userHost := sanitize(username + "@" + hostname)
	path = sanitize(path)
	if colors == nil {
		return userHost + ":" + path + " $ "
	}
	userColor := colors.user
	if root {
		userColor = colors.root
	}
	return colorize(userHost, userColor) + ":" + colorize(path, colors.path) + " $ "
}

// useColor reports whether the prompt should be colored: stdout must be a
// terminal, TERM must not be "dumb" and NO_COLOR (https://no-color.org) must
// not be set.
func useColor() bool {
	if os.Getenv("NO_COLOR") != "" || os.Getenv("TERM") == "dumb" {
		return false
	}
	return readline.IsTerminal(int(os.Stdout.Fd()))
}

// shellCompleter implements readline.AutoCompleter for contextual completion
type shellCompleter struct {
	commands []string
}

func newShellCompleter() *shellCompleter {
	return &shellCompleter{
		commands: []string{"cd", "pwd", "echo", "exit", "quit", "set", "export", "help"},
	}
}

// Do implements readline.AutoCompleter interface.
// readline expects the candidates to contain only the text that still has to
// be inserted after the cursor, and the returned length to be the number of
// runes of the current word that are already typed.
func (c *shellCompleter) Do(line []rune, pos int) ([][]rune, int) {
	word, start, quote := currentWord(string(line[:pos]))
	isFirstWord := strings.TrimSpace(string(line[:start])) == ""

	// Plain "~" -> "~/"
	if word == "~" {
		return [][]rune{[]rune("/")}, 1
	}

	if isFirstWord && !strings.Contains(word, "/") && !strings.HasPrefix(word, "~") {
		// Complete commands
		var candidates [][]rune
		seen := make(map[string]bool)
		for _, cmd := range append(c.commands, pathExecutables(word)...) {
			if strings.HasPrefix(cmd, word) && !seen[cmd] {
				seen[cmd] = true
				candidates = append(candidates, []rune(escapeWord(cmd[len(word):])+" "))
			}
		}
		return candidates, len([]rune(word))
	}

	// Complete file paths; only the part after the last slash is completed
	pattern := word[strings.LastIndex(word, "/")+1:]
	matches := completeFilePath(string(line[:pos]), word)
	var candidates [][]rune
	for _, m := range matches {
		rest := m[len(pattern):]
		if quote == 0 {
			rest = escapeWord(rest)
		}
		candidates = append(candidates, []rune(rest))
	}
	return candidates, len([]rune(pattern))
}

// pathExecutables returns the names of all executables in $PATH starting
// with prefix
func pathExecutables(prefix string) []string {
	var names []string
	for _, dir := range filepath.SplitList(os.Getenv("PATH")) {
		if dir == "" {
			dir = "."
		}
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue
		}
		for _, entry := range entries {
			name := entry.Name()
			if !strings.HasPrefix(name, prefix) || !isPrintable(name) {
				continue
			}
			// Follow symlinks, only regular files with an executable bit count
			info, err := os.Stat(filepath.Join(dir, name))
			if err != nil || !info.Mode().IsRegular() || info.Mode().Perm()&0111 == 0 {
				continue
			}
			names = append(names, name)
		}
	}
	return names
}

// currentWord returns the (unescaped, unquoted) word the cursor is in, the
// byte offset where it starts in s and the quote character that is still open
// (0 if none).
func currentWord(s string) (word string, start int, quote byte) {
	var b strings.Builder
	escaped := false
	for i := 0; i < len(s); i++ {
		ch := s[i]
		switch {
		case escaped:
			b.WriteByte(ch)
			escaped = false
		case ch == '\\' && quote != '\'':
			escaped = true
		case quote != 0:
			if ch == quote {
				quote = 0
			} else {
				b.WriteByte(ch)
			}
		case ch == '"' || ch == '\'':
			quote = ch
		case ch == ' ' || ch == '\t':
			b.Reset()
			start = i + 1
		default:
			b.WriteByte(ch)
		}
	}
	return b.String(), start, quote
}

// escapeWord escapes characters that would otherwise be interpreted by shlex
func escapeWord(s string) string {
	var b strings.Builder
	for _, r := range s {
		if strings.ContainsRune(" \t\\'\"", r) {
			b.WriteByte('\\')
		}
		b.WriteRune(r)
	}
	return b.String()
}

// completeFilePath handles file path completion.
// prefix is the unquoted path typed so far. It returns the names (with a
// trailing "/" for directories) of all entries in the directory part of
// prefix that start with the part after the last slash.
func completeFilePath(line, prefix string) []string {
	// Split textually so that "." or "~" are not normalized away
	slash := strings.LastIndex(prefix, "/")
	dir, pattern := prefix[:slash+1], prefix[slash+1:]

	switch {
	case dir == "":
		dir = "."
	case strings.HasPrefix(dir, "~"):
		dir = expandTilde(dir)
	}

	// Read the directory contents (relative paths resolve against the cwd)
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil
	}

	// Collect matching entries
	var matches []string
	for _, entry := range entries {
		name := entry.Name()
		if !strings.HasPrefix(name, pattern) {
			continue
		}
		// Hide dotfiles unless explicitly asked for
		if strings.HasPrefix(name, ".") && !strings.HasPrefix(pattern, ".") {
			continue
		}
		// Never insert control characters (e.g. escape sequences) into the line
		if !isPrintable(name) {
			continue
		}
		// Follow symlinks to decide whether this is a directory
		if info, err := os.Stat(filepath.Join(dir, name)); err == nil && info.IsDir() {
			name += "/"
		}
		matches = append(matches, name)
	}

	return matches
}

func main() {
	if len(os.Args) > 1 && (os.Args[1] == "--version" || os.Args[1] == "-V") {
		fmt.Println("gosh", version)
		return
	}

	// Ctrl+C sends SIGINT to the whole foreground process group. Catch it so
	// that it only stops the running command and not the shell. (signal.Ignore
	// would be inherited by the child processes.)
	signal.Notify(make(chan os.Signal, 1), os.Interrupt)

	histFile := historyFile()
	if histFile != "" {
		if err := prepareHistoryFile(histFile); err != nil {
			fmt.Fprintln(os.Stderr, "History disabled:", err)
			histFile = ""
		}
	}

	// Create readline instance
	l, err := readline.NewEx(&readline.Config{
		Prompt:                 getPrompt(),
		HistoryFile:            histFile,
		DisableAutoSaveHistory: true,
		AutoComplete:           newShellCompleter(),
		InterruptPrompt:        "^C",
		EOFPrompt:              "exit",
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, "Error initializing readline:", err)
		os.Exit(1)
	}
	defer l.Close()

	// Main loop
	for {
		// Update prompt dynamically (in case directory changed)
		l.SetPrompt(getPrompt())
		
		line, err := l.Readline()
		if err != nil {
			// Ctrl+C discards the current line
			if err == readline.ErrInterrupt {
				continue
			}
			// Handle EOF (Ctrl+D)
			if err == io.EOF {
				fmt.Println()
				break
			}
			fmt.Fprintln(os.Stderr, "Error reading input:", err)
			continue
		}

		// Skip empty input
		if strings.TrimSpace(line) == "" {
			continue
		}

		// Add to history, unless the line starts with a space (like bash's
		// HISTCONTROL=ignorespace) so that secrets can be kept out of it
		if !strings.HasPrefix(line, " ") {
			l.SaveHistory(line)
		}

		// Execute the input
		if err := execInput(line); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}
	}
}

func execInput(input string) error {
	// Split the input into the command and its arguments, expanding
	// variables and ~ on the way.
	args, err := splitWords(input)
	if err != nil {
		return fmt.Errorf("invalid input: %w", err)
	}
	if len(args) == 0 {
		return nil
	}

	// Check for built-in commands.
	switch args[0] {
	case "cd":
		// Save current directory before changing (for 'cd -')
		currentDir, _ := os.Getwd()
		
		// Handle different cd cases
		if len(args) < 2 {
			// cd with no args: go to home directory
			previousDir = currentDir
			return os.Chdir(getHomeDir())
		}
		
		path := args[1]
		
		// Handle 'cd -' : go to previous directory
		if path == "-" {
			if previousDir == "" {
				return fmt.Errorf("no previous directory")
			}
			// Save current directory as the new previous
			// so that next 'cd -' will come back here
			prevDir := previousDir
			previousDir = currentDir
			return os.Chdir(prevDir)
		}
		
		// Save current directory as previous before changing
		previousDir = currentDir
		
		// Change the directory and return the error.
		return os.Chdir(path)
	case "pwd":
		wd, err := os.Getwd()
		fmt.Println(wd)
		return err
	case "exit":
		// exit with optional status code
		if len(args) > 1 {
			var exitCode int
			fmt.Sscanf(args[1], "%d", &exitCode)
			os.Exit(exitCode)
		}
		os.Exit(0)
	case "quit":
		// quit with optional status code
		if len(args) > 1 {
			var exitCode int
			fmt.Sscanf(args[1], "%d", &exitCode)
			os.Exit(exitCode)
		}
		os.Exit(0)
	case "set":
		// set VAR=value - set environment variable
		if len(args) < 2 {
			return fmt.Errorf("usage: set VAR=value")
		}
		// Parse VAR=value
		parts := strings.SplitN(args[1], "=", 2)
		if len(parts) != 2 {
			return fmt.Errorf("usage: set VAR=value")
		}
		os.Setenv(parts[0], parts[1])
		return nil
	case "export":
		// export VAR=value - same as set but also marks for child processes
		if len(args) < 2 {
			return fmt.Errorf("usage: export VAR=value")
		}
		// Parse VAR=value
		parts := strings.SplitN(args[1], "=", 2)
		if len(parts) != 2 {
			return fmt.Errorf("usage: export VAR=value")
		}
		os.Setenv(parts[0], parts[1])
		return nil
	case "help":
		fmt.Println(`gosh - Go Shell

Built-in commands:
  cd [dir]       - Change directory (cd with no args goes to home)
  pwd           - Print working directory
  echo [args]   - Echo arguments
  exit [code]   - Exit shell with optional status code
  quit [code]   - Quit shell with optional status code
  set VAR=val   - Set environment variable
  export VAR=val - Export environment variable
  help          - Show this help

Features:
  - Argument parsing with quotes and escapes ('...' is literal,
    "..." expands variables, \ escapes the next character)
  - Ctrl+D to exit, Ctrl+C to cancel the current line or command
  - Command history (up/down arrows), stored in
    $XDG_STATE_HOME/gosh/history (default ~/.local/state/gosh/history,
    override with $GOSH_HISTFILE); lines starting with a space are not saved
  - Tab completion for commands and file paths
  - Tilde expansion (~, ~/path) at the start of unquoted words
  - Colored prompt (user@host green, red for root; directory blue),
    configurable with $GOSH_COLORS, e.g. GOSH_COLORS='user=1;33:path=36'
    (keys user, root, path; values are ANSI SGR codes, empty = no color);
    disabled with $NO_COLOR, TERM=dumb or when output is not a terminal
  - Environment variable expansion ($VAR, ${VAR}, ${VAR:-default});
    using an unset variable is an error, use ${VAR:-} for an empty default`)
		return nil
	}

	// Prepare the command to execute.
	cmd := exec.Command(args[0], args[1:]...)

	// Connect the command to the terminal. Without Stdin the command would
	// read from /dev/null, so interactive programs (su, vim, ssh, ...) would
	// get EOF right away. readline only reads from stdin while a line is
	// being edited, so it does not steal input from the command.
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// Execute the command return the error.
	return cmd.Run()
}