package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"os"
	"path/filepath"
	"strings"

	_ "golang.org/x/image/webp"
)

const (
	imageDirName  = "images"
	imageMaxBytes = 20 << 20
)

// ImageStore holds immutable images shared by the root, children, and
// subagents. Names are the SHA-256 of the bytes with the decoded format as
// extension. Publication is atomic and does not replace an existing image.
type ImageStore struct {
	dir string
}

// Save validates and stores data, and returns its filename. Only JPEG, PNG, and
// WebP images up to 20 MB are accepted.
func (s *ImageStore) Save(data []byte) (name string, err error) {
	format, err := imageFormat(data)
	if err != nil {
		return "", fmt.Errorf("check image: %w", err)
	}

	sum := sha256.Sum256(data)
	name = hex.EncodeToString(sum[:]) + "." + format
	err = os.MkdirAll(s.dir, 0o700)
	if err != nil {
		return "", fmt.Errorf("create image directory: %w", err)
	}

	file, err := os.CreateTemp(s.dir, ".image-*")
	if err != nil {
		return "", fmt.Errorf("create image file: %w", err)
	}

	defer func() {
		err = errors.Join(err, os.Remove(file.Name()))
	}()

	_, err = file.Write(data)
	if err == nil {
		err = file.Sync()
	}

	err = errors.Join(err, file.Close())
	if err != nil {
		return "", fmt.Errorf("write image: %w", err)
	}

	err = os.Link(file.Name(), filepath.Join(s.dir, name))
	if err != nil && !errors.Is(err, os.ErrExist) {
		return "", fmt.Errorf("publish image: %w", err)
	}

	return name, nil
}

// Read returns the bytes of a stored image after checking its name and content.
func (s *ImageStore) Read(name string) ([]byte, error) {
	if !imageName(name) {
		return nil, errors.New("invalid image name")
	}

	data, err := imageReadFile(filepath.Join(s.dir, name))
	if err != nil {
		return nil, fmt.Errorf("read stored image: %w", err)
	}

	format, err := imageFormat(data)
	if err != nil {
		return nil, fmt.Errorf("check stored image: %w", err)
	}

	sum := sha256.Sum256(data)
	if name != hex.EncodeToString(sum[:])+"."+format {
		return nil, errors.New("image content does not match its name")
	}

	return data, nil
}

func imageFormat(data []byte) (string, error) {
	if len(data) > imageMaxBytes {
		return "", errors.New("image exceeds 20 MB")
	}

	config, format, err := image.DecodeConfig(bytes.NewReader(data))
	if err != nil || config.Width <= 0 || config.Height <= 0 {
		return "", errors.New("use a JPEG, PNG, or WebP image")
	}

	switch format {
	case "jpeg":
		return "jpg", nil
	case "png", "webp":
		return format, nil
	default:
		return "", errors.New("use a JPEG, PNG, or WebP image")
	}
}

func imageName(name string) bool {
	hash, ext, found := strings.Cut(name, ".")
	return found && len(hash) == 64 && strings.Trim(hash, "0123456789abcdef") == "" && (ext == "jpg" || ext == "png" || ext == "webp")
}

func imageReadFile(path string) (data []byte, err error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open image: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	data, err = io.ReadAll(io.LimitReader(file, imageMaxBytes+1))
	if err != nil {
		return nil, fmt.Errorf("read image: %w", err)
	}

	if len(data) > imageMaxBytes {
		return nil, errors.New("image exceeds 20 MB")
	}

	return data, nil
}

// imageFiles reads only complete marker lines with valid store filenames.
func imageFiles(text string) []string {
	var names []string

	for line := range strings.SplitSeq(text, "\n") {
		name, found := strings.CutPrefix(line, "[image ")
		if !found || !strings.HasSuffix(name, "]") {
			continue
		}

		name = strings.TrimSuffix(name, "]")
		if imageName(name) {
			names = append(names, name)
		}
	}

	return names
}

func imageText(names []string, caption string) string {
	var text strings.Builder

	for _, name := range names {
		text.WriteString("[image " + name + "]\n")
	}

	text.WriteString(caption)

	return strings.TrimSuffix(text.String(), "\n")
}