// pico — a minimal single-binary CLI coding agent.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"
)

// ------------------------------------------------------------------- palette

const (
	reset = "\x1b[0m"
	bold  = "\x1b[1m"
	dim   = "\x1b[2m"
	blue  = "\x1b[34m"
	cyan  = "\x1b[36m"
	green = "\x1b[32m"
	red   = "\x1b[31m"
)

// -------------------------------------------------------------------- config

var (
	apiURL = firstEnv("API_URL", "ANTHROPIC_BASE_URL", "http://localhost:8000")
	apiKey = firstEnv("API_KEY", "ANTHROPIC_API_KEY", "")
	model  = firstEnv("MODEL", "ANTHROPIC_MODEL", "gpt-oss-120b")
)

func firstEnv(keys ...string) string {
	def := keys[len(keys)-1]
	for _, k := range keys[:len(keys)-1] {
		if k == "" {
			continue
		}
		if v := os.Getenv(k); v != "" {
			return strings.TrimRight(v, "/")
		}
	}
	return def
}

// --------------------------------------------------------------------- tools

type param struct {
	name string
	typ  string // trailing "?" marks optional
}

type tool struct {
	name   string
	desc   string
	params []param
	run    func(in map[string]any) (string, error)
}

var tools = []tool{
	{"read", "Read file with line numbers (file path, not directory)",
		[]param{{"path", "string"}, {"offset", "integer?"}, {"limit", "integer?"}}, toolRead},
	{"write", "Write content to file",
		[]param{{"path", "string"}, {"content", "string"}}, toolWrite},
	{"edit", "Replace old with new in file (old must be unique unless all=true)",
		[]param{{"path", "string"}, {"old", "string"}, {"new", "string"}, {"all", "boolean?"}}, toolEdit},
	{"glob", "Find files by pattern, sorted by mtime",
		[]param{{"pat", "string"}, {"path", "string?"}}, toolGlob},
	{"grep", "Search files for regex pattern",
		[]param{{"pat", "string"}, {"path", "string?"}}, toolGrep},
	{"bash", "Run shell command",
		[]param{{"cmd", "string"}}, toolBash},
}

var byName = func() map[string]*tool {
	m := map[string]*tool{}
	for i := range tools {
		m[tools[i].name] = &tools[i]
	}
	return m
}()

// schemas is built once at startup and reused verbatim on every request, so
// the serialized prefix stays byte-identical and KV cache holds.
// Properties are emitted in declaration order, which map marshaling would not
// preserve. Each property carries "type" and nothing else.
var schemas = func() []json.RawMessage {
	out := make([]json.RawMessage, len(tools))
	for i, t := range tools {
		var props, required []string
		for _, p := range t.params {
			props = append(props, fmt.Sprintf(`%s:{"type":%s}`, jstr(p.name), jstr(strings.TrimSuffix(p.typ, "?"))))
			if !strings.HasSuffix(p.typ, "?") {
				required = append(required, jstr(p.name))
			}
		}
		out[i] = json.RawMessage(fmt.Sprintf(
			`{"name":%s,"description":%s,"input_schema":{"type":"object","properties":{%s},"required":[%s]}}`,
			jstr(t.name), jstr(t.desc), strings.Join(props, ","), strings.Join(required, ",")))
	}
	return out
}()

func jstr(s string) string {
	enc, _ := json.Marshal(s)
	return string(enc)
}

// ------------------------------------------------------------- arg coercion

func argStr(in map[string]any, k, def string) string {
	if s, ok := in[k].(string); ok {
		return s
	}
	return def
}

func argInt(in map[string]any, k string, def int) int {
	switch v := in[k].(type) {
	case float64:
		return int(v)
	case string:
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
	}
	return def
}

func argBool(in map[string]any, k string, def bool) bool {
	if b, ok := in[k].(bool); ok {
		return b
	}
	return def
}

// ---------------------------------------------------------------------- read

func toolRead(in map[string]any) (string, error) {
	data, err := os.ReadFile(argStr(in, "path", ""))
	if err != nil {
		return "", err
	}
	// SplitAfter keeps the newline on each line; drop the empty tail element.
	lines := strings.SplitAfter(string(data), "\n")
	if lines[len(lines)-1] == "" {
		lines = lines[:len(lines)-1]
	}
	offset := argInt(in, "offset", 0)
	limit := argInt(in, "limit", len(lines))
	start := min(max(offset, 0), len(lines))
	end := min(max(offset+limit, start), len(lines))

	var b strings.Builder
	for i, ln := range lines[start:end] {
		fmt.Fprintf(&b, "%4d| %s", start+i+1, ln)
	}
	return b.String(), nil
}

// ---------------------------------------------------------------- write, edit

func toolWrite(in map[string]any) (string, error) {
	if err := os.WriteFile(argStr(in, "path", ""), []byte(argStr(in, "content", "")), 0644); err != nil {
		return "", err
	}
	return "ok", nil
}

func toolEdit(in map[string]any) (string, error) {
	path := argStr(in, "path", "")
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	s, old, new := string(data), argStr(in, "old", ""), argStr(in, "new", "")

	// Both of these are ordinary results, not errors: the model reads them
	// and retries with a longer `old`.
	switch n := strings.Count(s, old); {
	case n == 0:
		return "error: old_string not found", nil
	case n > 1 && !argBool(in, "all", false):
		return fmt.Sprintf("error: old_string appears %d times, must be unique (use all=true)", n), nil
	}

	if argBool(in, "all", false) {
		s = strings.ReplaceAll(s, old, new)
	} else {
		s = strings.Replace(s, old, new, 1)
	}
	if err := os.WriteFile(path, []byte(s), 0644); err != nil {
		return "", err
	}
	return "ok", nil
}

// ---------------------------------------------------------------- glob, grep

func argPath(in map[string]any) string {
	if p := argStr(in, "path", "."); p != "" {
		return p
	}
	return "."
}

func toolGlob(in map[string]any) (string, error) {
	matches := globExpand(strings.ReplaceAll(argPath(in)+"/"+argStr(in, "pat", ""), "//", "/"))
	if len(matches) == 0 {
		return "none", nil
	}
	mtime := make(map[string]int64, len(matches))
	for _, m := range matches {
		if fi, err := os.Stat(m); err == nil && !fi.IsDir() {
			mtime[m] = fi.ModTime().UnixNano()
		}
	}
	sort.SliceStable(matches, func(i, j int) bool { return mtime[matches[i]] > mtime[matches[j]] })
	return strings.Join(matches, "\n"), nil
}

func toolGrep(in map[string]any) (string, error) {
	re, err := regexp.Compile(argStr(in, "pat", ""))
	if err != nil {
		return "", err
	}
	var hits []string
	for _, p := range globExpand(argPath(in) + "/**") {
		f, err := os.Open(p)
		if err != nil {
			continue
		}
		// Directories appear in the ** expansion; scanning one yields nothing.
		sc := bufio.NewScanner(f)
		sc.Buffer(make([]byte, 0, 64*1024), 8<<20)
		for n := 1; sc.Scan(); n++ {
			if line := sc.Text(); re.MatchString(line) {
				hits = append(hits, fmt.Sprintf("%s:%d:%s", p, n, strings.TrimRight(line, " \t\r")))
			}
		}
		f.Close()
		if len(hits) >= 50 {
			break
		}
	}
	if len(hits) == 0 {
		return "none", nil
	}
	if len(hits) > 50 {
		hits = hits[:50]
	}
	return strings.Join(hits, "\n"), nil
}

// ---------------------------------------------------------------------- bash

func toolBash(in map[string]any) (string, error) {
	pr, pw, err := os.Pipe()
	if err != nil {
		return "", err
	}
	cmd := exec.Command("sh", "-c", argStr(in, "cmd", ""))
	cmd.Stdout, cmd.Stderr = pw, pw // one pipe: interleaved in real order
	// Own process group, so the timeout below can kill the whole tree.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if err := cmd.Start(); err != nil {
		pw.Close()
		pr.Close()
		return "", err
	}
	// The parent MUST drop its own copy of the write end, or the pipe never
	// reaches EOF and the reader below blocks forever -- even for `true`.
	pw.Close()

	done := make(chan string, 1)
	go func() {
		defer pr.Close()
		var acc strings.Builder
		sc := bufio.NewScanner(pr)
		sc.Buffer(make([]byte, 0, 64*1024), 8<<20)
		for sc.Scan() {
			line := sc.Text()
			fmt.Printf("  %s│ %s%s\n", dim, strings.TrimRight(line, " \t\r"), reset)
			acc.WriteString(line + "\n")
		}
		done <- acc.String()
	}()

	var out string
	select {
	case out = <-done:
	case <-time.After(30 * time.Second):
		// Kill the GROUP, not just sh: `sh -c "sleep 60"` forks, and the
		// grandchild inherits the pipe's write end.
		if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
			cmd.Process.Kill()
		}
		out = <-done // let the reader drain before reporting
		out += "\n(timed out after 30s)"
	}
	cmd.Wait() // always reap

	if out = strings.TrimSpace(out); out == "" {
		return "(empty)", nil
	}
	return out, nil
}

// ------------------------------------------------------------------ dispatch

func runTool(name string, input json.RawMessage) string {
	t, ok := byName[name]
	if !ok {
		return fmt.Sprintf("error: unknown tool %q", name)
	}
	var in map[string]any
	json.Unmarshal(input, &in)
	out, err := t.run(in)
	if err != nil {
		return "error: " + err.Error()
	}
	return out
}

// --------------------------------------------------------------- glob engine

func globExpand(pattern string) []string {
	root, rest := ".", pattern
	if strings.HasPrefix(pattern, "/") {
		root, rest = "/", pattern[1:]
	}
	var out []string
	globWalk(root, strings.Split(rest, "/"), &out)
	if root == "." {
		for i := range out {
			out[i] = strings.TrimPrefix(out[i], "./")
		}
	}
	return out
}

func globWalk(dir string, segs []string, out *[]string) {
	if len(segs) == 0 {
		*out = append(*out, dir)
		return
	}
	seg, rest := segs[0], segs[1:]

	if seg == "" || seg == "." {
		globWalk(dir, rest, out)
		return
	}

	if seg == "**" {
		globWalk(dir, rest, out) // ** matches zero path segments
		for _, e := range readDir(dir) {
			if strings.HasPrefix(e.Name(), ".") {
				continue
			}
			switch p := join(dir, e.Name()); {
			case e.IsDir():
				globWalk(p, segs, out) // keep ** at the head
			case len(rest) == 0:
				*out = append(*out, p)
			}
		}
		return
	}

	for _, e := range readDir(dir) {
		name := e.Name()
		if strings.HasPrefix(name, ".") && !strings.HasPrefix(seg, ".") {
			continue
		}
		if ok, _ := filepath.Match(seg, name); !ok {
			continue
		}
		p := join(dir, name)
		if len(rest) == 0 {
			*out = append(*out, p) // file or directory
		} else if e.IsDir() {
			globWalk(p, rest, out)
		}
	}
}

// readDir returns entries sorted by name; unreadable directories yield none.
func readDir(dir string) []os.DirEntry {
	es, err := os.ReadDir(dir)
	if err != nil {
		return nil
	}
	return es
}

func join(dir, name string) string {
	if strings.HasSuffix(dir, "/") {
		return dir + name
	}
	return dir + "/" + name
}

// ------------------------------------------------------------------ API call

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

type toolResult struct {
	Type      string `json:"type"`
	ToolUseID string `json:"tool_use_id"`
	Content   string `json:"content"`
}

type block struct {
	Type  string          `json:"type"`
	Text  string          `json:"text"`
	ID    string          `json:"id"`
	Name  string          `json:"name"`
	Input json.RawMessage `json:"input"`
}

var systemPrompt = func() string {
	cwd, _ := os.Getwd()
	return "Concise coding assistant. cwd: " + cwd
}()

func callAPI(messages []message) ([]json.RawMessage, error) {
	body, err := json.Marshal(map[string]any{
		"model":      model,
		"max_tokens": 8192,
		"system":     systemPrompt,
		"messages":   messages,
		"tools":      schemas,
	})
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", apiURL+"/v1/messages", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	// Every header is byte-identical across turns: a varying one would
	// invalidate llama-swap's KV cache on every request.
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("anthropic-version", "2023-06-01")
	req.Header.Set("x-api-key", apiKey)

	// No timeout: a cold llama-swap model load costs ~10-25s before prefill
	// even begins, and llama-swap's own responseHeader timeout is 1440s.
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(raw), 500))
	}
	// Content blocks stay raw: thinking/text/tool_use and anything unknown
	// all round-trip verbatim into the next request.
	var parsed struct {
		Content []json.RawMessage `json:"content"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, err
	}
	return parsed.Content, nil
}

// ---------------------------------------------------------------- agent loop

func agentLoop(messages *[]message) error {
	for {
		blocks, err := callAPI(*messages)
		if err != nil {
			return err
		}
		var results []toolResult
		for _, raw := range blocks {
			var b block
			if err := json.Unmarshal(raw, &b); err != nil {
				return err
			}
			switch b.Type {
			case "text":
				if bareMode {
					// In bare mode, output raw text without leading newline or markers.
					fmt.Print(b.Text)
					if !strings.HasSuffix(b.Text, "\n") {
						fmt.Print("\n")
					}
				} else {
					fmt.Printf("\n%s⏺%s %s\n", cyan, reset, renderMarkdown(b.Text))
				}
			case "tool_use":
				if !bareMode {
					fmt.Printf("\n%s⏺ %s%s(%s%s%s)\n",
						green, capitalize(b.Name), reset, dim, argPreview(b.Name, b.Input), reset)
				}
				out := runTool(b.Name, b.Input)
				printResult(out)
				results = append(results, toolResult{"tool_result", b.ID, out})
			}
			// Other block types (thinking) are not displayed but still echoed.
		}
		*messages = append(*messages, message{"assistant", blocks})
		if len(results) == 0 {
			return nil
		}
		*messages = append(*messages, message{"user", results})
	}
}

// ------------------------------------------------------------------- display

var boldRe = regexp.MustCompile(`(?s)\*\*(.*?)\*\*`)

func renderMarkdown(s string) string { return boldRe.ReplaceAllString(s, bold+"$1"+reset) }

func capitalize(s string) string {
	if s == "" {
		return s
	}
	return strings.ToUpper(s[:1]) + s[1:]
}

// truncate cuts to n code points, never bytes, and never adds an ellipsis.
func truncate(s string, n int) string {
	if utf8.RuneCountInString(s) <= n {
		return s
	}
	return string([]rune(s)[:n])
}

func argPreview(name string, input json.RawMessage) string {
	t, ok := byName[name]
	if !ok {
		return ""
	}
	var in map[string]any
	json.Unmarshal(input, &in)
	for _, p := range t.params {
		v, present := in[p.name]
		if !present {
			continue
		}
		if s, isStr := v.(string); isStr {
			return truncate(s, 50)
		}
		enc, _ := json.Marshal(v)
		return truncate(string(enc), 50)
	}
	return ""
}

// Global flag indicating bare (non‑interactive) mode.
var bareMode bool

func printResult(out string) {
	if bareMode {
		// In bare mode output the raw result.
		fmt.Print(out)
		if !strings.HasSuffix(out, "\n") {
			fmt.Print("\n")
		}
		return
	}
	lines := strings.Split(out, "\n")
	preview := truncate(lines[0], 60)
	if len(lines) > 1 {
		preview += fmt.Sprintf(" ... +%d lines", len(lines)-1)
	} else if utf8.RuneCountInString(lines[0]) > 60 {
		preview += "..."
	}
	fmt.Printf("  %s⎿  %s%s\n", dim, preview, reset)
}

// ---------------------------------------------------------------------- REPL

func main() {
	fi, _ := os.Stdin.Stat()
	bareMode = (fi.Mode() & os.ModeCharDevice) == 0

	if bareMode {
		data, _ := io.ReadAll(os.Stdin)
		line := strings.TrimSpace(string(data))
		if line == "" {
			return
		}
		messages := []message{{"user", line}}
		_ = agentLoop(&messages) // errors are printed inside agentLoop via fmt.Printf in interactive mode; here we ignore.
		return
	}

	cwd, _ := os.Getwd()
	fmt.Printf("%spico%s | %s%s | %s%s\n\n", bold, reset, dim, model, cwd, reset)

	sep := dim + strings.Repeat("─", 80) + reset
	stdin := bufio.NewScanner(os.Stdin)
	stdin.Buffer(make([]byte, 0, 64*1024), 8<<20)

	var messages []message
	for {
		fmt.Println(sep)
		fmt.Printf("%s%s❯%s ", bold, blue, reset)
		if !stdin.Scan() {
			fmt.Println()
			return
		}
		fmt.Println(sep)

		switch line := strings.TrimSpace(stdin.Text()); line {
		case "":
			continue
		case "/q", "exit":
			return
		case "/c":
			messages = nil
			fmt.Printf("%s⏺ Cleared conversation%s\n", green, reset)
		default:
			messages = append(messages, message{"user", line})
			if err := agentLoop(&messages); err != nil {
				fmt.Printf("%s⏺ Error: %v%s\n", red, err, reset)
				continue
			}
			fmt.Println()
		}
	}
}