source.go

package visionapi

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/textproto"
	"os"
	"path/filepath"
	"strings"
)

// Source is one of the four ways to name a file. Exactly one is required per request, and
// the constructors below are the only implementations:
//
//	visionapi.FilePath("invoice.pdf")            // read from disk
//	visionapi.FileBytes("scan.png", raw)         // bytes you already hold
//	visionapi.FileReader("scan.png", r)          // anything that reads
//	visionapi.FileURL("https://…/invoice.pdf")   // the API fetches it
//	visionapi.FileBase64(encoded)                // base64, "data:" prefix optional
//
// Local bytes go out as multipart/form-data; a URL or a base64 payload goes out as JSON,
// which is smaller and avoids a pointless encode step. Both forms use identical field
// names, so the endpoint cannot tell them apart.
type Source interface {
	build(fields map[string]any) (body []byte, contentType string, err error)
}

type fileSource struct {
	name string
	open func() (io.ReadCloser, error)
}

// FilePath uploads a file from disk. The filename is sent for readability only — the
// server detects the type from magic bytes and ignores what we declare.
func FilePath(path string) Source {
	return fileSource{
		name: filepath.Base(path),
		open: func() (io.ReadCloser, error) { return os.Open(path) },
	}
}

// FileBytes uploads bytes you already hold.
func FileBytes(filename string, data []byte) Source {
	if filename == "" {
		filename = "upload"
	}
	return fileSource{
		name: filename,
		open: func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(data)), nil },
	}
}

// FileReader uploads from any reader — an HTTP body, an S3 object, a decompressor.
//
// The reader is consumed once and buffered, because a retry has to send the same bytes
// again. For a file on disk prefer FilePath, which reopens instead of buffering.
func FileReader(filename string, r io.Reader) Source {
	if filename == "" {
		filename = "upload"
	}
	return fileSource{
		name: filename,
		open: func() (io.ReadCloser, error) { return io.NopCloser(r), nil },
	}
}

func (f fileSource) build(fields map[string]any) ([]byte, string, error) {
	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)

	for name, value := range fields {
		switch v := value.(type) {
		case string:
			if err := writer.WriteField(name, v); err != nil {
				return nil, "", err
			}
		case bool:
			if err := writer.WriteField(name, fmt.Sprintf("%t", v)); err != nil {
				return nil, "", err
			}
		case []string:
			// Repeated parts, which is how `questions` is meant to arrive over multipart.
			for _, item := range v {
				if err := writer.WriteField(name, item); err != nil {
					return nil, "", err
				}
			}
		default:
			encoded, err := json.Marshal(v)
			if err != nil {
				return nil, "", &UsageError{Message: fmt.Sprintf("could not encode %q: %v", name, err)}
			}
			if err := writer.WriteField(name, string(encoded)); err != nil {
				return nil, "", err
			}
		}
	}

	reader, err := f.open()
	if err != nil {
		return nil, "", &UsageError{Message: fmt.Sprintf("could not read %q: %v", f.name, err)}
	}
	defer reader.Close()

	header := make(textproto.MIMEHeader)
	header.Set("Content-Disposition",
		fmt.Sprintf(`form-data; name="file"; filename=%q`, f.name))
	header.Set("Content-Type", contentTypeFor(f.name))
	part, err := writer.CreatePart(header)
	if err != nil {
		return nil, "", err
	}
	if _, err := io.Copy(part, reader); err != nil {
		return nil, "", &UsageError{Message: fmt.Sprintf("could not read %q: %v", f.name, err)}
	}
	if err := writer.Close(); err != nil {
		return nil, "", err
	}
	return buf.Bytes(), writer.FormDataContentType(), nil
}

type jsonSource struct {
	key   string
	value string
}

// FileURL has the API fetch a publicly reachable file itself.
func FileURL(u string) Source { return jsonSource{key: "file_url", value: u} }

// FileBase64 sends base64-encoded bytes. A "data:…;base64," prefix is accepted.
func FileBase64(encoded string) Source { return jsonSource{key: "file_base64", value: encoded} }

func (j jsonSource) build(fields map[string]any) ([]byte, string, error) {
	if j.value == "" {
		return nil, "", &UsageError{Message: "empty " + j.key}
	}
	payload := make(map[string]any, len(fields)+1)
	for k, v := range fields {
		payload[k] = v
	}
	payload[j.key] = j.value
	body, err := json.Marshal(payload)
	if err != nil {
		return nil, "", &UsageError{Message: "could not encode request body: " + err.Error()}
	}
	return body, "application/json", nil
}

func contentTypeFor(name string) string {
	switch strings.ToLower(filepath.Ext(name)) {
	case ".pdf":
		return "application/pdf"
	case ".png":
		return "image/png"
	case ".jpg", ".jpeg":
		return "image/jpeg"
	case ".webp":
		return "image/webp"
	case ".tif", ".tiff":
		return "image/tiff"
	}
	return "application/octet-stream"
}

func randomKey() string {
	buf := make([]byte, 16)
	if _, err := rand.Read(buf); err != nil {
		// crypto/rand does not fail in practice; an empty key simply means the request
		// goes out without idempotency rather than not going out at all.
		return ""
	}
	return hex.EncodeToString(buf)
}