// Package attach holds the pictures a prompt carries.
//
// A terminal cannot paste an image. Ctrl-V over a running full-screen program
// delivers whatever text form the clipboard has — for a screenshot, usually
// nothing at all, sometimes a file path the window manager put there. So the
// image is not taken from the keystroke: the keystroke is the signal to go and
// ask the system clipboard directly, with the same helper a person would run
// by hand (wl-paste, xclip, pngpaste, PowerShell).
//
// What lands in the prompt box is a marker, [image 1]. It is text like any
// other text: it can be moved, retyped, or deleted, and deleting it is how a
// paste is taken back — only markers still standing when enter is pressed turn
// into image blocks, in the place where they stand.
package attach

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/johannSo/codey/internal/llm"
)

// MaxBytes caps one image. The API refuses anything over 5MB, and it counts
// the encoded payload, so the raw ceiling sits below that by the 4/3 base64
// grows by.
const MaxBytes = 3_750_000

// ErrNoImage means the clipboard held no picture. It is the ordinary answer to
// ctrl-V with text on the clipboard, not a fault, and the caller should fall
// through to pasting text rather than complain.
var ErrNoImage = errors.New("no image on the clipboard")

// ErrTooBig is an image found and refused. Unlike a format that is not an
// image at all, this one is worth interrupting someone over: they pasted a
// picture and it is not going to arrive.
var ErrTooBig = errors.New("image too large")

// Image is one picture waiting for the next prompt. Data is the raw file, not
// base64: encoding happens once, on the way out.
type Image struct {
	Media  string
	Data   []byte
	Origin string // "clipboard", or the path it was read from — for the status line
}

// Block renders the image as an API content block.
func (im Image) Block() llm.ContentBlock {
	return llm.ImageBlock(im.Media, base64.StdEncoding.EncodeToString(im.Data))
}

// Sniff identifies an image by its leading bytes. Extensions lie — a .png that
// is really a JPEG is common enough, and the API checks the bytes — so the
// magic number is what decides, and a format the API does not take is not an
// image as far as this package is concerned.
func Sniff(b []byte) (string, bool) {
	switch {
	case bytes.HasPrefix(b, []byte("\x89PNG\r\n\x1a\n")):
		return "image/png", true
	case bytes.HasPrefix(b, []byte{0xff, 0xd8, 0xff}):
		return "image/jpeg", true
	case bytes.HasPrefix(b, []byte("GIF87a")), bytes.HasPrefix(b, []byte("GIF89a")):
		return "image/gif", true
	case len(b) >= 12 && bytes.HasPrefix(b, []byte("RIFF")) && bytes.Equal(b[8:12], []byte("WEBP")):
		return "image/webp", true
	}
	return "", false
}

// New wraps bytes that are already in hand, checking them the way anything
// else that becomes an Image is checked.
func New(data []byte, origin string) (Image, error) {
	if len(data) == 0 {
		return Image{}, ErrNoImage
	}
	if len(data) > MaxBytes {
		return Image{}, fmt.Errorf("%w: it is %s, over the %s the API takes", ErrTooBig, size(len(data)), size(MaxBytes))
	}
	media, ok := Sniff(data)
	if !ok {
		return Image{}, errors.New("not a PNG, JPEG, GIF or WebP")
	}
	return Image{Media: media, Data: data, Origin: origin}, nil
}

// FromFile reads an image off disk. This is the path a drag-and-drop takes:
// most terminals answer a dropped file by pasting its name.
func FromFile(path string) (Image, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return Image{}, err
	}
	im, err := New(data, filepath.Base(path))
	if err != nil {
		return Image{}, fmt.Errorf("%s: %w", filepath.Base(path), err)
	}
	return im, nil
}

func size(n int) string {
	if n < 1<<20 {
		return fmt.Sprintf("%dKB", n>>10)
	}
	return fmt.Sprintf("%.1fMB", float64(n)/(1<<20))
}

// imageExts are the endings worth trying to open. A pasted path is only read
// from disk when it looks like an image; every other paste is text, and text
// that happens to name a file should stay text.
var imageExts = map[string]bool{
	".png": true, ".jpg": true, ".jpeg": true, ".gif": true, ".webp": true,
}

// Path pulls a local image path out of pasted text, and reports whether the
// paste was one. Terminals hand over a dropped file in several dresses: bare,
// quoted, percent-escaped behind file://, or with the spaces backslashed the
// way a shell would want them.
func Path(s string) (string, bool) {
	s = strings.TrimSpace(s)
	if s == "" || strings.ContainsAny(s, "\n\r") {
		return "", false // more than one line is prose, not a file name
	}
	s = strings.Trim(s, "'\"")
	if rest, ok := strings.CutPrefix(s, "file://"); ok {
		if i := strings.Index(rest, "/"); i > 0 {
			rest = rest[i:] // drop a host, including the empty one in file:///
		}
		s = unescapeURL(rest)
	} else if !strings.Contains(s, `\\`) {
		// A dropped name arrives shell-quoted: "my holiday.png" as
		// my\ holiday.png. Undoing that is only safe where the backslash is
		// not itself a path separator, so Windows paths are left alone.
		s = strings.ReplaceAll(s, `\ `, " ")
	}
	if !imageExts[strings.ToLower(filepath.Ext(s))] {
		return "", false
	}
	if strings.HasPrefix(s, "~/") {
		if home, err := os.UserHomeDir(); err == nil {
			s = filepath.Join(home, s[2:])
		}
	}
	if fi, err := os.Stat(s); err != nil || fi.IsDir() {
		return "", false
	}
	return s, true
}

func unescapeURL(s string) string {
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		if s[i] == '%' && i+2 < len(s) {
			if n, err := strconv.ParseUint(s[i+1:i+3], 16, 8); err == nil {
				b.WriteByte(byte(n))
				i += 2
				continue
			}
		}
		b.WriteByte(s[i])
	}
	return b.String()
}

// ---- markers ------------------------------------------------------------

var markerRE = regexp.MustCompile(`\[image (\d+)\]`)

// Marker is what stands in the prompt box for image n, counting from one.
func Marker(n int) string { return "[image " + strconv.Itoa(n) + "]" }

// Set is the images pasted into the prompt being typed. It is emptied when
// that prompt is sent.
type Set struct {
	imgs []Image
}

// ExpandPaths turns image paths left in a prompt into attached images and
// returns the text with each one replaced by its marker.
//
// This is the backstop under the paste key, and it is what actually carries
// most terminals. Catching a paste as it happens needs the terminal to bracket
// it — to wrap it in an escape sequence meaning "this was pasted" — and plenty
// of them do not. Text that arrives as ordinary keystrokes cannot be told from
// typing, so it is not caught on the way in; it is caught here, on the way
// out, where how it arrived has stopped mattering.
//
// Only a path that really is a readable image on disk is taken. A filename in
// a sentence, a path to source, a file that is not there: all left exactly as
// written, because a prompt that talks about a file is the ordinary case and
// has to survive untouched.
func (s *Set) ExpandPaths(text string) string {
	// A line that is nothing but one path is tried whole first, so a name with
	// spaces in it survives the splitting below.
	if p, ok := Path(text); ok {
		if im, err := FromFile(p); err == nil {
			return Marker(s.Add(im))
		}
		return text
	}

	// Replacements are gathered before any is made: rewriting the text while
	// walking it would shift everything after the first hit.
	type swap struct{ from, to string }
	var swaps []swap
	seen := map[string]bool{}
	for _, f := range strings.Fields(text) {
		if seen[f] {
			continue // the same path twice is one picture, mentioned twice
		}
		p, ok := Path(f)
		if !ok {
			continue
		}
		im, err := FromFile(p)
		if err != nil {
			continue
		}
		seen[f] = true
		swaps = append(swaps, swap{from: f, to: Marker(s.Add(im))})
	}
	for _, sw := range swaps {
		text = strings.ReplaceAll(text, sw.from, sw.to)
	}
	return text
}

// Add takes an image and returns its number, which is what the marker says.
func (s *Set) Add(im Image) int {
	s.imgs = append(s.imgs, im)
	return len(s.imgs)
}

func (s *Set) Len() int { return len(s.imgs) }

func (s *Set) Reset() { s.imgs = nil }

// Blocks turns a typed prompt into the content of one user message, with each
// surviving marker replaced in place by the picture it names.
//
// A marker whose number names no image is left alone: it is then literally
// what someone typed, and typing "[image 7]" at the model is allowed to mean
// what it says. An image whose marker was deleted is dropped, which is what
// makes deleting the marker an undo.
func (s *Set) Blocks(text string) []llm.ContentBlock {
	var out []llm.ContentBlock
	words := false
	// The marker stood inside a sentence with spaces around it. Those spaces
	// held it apart from the words and have nothing to do once it is gone, so
	// each run of text is trimmed at the edge where a picture was lifted out.
	addText := func(t string) {
		if t = strings.TrimSpace(t); t != "" {
			out = append(out, llm.TextBlock(t))
			words = true
		}
	}

	last := 0
	for _, loc := range markerRE.FindAllStringSubmatchIndex(text, -1) {
		n, err := strconv.Atoi(text[loc[2]:loc[3]])
		if err != nil || n < 1 || n > len(s.imgs) {
			continue // not one of ours; it stays in the text around it
		}
		addText(text[last:loc[0]])
		out = append(out, s.imgs[n-1].Block())
		last = loc[1]
	}
	addText(text[last:])

	if !words {
		// Nothing but markers. The message still needs words, because a user
		// turn of images alone reads as a question with no question in it —
		// and an empty content list is an API error besides.
		out = append(out, llm.TextBlock("(see the image)"))
	}
	return out
}