discordo-plus/internal/ui/chat/attachment_handler.go

package chat

import (
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/ayn2op/discordo/internal/consts"
	"github.com/ayn2op/discordo/internal/ui"
	"github.com/ayn2op/tview/layers"
	"github.com/diamondburned/arikawa/v3/discord"
	"github.com/skratchdot/open-golang/open"
)

const maxAttachmentSize = 100 * 1024 * 1024 // 100 MB

var supportedImageTypes = map[string]bool{
	"image/jpeg": true,
	"image/png":  true,
	"image/webp": true,
	"image/gif":  true,
}

func (ml *messagesList) downloadToCache(attachment discord.Attachment) (string, error) {
	parsed, err := url.Parse(attachment.URL)
	if err != nil {
		return "", fmt.Errorf("invalid attachment URL: %w", err)
	}
	if parsed.Scheme != "https" {
		return "", fmt.Errorf("refusing non-HTTPS attachment URL: %s", parsed.Scheme)
	}

	resp, err := http.Get(attachment.URL)
	if err != nil {
		return "", fmt.Errorf("failed to fetch attachment: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("unexpected status %d fetching attachment", resp.StatusCode)
	}

	dir := filepath.Join(consts.CacheDir(), "attachments")
	if err := os.MkdirAll(dir, 0700); err != nil {
		return "", fmt.Errorf("failed to create attachments dir: %w", err)
	}

	safeName := filepath.Base(attachment.Filename)
	if safeName == "." || safeName == ".." {
		safeName = "attachment"
	}
	path := filepath.Join(dir, safeName)
	file, err := os.Create(path)
	if err != nil {
		return "", fmt.Errorf("failed to create attachment file: %w", err)
	}
	defer file.Close()

	if _, err := io.Copy(file, io.LimitReader(resp.Body, maxAttachmentSize)); err != nil {
		return "", fmt.Errorf("failed to write attachment file: %w", err)
	}

	return path, nil
}

func terminalGeometry() string {
	out, err := exec.Command("xdotool", "getactivewindow", "getwindowgeometry", "--shell").Output()
	if err != nil {
		return ""
	}
	var x, y, w, h string
	for _, line := range strings.Split(string(out), "\n") {
		switch {
		case strings.HasPrefix(line, "X="):
			x = strings.TrimPrefix(line, "X=")
		case strings.HasPrefix(line, "Y="):
			y = strings.TrimPrefix(line, "Y=")
		case strings.HasPrefix(line, "WIDTH="):
			w = strings.TrimPrefix(line, "WIDTH=")
		case strings.HasPrefix(line, "HEIGHT="):
			h = strings.TrimPrefix(line, "HEIGHT=")
		}
	}
	if w != "" && h != "" && x != "" && y != "" {
		return w + "x" + h + "+" + x + "+" + y
	}
	return ""
}

func viewerArgs(viewer, path string, customArgs []string) []string {
	if len(customArgs) > 0 {
		args := make([]string, 0, len(customArgs)+2)
		args = append(args, viewer)
		args = append(args, customArgs...)
		args = append(args, path)
		return args
	}
	if strings.HasSuffix(viewer, "mpv") || strings.Contains(viewer, "mpv ") {
		args := []string{viewer, "--force-window", "--loop-file=inf"}
		if geom := terminalGeometry(); geom != "" {
			args = append(args, "--geometry="+geom)
		}
		args = append(args, path)
		return args
	}
	return []string{viewer, path}
}

func (ml *messagesList) openAttachment(attachment discord.Attachment) {
	path, err := ml.downloadToCache(attachment)
	if err != nil {
		slog.Error("failed to download attachment", "err", err, "url", attachment.URL)
		return
	}

	viewer := ml.cfg.ImageViewer
	if viewer != "" && supportedImageTypes[attachment.ContentType] {
		if _, err := exec.LookPath(viewer); err != nil {
			slog.Error("image viewer not found in PATH", "viewer", viewer, "err", err)
			return
		}
		args := viewerArgs(viewer, path, ml.cfg.ImageViewerArgs)
		cmd := exec.Command(args[0], args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		ml.chatView.app.Suspend(func() {
			if err := cmd.Run(); err != nil {
				slog.Error("failed to run image viewer", "viewer", viewer, "err", err)
			}
		})
		return
	}

	if err := open.Start(path); err != nil {
		slog.Error("failed to open attachment file", "err", err, "path", path)
	}
}

func (ml *messagesList) saveAttachmentImage(attachment discord.Attachment) {
	path, err := ml.downloadToCache(attachment)
	if err != nil {
		slog.Error("failed to download attachment", "err", err, "url", attachment.URL)
		return
	}

	saveDir := ml.cfg.ImageSaveDir
	if saveDir == "" {
		saveDir = "."
	}
	if err := os.MkdirAll(saveDir, 0700); err != nil {
		slog.Error("failed to create save directory", "err", err, "path", saveDir)
		return
	}

	safeName := filepath.Base(attachment.Filename)
	if safeName == "." || safeName == ".." {
		safeName = "attachment"
	}
	destPath := filepath.Join(saveDir, safeName)
	src, err := os.Open(path)
	if err != nil {
		slog.Error("failed to open cached file", "err", err, "path", path)
		return
	}
	defer src.Close()

	dst, err := os.Create(destPath)
	if err != nil {
		slog.Error("failed to create save file", "err", err, "path", destPath)
		return
	}
	defer dst.Close()

	if _, err := io.Copy(dst, src); err != nil {
		slog.Error("failed to save image", "err", err)
		return
	}

	slog.Info("image saved", "path", destPath)
}

func (ml *messagesList) saveImage() {
	msg, err := ml.selectedMessage()
	if err != nil {
		slog.Error("failed to get selected message", "err", err)
		return
	}

	var images []discord.Attachment
	for _, a := range msg.Attachments {
		if supportedImageTypes[a.ContentType] {
			images = append(images, a)
		}
	}

	if len(images) == 0 {
		return
	}

	if len(images) == 1 {
		go ml.saveAttachmentImage(images[0])
		return
	}

	var items []attachmentItem
	for _, a := range images {
		attachment := a
		items = append(items, attachmentItem{
			label: attachment.Filename,
			open:  func() { go ml.saveAttachmentImage(attachment) },
		})
	}
	ml.attachmentsPicker.SetItems(items)
	ml.showAttachmentsOverlay()
}

func (ml *messagesList) openURL(url string) {
	if err := open.Start(url); err != nil {
		slog.Error("failed to open URL", "err", err, "url", url)
	}
}

func (ml *messagesList) showAttachmentsList(urls []string, attachments []discord.Attachment) {
	var items []attachmentItem
	for _, a := range attachments {
		attachment := a
		action := func() {
			if strings.HasPrefix(attachment.ContentType, "image/") {
				go ml.openAttachment(attachment)
			} else {
				go ml.openURL(attachment.URL)
			}
		}
		items = append(items, attachmentItem{
			label: attachment.Filename,
			open:  action,
		})
	}
	for _, u := range urls {
		url := u
		items = append(items, attachmentItem{
			label: url,
			open:  func() { go ml.openURL(url) },
		})
	}
	ml.attachmentsPicker.SetItems(items)
	ml.showAttachmentsOverlay()
}

func (ml *messagesList) showAttachmentsOverlay() {
	ml.chatView.
		AddLayer(
			ui.Centered(ml.attachmentsPicker, ml.cfg.Picker.Width, ml.cfg.Picker.Height),
			layers.WithName(attachmentsListLayerName),
			layers.WithResize(true),
			layers.WithVisible(true),
			layers.WithOverlay(),
		).
		SendToFront(attachmentsListLayerName)
	ml.chatView.app.SetFocus(ml.attachmentsPicker)
}