ask/main.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const defaultBaseURL = "http://10.0.2.145:8090/v1"

type chatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
}

type chatResponse struct {
	Choices []struct {
		Message chatMessage `json:"message"`
	} `json:"choices"`
}

func main() {
	help := flag.Bool("h", false, "Show help")
	history := flag.Bool("H", false, "Answer questions about shell history")
	question := flag.Bool("q", false, "Answer general questions, render with glow")
	paginate := flag.Bool("p", false, "Paginate output using less")
	flag.Parse()

	if *help {
		printHelp()
		return
	}

	args := flag.Args()
	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "Error: no question provided")
		os.Exit(1)
	}
	query := strings.Join(args, " ")

	if *history {
		handleHistory(query, *paginate)
		return
	}

	if *question {
		handleQuestion(query, *paginate)
		return
	}

	handleCommand(query, *paginate)
}

func printHelp() {
	fmt.Println(`Usage: ask [options] <question>

Options:
  -h    Show this help
  -H    Answer questions about your shell history
  -q    Answer general command questions (output rendered with glow)
  -p    Paginate output using less

Examples:
  ask how to find large files
  ask -H what command did I use for docker yesterday
  ask -q how does rsync work
  ask -p how to compress a directory`)
}

func getBaseURL() string {
	if v := os.Getenv("OPENAI_BASE_URL"); v != "" {
		return strings.TrimRight(v, "/")
	}
	return defaultBaseURL
}

func getAPIKey() string {
	return os.Getenv("OPENAI_API_KEY")
}

func getModel() string {
	if v := os.Getenv("OPENAI_MODEL"); v != "" {
		return v
	}
	return "gpt-4o-mini"
}

func chat(messages []chatMessage) (string, error) {
	reqBody := chatRequest{
		Model:    getModel(),
		Messages: messages,
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return "", fmt.Errorf("marshal request: %w", err)
	}

	url := getBaseURL() + "/chat/completions"
	req, err := http.NewRequest("POST", url, bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	if key := getAPIKey(); key != "" {
		req.Header.Set("Authorization", "Bearer "+key)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("http request: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("read response: %w", err)
	}

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("API error (%d): %s", resp.StatusCode, string(respBody))
	}

	var cr chatResponse
	if err := json.Unmarshal(respBody, &cr); err != nil {
		return "", fmt.Errorf("parse response: %w", err)
	}

	if len(cr.Choices) == 0 {
		return "", fmt.Errorf("no choices in response")
	}

	return cr.Choices[0].Message.Content, nil
}

func runCmd(name string, arg ...string) (string, error) {
	var out, stderr bytes.Buffer
	cmd := exec.Command(name, arg...)
	cmd.Stdout = &out
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		return strings.TrimSpace(out.String()), fmt.Errorf("%s: %w\n%s", name, err, strings.TrimSpace(stderr.String()))
	}
	return strings.TrimSpace(out.String()), nil
}

func getTLDR(tool string) string {
	out, err := runCmd("tldr", tool)
	if err != nil {
		return ""
	}
	return out
}

func getManPage(tool string) string {
	// man pages can be long; limit via col -b for cleaner text
	cmd := exec.Command("sh", "-c", fmt.Sprintf("man %s 2>/dev/null | col -b", tool))
	var out bytes.Buffer
	cmd.Stdout = &out
	_ = cmd.Run()
	result := strings.TrimSpace(out.String())
	if result == "" {
		return ""
	}
	if len(result) > 6000 {
		result = result[:6000] + "\n... (truncated)"
	}
	return result
}

func getShellHistory() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return ""
	}

	candidates := []string{
		filepath.Join(home, ".zsh_history"),
		filepath.Join(home, ".bash_history"),
		filepath.Join(home, ".history"),
	}

	for _, p := range candidates {
		data, err := os.ReadFile(p)
		if err == nil && len(data) > 0 {
			lines := strings.Split(string(data), "\n")
			if len(lines) > 200 {
				lines = lines[len(lines)-200:]
			}
			return strings.Join(lines, "\n")
		}
	}
	return ""
}

func writeOutput(text string, paginate bool) {
	if paginate {
		cmd := exec.Command("less", "-R")
		cmd.Stdin = strings.NewReader(text)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		_ = cmd.Run()
	} else {
		fmt.Print(text)
	}
}

func handleCommand(query string, paginate bool) {
	messages := []chatMessage{
		{Role: "system", Content: "You identify Unix/Linux command-line tools from user questions. Respond with ONLY the tool name, nothing else. Example: 'how to find large files' -> 'find'. 'how to compress a directory' -> 'tar'. If no single tool fits, respond with 'general'."},
		{Role: "user", Content: query},
	}

	tool, err := chat(messages)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error identifying tool: %v\n", err)
		os.Exit(1)
	}
	tool = strings.TrimSpace(strings.ToLower(tool))

	var contextParts []string
	if tool != "general" && tool != "" {
		if t := getTLDR(tool); t != "" {
			contextParts = append(contextParts, "=== TLDR ===\n"+t)
		}
		if m := getManPage(tool); m != "" {
			contextParts = append(contextParts, "=== MAN PAGE ===\n"+m)
		}
	}
	context := strings.Join(contextParts, "\n\n")

	systemPrompt := "You are a command-line expert. Given a user's question, provide the exact command line they need.\n"
	if context != "" {
		systemPrompt += "Use the context below for reference.\n\n" + context
	}
	systemPrompt += "\n\nIMPORTANT: Provide ONLY the command line, nothing else. Use placeholders like /path/to/dir or filename for things the user needs to customize. The command should be directly runnable after replacing placeholders. Do NOT include any explanation, markdown formatting, or backticks."

	cmdMessages := []chatMessage{
		{Role: "system", Content: systemPrompt},
		{Role: "user", Content: query},
	}

	result, err := chat(cmdMessages)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	writeOutput(strings.TrimSpace(result)+"\n", paginate)
}

func handleHistory(query string, paginate bool) {
	history := getShellHistory()
	if history == "" {
		fmt.Fprintln(os.Stderr, "Error: no shell history found")
		os.Exit(1)
	}

	messages := []chatMessage{
		{Role: "system", Content: fmt.Sprintf("You are analyzing a user's shell history. Answer their question about their command history based on the data below. Be specific and reference actual commands from their history.\n\nShell history (last 500 lines):\n%s", history)},
		{Role: "user", Content: query},
	}

	result, err := chat(messages)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	writeOutput(strings.TrimSpace(result)+"\n", paginate)
}

func handleQuestion(query string, paginate bool) {
	messages := []chatMessage{
		{Role: "system", Content: "You are a command-line expert. Answer the user's question about commands and tools. Provide a terse, informative response in markdown format. Include specific command examples where relevant. Keep it concise but thorough."},
		{Role: "user", Content: query},
	}

	result, err := chat(messages)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	if paginate {
		cmd := exec.Command("less", "-R")
		cmd.Stdin = strings.NewReader(result)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		_ = cmd.Run()
	} else {
		cmd := exec.Command("glow")
		cmd.Stdin = strings.NewReader(result)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		_ = cmd.Run()
	}
}