Pete/internal/matrix/client.go

package matrix

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"log/slog"
	"net/http"
	"os"
	"strings"
	"path/filepath"
	"sync"
	"time"

	"pete/internal/config"

	"maunium.net/go/mautrix"
	"maunium.net/go/mautrix/crypto/cryptohelper"
	"maunium.net/go/mautrix/event"
	"maunium.net/go/mautrix/id"
)

// ReactionHandler is called when a reaction is received on a post.
type ReactionHandler func(roomID id.RoomID, eventID id.EventID, targetEventID id.EventID, emoji string, userID id.UserID)

// deviceInfo holds persisted device credentials.
type deviceInfo struct {
	AccessToken string `json:"access_token"`
	DeviceID    string `json:"device_id"`
	UserID      string `json:"user_id"`
}

// Client wraps a mautrix client for Pete's needs: posting stories and tracking reactions.
type Client struct {
	mx        *mautrix.Client
	channels  map[string]id.RoomID
	adminRoom *id.RoomID
	userID    id.UserID
	cfg       config.MatrixConfig
	crypto    *cryptohelper.CryptoHelper

	mu         sync.RWMutex
	onReact    ReactionHandler
	cancelSync context.CancelFunc
}

// New creates a Matrix client with password login, device persistence, and E2EE.
func New(cfg config.MatrixConfig) (*Client, error) {
	if err := os.MkdirAll(cfg.DataDir, 0o755); err != nil {
		return nil, fmt.Errorf("create data dir: %w", err)
	}

	devicePath := filepath.Join(cfg.DataDir, "device.json")

	// Try to load existing device credentials
	device, err := loadDevice(devicePath)
	if err != nil {
		slog.Info("no existing device found, will login fresh")
	}

	var mx *mautrix.Client

	if device != nil {
		valid := isTokenValid(cfg.Homeserver, device.AccessToken)
		if valid {
			slog.Info("existing device credentials valid", "device_id", device.DeviceID)
			userID := id.UserID(device.UserID)
			mx, err = mautrix.NewClient(cfg.Homeserver, userID, device.AccessToken)
			if err != nil {
				return nil, fmt.Errorf("create client with existing token: %w", err)
			}
			mx.DeviceID = id.DeviceID(device.DeviceID)
		} else {
			slog.Warn("existing device credentials invalid, logging in again")
			device = nil
		}
	}

	if device == nil {
		mx, err = mautrix.NewClient(cfg.Homeserver, "", "")
		if err != nil {
			return nil, fmt.Errorf("create client: %w", err)
		}

		resp, err := mx.Login(context.Background(), &mautrix.ReqLogin{
			Type: mautrix.AuthTypePassword,
			Identifier: mautrix.UserIdentifier{
				Type: mautrix.IdentifierTypeUser,
				User: cfg.UserID,
			},
			Password:                 cfg.Password,
			InitialDeviceDisplayName: cfg.DisplayName,
		})
		if err != nil {
			return nil, fmt.Errorf("login: %w", err)
		}

		mx.AccessToken = resp.AccessToken
		mx.UserID = resp.UserID
		mx.DeviceID = resp.DeviceID

		device = &deviceInfo{
			AccessToken: resp.AccessToken,
			DeviceID:    string(resp.DeviceID),
			UserID:      string(resp.UserID),
		}
		if err := saveDevice(devicePath, device); err != nil {
			slog.Warn("failed to save device info", "err", err)
		}

		slog.Info("logged in successfully",
			"user_id", resp.UserID,
			"device_id", resp.DeviceID,
		)
	}

	// Set up E2EE via cryptohelper — stores crypto state in its own SQLite DB.
	// Persists device keys, olm/megolm sessions, cross-signing keys across restarts.
	cryptoDBPath := filepath.Join(cfg.DataDir, "crypto.db")
	ch, err := cryptohelper.NewCryptoHelper(mx, []byte(cfg.PickleKey), cryptoDBPath)
	if err != nil {
		return nil, fmt.Errorf("init crypto helper: %w", err)
	}

	// LoginAs enables the cryptohelper to re-login if the token expires
	ch.LoginAs = &mautrix.ReqLogin{
		Type: mautrix.AuthTypePassword,
		Identifier: mautrix.UserIdentifier{
			Type: mautrix.IdentifierTypeUser,
			User: cfg.UserID,
		},
		Password:                 cfg.Password,
		InitialDeviceDisplayName: cfg.DisplayName,
	}

	if err := ch.Init(context.Background()); err != nil {
		return nil, fmt.Errorf("crypto helper init: %w", err)
	}

	mx.Crypto = ch

	// Bootstrap cross-signing only if not already set up.
	mach := ch.Machine()
	existingKeys, err := mach.GetOwnCrossSigningPublicKeys(context.Background())
	if err != nil {
		slog.Warn("cross-signing: failed to fetch existing keys", "err", err)
	}
	if existingKeys == nil || existingKeys.MasterKey == "" {
		_, _, err = mach.GenerateAndUploadCrossSigningKeys(context.Background(), func(ui *mautrix.RespUserInteractive) interface{} {
			return map[string]interface{}{
				"type": mautrix.AuthTypePassword,
				"identifier": map[string]interface{}{
					"type": mautrix.IdentifierTypeUser,
					"user": cfg.UserID,
				},
				"password": cfg.Password,
				"session":  ui.Session,
			}
		}, "")
		if err != nil {
			slog.Warn("cross-signing: key upload failed", "err", err)
		} else {
			slog.Info("cross-signing: keys uploaded")
		}

		if err := mach.SignOwnDevice(context.Background(), mach.OwnIdentity()); err != nil {
			slog.Warn("cross-signing: sign own device failed", "err", err)
		} else {
			slog.Info("cross-signing: own device signed")
		}

		if err := mach.SignOwnMasterKey(context.Background()); err != nil {
			slog.Warn("cross-signing: sign master key failed", "err", err)
		} else {
			slog.Info("cross-signing: master key signed")
		}
	} else {
		slog.Info("cross-signing: already configured, skipping bootstrap")
	}

	slog.Info("E2EE initialized",
		"device_id", mx.DeviceID,
		"crypto_store", "sqlite-persistent",
	)

	channels := make(map[string]id.RoomID, len(cfg.Channels))
	for name, roomID := range cfg.Channels {
		channels[name] = id.RoomID(roomID)
	}

	c := &Client{
		mx:       mx,
		channels: channels,
		userID:   mx.UserID,
		cfg:      cfg,
		crypto:   ch,
	}

	if cfg.AdminRoom != "" {
		adminRoom := id.RoomID(cfg.AdminRoom)
		c.adminRoom = &adminRoom
	}

	return c, nil
}

// SetReactionHandler registers a callback for reaction events.
func (c *Client) SetReactionHandler(fn ReactionHandler) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.onReact = fn
}

// Start begins the Matrix sync loop in the background.
func (c *Client) Start(ctx context.Context) {
	syncer := c.mx.Syncer.(*mautrix.DefaultSyncer)

	// Listen for reactions
	syncer.OnEventType(event.EventReaction, func(ctx context.Context, evt *event.Event) {
		if evt.Sender == c.userID {
			return
		}

		content := evt.Content.AsReaction()
		if content == nil || content.RelatesTo.EventID == "" {
			return
		}

		c.mu.RLock()
		handler := c.onReact
		c.mu.RUnlock()

		if handler != nil {
			handler(
				evt.RoomID,
				evt.ID,
				content.RelatesTo.EventID,
				content.RelatesTo.Key,
				evt.Sender,
			)
		}
	})

	// Auto-join on invite
	syncer.OnEventType(event.StateMember, func(ctx context.Context, evt *event.Event) {
		if evt.GetStateKey() == string(c.userID) && evt.Content.AsMember().Membership == event.MembershipInvite {
			_, err := c.mx.JoinRoomByID(ctx, evt.RoomID)
			if err != nil {
				slog.Error("failed to auto-join room", "room", evt.RoomID, "err", err)
			} else {
				slog.Info("auto-joined room", "room", evt.RoomID)
			}
		}
	})

	syncCtx, cancel := context.WithCancel(ctx)
	c.cancelSync = cancel

	go func() {
		for {
			err := c.mx.SyncWithContext(syncCtx)
			if syncCtx.Err() != nil {
				return
			}
			if err != nil {
				slog.Error("matrix sync stopped, restarting in 5s", "err", err)
			} else {
				slog.Warn("matrix sync returned without error, restarting in 5s")
			}
			select {
			case <-time.After(5 * time.Second):
			case <-syncCtx.Done():
				return
			}
		}
	}()
}

// Stop halts the sync loop and closes the crypto store.
func (c *Client) Stop() {
	if c.cancelSync != nil {
		c.cancelSync()
	}
	if c.crypto != nil {
		if err := c.crypto.Close(); err != nil {
			slog.Error("failed to close crypto helper", "err", err)
		}
	}
}

// PostThreadedReply sends a plain/HTML message into the thread rooted at
// rootEventID. Resolves channel name → room ID like PostStory. The is_falling_back
// flag plus m.in_reply_to gives non-thread-aware clients a sensible quoted reply.
func (c *Client) PostThreadedReply(channel string, rootEventID id.EventID, plain, htmlBody string) error {
	roomID, ok := c.channels[channel]
	if !ok {
		return fmt.Errorf("unknown channel: %s", channel)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          plain,
		Format:        event.FormatHTML,
		FormattedBody: htmlBody,
		RelatesTo: &event.RelatesTo{
			Type:    event.RelThread,
			EventID: rootEventID,
			InReplyTo: &event.InReplyTo{
				EventID: rootEventID,
			},
			IsFallingBack: true,
		},
	}
	_, err := c.mx.SendMessageEvent(ctx, roomID, event.EventMessage, content)
	return err
}

// PostStory sends a story to a channel. Returns the text event ID for
// reaction tracking and a flag indicating whether the m.image event was
// successfully sent. Callers retrying after a failed text send MUST clear
// the ImageURL on the next attempt if imageSent is true, otherwise the
// image will be posted twice.
func (c *Client) PostStory(channel string, story *PostableStory) (eventID id.EventID, imageSent bool, err error) {
	roomID, ok := c.channels[channel]
	if !ok {
		return "", false, fmt.Errorf("unknown channel: %s", channel)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	// Send image first if present
	if story.ImageURL != "" {
		if ierr := c.sendImage(ctx, roomID, story.ImageURL); ierr != nil {
			slog.Warn("failed to send image, posting text only", "url", story.ImageURL, "err", ierr)
		} else {
			imageSent = true
		}
	}

	// Send text message
	plain, html := formatPost(story)
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          plain,
		Format:        event.FormatHTML,
		FormattedBody: html,
	}

	resp, err := c.mx.SendMessageEvent(ctx, roomID, event.EventMessage, content)
	if err != nil {
		return "", imageSent, fmt.Errorf("send message: %w", err)
	}

	return resp.EventID, imageSent, nil
}

// sendImage downloads an image, uploads it to Matrix, and sends it as m.image.
func (c *Client) sendImage(ctx context.Context, roomID id.RoomID, imageURL string) error {
	req, err := http.NewRequestWithContext(ctx, "GET", imageURL, nil)
	if err != nil {
		return fmt.Errorf("create image request: %w", err)
	}
	httpClient := &http.Client{Timeout: 15 * time.Second}
	resp, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("download image: %w", err)
	}
	defer resp.Body.Close()

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

	data, err := io.ReadAll(io.LimitReader(resp.Body, 10*1024*1024))
	if err != nil {
		return fmt.Errorf("read image: %w", err)
	}

	contentType := resp.Header.Get("Content-Type")
	if contentType == "" {
		contentType = "image/jpeg"
	}

	// Reject non-image content types
	if !strings.HasPrefix(contentType, "image/") {
		return fmt.Errorf("not an image: %s", contentType)
	}

	// Decode image dimensions so clients render it inline instead of as a download
	var width, height int
	if img, _, decErr := image.DecodeConfig(bytes.NewReader(data)); decErr == nil {
		width = img.Width
		height = img.Height
	}

	uploadResp, err := c.mx.UploadBytes(ctx, data, contentType)
	if err != nil {
		return fmt.Errorf("upload image: %w", err)
	}

	// Derive a filename with proper extension — some clients use Body as filename
	// and won't render inline without a recognized image extension
	ext := ".jpg"
	switch contentType {
	case "image/png":
		ext = ".png"
	case "image/gif":
		ext = ".gif"
	case "image/webp":
		ext = ".webp"
	}

	imgContent := &event.MessageEventContent{
		MsgType:  event.MsgImage,
		Body:     "image" + ext,
		FileName: "image" + ext,
		URL:      uploadResp.ContentURI.CUString(),
		Info: &event.FileInfo{
			MimeType: contentType,
			Size:     len(data),
			Width:    width,
			Height:   height,
		},
	}

	_, err = c.mx.SendMessageEvent(ctx, roomID, event.EventMessage, imgContent)
	return err
}

// SendAdminWarning sends a warning message to the admin room, if configured.
func (c *Client) SendAdminWarning(msg string) {
	if c.adminRoom == nil {
		slog.Warn("admin warning (no admin room configured)", "msg", msg)
		return
	}

	content := &event.MessageEventContent{
		MsgType: event.MsgText,
		Body:    msg,
	}

	if _, err := c.mx.SendMessageEvent(context.Background(), *c.adminRoom, event.EventMessage, content); err != nil {
		slog.Error("failed to send admin warning", "err", err, "msg", msg)
	}
}

// ChannelRoomID returns the room ID for a named channel.
// Returns false if the channel is not configured or has an empty room ID.
func (c *Client) ChannelRoomID(channel string) (id.RoomID, bool) {
	roomID, ok := c.channels[channel]
	if !ok || roomID == "" {
		return "", false
	}
	return roomID, true
}

func loadDevice(path string) (*deviceInfo, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var info deviceInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return nil, err
	}
	return &info, nil
}

func saveDevice(path string, info *deviceInfo) error {
	data, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0o600)
}

func isTokenValid(homeserver, accessToken string) bool {
	url := homeserver + "/_matrix/client/v3/account/whoami"
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return false
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == http.StatusOK
}