ygggo/internal/remote/handlers.go

/*
 *  Copyright (c) 2026 Kaiy Ragur
 *
 *  This Source Code Form is subject to the terms of the Mozilla Public
 *  License, v. 2.0. If a copy of the MPL was not distributed with this
 *  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

package remote

import (
	"bytes"
	"context"
	"encoding/hex"
	"io"
	"time"

	"github.com/neilalexander/yggmail/internal/notify"
	pb "github.com/neilalexander/yggmail/internal/remote/types"
	"github.com/neilalexander/yggmail/internal/utils"
)

// --- MAILBOXES ---

func (srv *Server) handleMailboxList(ctx context.Context, pk string, call *pb.MailboxListRequest) (*pb.Response, error) {
	list, err := srv.storage.MailboxList(pk, call.OnlySubscribed)
	if err != nil {
		return nil, err
	}
	return &pb.Response{Success: true, Data: &pb.Response_MboxList{MboxList: &pb.MailboxList{Names: list}}}, nil
}

func (srv *Server) handleMailboxCreate(ctx context.Context, pk string, call *pb.MailboxCreateRequest) (*pb.Response, error) {
	err := srv.storage.MailboxCreate(pk, call.Name)
	srv.log.Printf("RPC Server:: MailboxCreate: Name: %s: Pk: %s", call.Name, pk)
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailboxDelete(ctx context.Context, pk string, call *pb.MailboxDeleteRequest) (*pb.Response, error) {
	err := srv.storage.MailboxDelete(pk, call.Name)
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailboxRename(ctx context.Context, pk string, call *pb.MailboxRenameRequest) (*pb.Response, error) {
	err := srv.storage.MailboxRename(pk, call.OldName, call.NewName)
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailboxSubscribe(ctx context.Context, pk string, call *pb.MailboxSubscribeRequest) (*pb.Response, error) {
	err := srv.storage.MailboxSubscribe(pk, call.Name, call.Subscribed)
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailboxStatus(ctx context.Context, pk string, call *pb.MailboxStatusRequest) (*pb.Response, error) {
	exists, err := srv.storage.MailboxSelect(pk, call.Name)
	if err != nil || !exists {
		return &pb.Response{Success: false, Error: "Mailbox not found"}, nil
	}

	count, _ := srv.storage.MailCount(pk, call.Name)
	unseen, _ := srv.storage.MailUnseen(pk, call.Name)
	nextID, _ := srv.storage.MailNextID(pk, call.Name)

	return &pb.Response{
		Success: true,
		Data: &pb.Response_MboxStatus{
			MboxStatus: &pb.MailboxStatus{
				Name: call.Name, MessagesCount: uint32(count), UnseenCount: uint32(unseen), UidNext: uint32(nextID), UidValidity: 1,
			},
		},
	}, nil
}

// --- MAILS ---

// handleMailFetch processes fetching metadata and optionally streams a single message body.
func (srv *Server) handleMailFetch(ctx context.Context, pk string, call *pb.MailFetchRequest) (*pb.Response, *Payload, error) {
	pbList := &pb.MailList{}
	var messagePayload *Payload

	for _, id := range call.Ids {
		seq, mail, err := srv.storage.MailSelect(pk, call.Mailbox, int(id))
		if err != nil {
			continue
		}

		pbList.Items = append(pbList.Items, &pb.MailMessage{
			Id:       uint32(mail.ID),
			Seq:      uint32(seq),
			Date:     mail.Date.Unix(),
			Seen:     mail.Seen,
			Answered: mail.Answered,
			Flagged:  mail.Flagged,
			Deleted:  mail.Deleted,
			Size:     uint32(mail.Size),
		})

		// If fetching a single message, prepare the stream for the body payload
		if len(call.Ids) == 1 && mail.Body != nil {
			messagePayload = &Payload{
				R:    mail.Body,
				Size: uint32(mail.Size),
			}
		} else if mail.Body != nil {
			mail.Body.Close()
		}
	}

	return &pb.Response{
		Success: true,
		Data:    &pb.Response_MailList{MailList: pbList},
	}, messagePayload, nil
}

func (srv *Server) handleMailFlags(ctx context.Context, pk string, call *pb.MailUpdateFlagsRequest) (*pb.Response, error) {
	err := srv.storage.MailUpdateFlags(pk, call.Mailbox, int(call.Id), call.Seen, call.Answered, call.Flagged, call.Deleted)
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailMove(ctx context.Context, pk string, call *pb.MailMoveRequest) (*pb.Response, error) {
	_, err := srv.storage.MailMove(pk, call.Mailbox, int(call.Id), call.DestinationMailbox)
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailDelete(ctx context.Context, pk string, call *pb.MailDeleteRequest) (*pb.Response, error) {
	err := srv.storage.MailDelete(pk, call.Mailbox, int(call.Id))
	return &pb.Response{Success: err == nil}, err
}

func (srv *Server) handleMailExpunge(ctx context.Context, pk string, call *pb.MailExpungeRequest) (*pb.Response, error) {
	err := srv.storage.MailExpunge(pk, call.Mailbox)
	return &pb.Response{Success: err == nil}, err
}

// --- QUEUE & EXCHANGE ---

func (srv *Server) handleQueueInsert(ctx context.Context, pk string, call *pb.QueueInsertRequest) (*pb.Response, error) {
	err := srv.storage.QueueInsertDestinationForID(pk, call.Destination, int(call.Id), call.From, call.Rcpt)
	return &pb.Response{Success: err == nil}, err
}

// handleMailPush handles incoming mail streams and routes them (Incoming, Outgoing, or Save).
func (srv *Server) handleMailPush(ctx context.Context, senderPkStr string, call *pb.MailPushRequest, payload *Payload) (*pb.Response, error) {
	if payload == nil || payload.R == nil {
		return &pb.Response{Success: false, Error: "Empty payload"}, nil
	}
	defer payload.R.Close()

	senderPkBytes, _ := utils.DecodeString(senderPkStr)
	fromDom, fromAcc, _ := utils.ParseAddress(call.From)
	toDom, toAcc, _ := utils.ParseAddress(call.To)

	isFromMe := bytes.Equal(senderPkBytes, fromAcc) && fromDom.Equal(srv.nodePublickey)
	isFromPeerServer := bytes.Equal(senderPkBytes, fromDom)
	isToMe := toDom.Equal(srv.nodePublickey)

	if isFromMe && call.Mailbox != "" {
		return srv.saveToFolder(fromAcc, toDom, call.Mailbox, payload.R, payload.Size)
	}

	if isFromMe {
		return srv.processOutgoing(fromAcc, toDom, payload.R, payload.Size)
	}

	if isFromPeerServer && isToMe {
		accTo, _ := srv.accMgr.GetAccount(utils.EncodeString(toAcc))
		if accTo != nil {
			return srv.processIncoming(accTo.PublicKey, "INBOX", payload.R, payload.Size)
		}
	}

	return &pb.Response{Success: false, Error: "Relay access denied"}, nil
}

func (srv *Server) processIncoming(accPK, mailbox string, content io.Reader, size uint32) (*pb.Response, error) {
	id, err := srv.storage.MailCreate(accPK, mailbox, content, int(size))
	if err != nil {
		return nil, err
	}
	if count, err := srv.storage.MailCount(accPK, mailbox); err == nil {
		srv.notify.NotifyNew(accPK, notify.NewEmailEvent(accPK, mailbox, int64(id), int64(count)))
	}
	srv.log.Printf("RPC Server: MailPush: Incoming: Mailbox: %s for: %s\n", mailbox, accPK)
	return &pb.Response{Success: true, Data: &pb.Response_IdResult{IdResult: uint32(id)}}, nil
}

func (srv *Server) processOutgoing(accPK []byte, targetDom []byte, content io.Reader, size uint32) (*pb.Response, error) {
	pkStr := utils.EncodeString(accPK)
	_ = srv.storage.MailboxCreate(pkStr, "Outbox")

	id, err := srv.storage.MailCreate(pkStr, "Outbox", content, int(size))
	if err != nil {
		return nil, err
	}

	srv.log.Printf("RPC Server: MailPush: Outgoing: From: %s\n", accPK)
	srv.queue.QueueForServer(hex.EncodeToString(targetDom))

	return &pb.Response{Success: true, Data: &pb.Response_IdResult{IdResult: uint32(id)}}, nil
}

func (srv *Server) saveToFolder(accPK []byte, targetDom []byte, mailbox string, content io.Reader, size uint32) (*pb.Response, error) {
	pkStr := utils.EncodeString(accPK)
	_ = srv.storage.MailboxCreate(pkStr, mailbox)

	id, err := srv.storage.MailCreate(pkStr, mailbox, content, int(size))
	if err != nil {
		return nil, err
	}

	// Trigger queue processing if saving directly to Outbox
	if mailbox == "Outbox" {
		srv.queue.QueueForServer(hex.EncodeToString(targetDom))
	}

	srv.log.Printf("RPC Server: MailPush: Save: Mailbox: %s for: %s\n", mailbox, pkStr)
	return &pb.Response{Success: true, Data: &pb.Response_IdResult{IdResult: uint32(id)}}, nil
}

// --- UTILS ---

func (srv *Server) handleMailIdForSeq(ctx context.Context, pk string, call *pb.MailIdForSeqRequest) (*pb.Response, error) {
	id, err := srv.storage.MailIDForSeq(pk, call.Mailbox, int(call.Seq))
	if err != nil {
		return nil, err
	}

	return &pb.Response{
		Success: true,
		Data:    &pb.Response_IdResult{IdResult: uint32(id)},
	}, nil
}

func (srv *Server) handleMailSearch(ctx context.Context, pk string, call *pb.MailSearchRequest) (*pb.Response, error) {
	ids, err := srv.storage.MailSearch(pk, call.Mailbox)
	if err != nil {
		return nil, err
	}

	pbList := &pb.MailList{}
	for _, id := range ids {
		pbList.Items = append(pbList.Items, &pb.MailMessage{Id: id})
	}

	return &pb.Response{
		Success: true,
		Data:    &pb.Response_MailList{MailList: pbList},
	}, nil
}

func (srv *Server) handleMailWatch(ctx context.Context, pk string, call *pb.WatchRequest) (*pb.Response, error) {
	notifyChan := srv.notify.Subscribe(pk)
	defer srv.notify.Unsubscribe(pk, notifyChan)

	select {
	case ev := <-notifyChan:
		return srv.handleNotify(pk, ev)

	case <-time.After(40 * time.Second):
		return &pb.Response{Success: true, Error: "HEARTBEAT"}, nil

	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (srv *Server) handleNotify(pk string, event notify.Event) (*pb.Response, error) {
	srv.log.Println("RPC: Notify Subscribe: ", pk)
	switch ev := event.(type) {
	case *notify.EmailEvent:
		srv.log.Println("DEBUG:: RPC:: handleNotify:: Send:: EmailEvent: ", pk)
		return &pb.Response{
			Success: true,
			Data: &pb.Response_MboxStatus{
				MboxStatus: &pb.MailboxStatus{
					MessagesCount: uint32(ev.Count),
					UidNext:       uint32(ev.MessageId),
				},
			},
		}, nil
	case *notify.FlagsEvent:
		return &pb.Response{
			Success: true,
			Data: &pb.Response_MboxStatus{
				MboxStatus: &pb.MailboxStatus{},
			},
		}, nil

	default:
		return &pb.Response{Success: true, Error: "HEARTBEAT"}, nil
	}
}