workspace/apps/whatsme/main.go

package main

import (
	"context"
	"database/sql"
	"embed"
	"encoding/json"
	"fmt"
	"html/template"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/go-chi/chi/v5"
	_ "github.com/glebarez/sqlite"
	"go.mau.fi/whatsmeow"
	"go.mau.fi/whatsmeow/store/sqlstore"
	"go.mau.fi/whatsmeow/types"
	"go.mau.fi/whatsmeow/types/events"
	waLog "go.mau.fi/whatsmeow/util/log"
)

//go:embed templates/*
var templateFS embed.FS

type ConnectionState string

const (
	StateDisconnected ConnectionState = "disconnected"
	StateConnecting   ConnectionState = "connecting"
	StateWaitingQR    ConnectionState = "waiting_qr"
	StateConnected    ConnectionState = "connected"
	StateFailed       ConnectionState = "failed"
)

type Message struct {
	ID        int64
	Phone     string
	Name      string
	Message   string
	MessageID string
	Timestamp int64
}

type Event struct {
	Type    string `json:"type"`
	State   string `json:"state,omitempty"`
	QR      string `json:"qr,omitempty"`
	Error   string `json:"error,omitempty"`
	Message *Message `json:"message,omitempty"`
	Log     string `json:"log,omitempty"`
}

type App struct {
	db          *sql.DB
	client      *whatsmeow.Client
	container   *sqlstore.Container
	mu          sync.RWMutex
	state       ConnectionState
	currentQR   string
	eventChan   chan Event
	logger      *slog.Logger
	ctx         context.Context
	cancel      context.CancelFunc
	dataDir     string
}

func main() {
	port := os.Getenv("WHATSME_PORT")
	if port == "" {
		port = "8080"
	}
	dataDir := os.Getenv("WHATSME_DATA_DIR")
	if dataDir == "" {
		dataDir = "./whatsme"
	}

	if err := os.MkdirAll(dataDir, 0755); err != nil {
		slog.Error("failed to create data dir", "error", err)
		os.Exit(1)
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))

	ctx, cancel := context.WithCancel(context.Background())

	app := &App{
		eventChan: make(chan Event, 100),
		logger:    logger,
		ctx:       ctx,
		cancel:    cancel,
		dataDir:   dataDir,
		state:     StateDisconnected,
	}

	if err := app.initDB(); err != nil {
		logger.Error("failed to init DB", "error", err)
		os.Exit(1)
	}
	defer app.db.Close()

	if err := app.initWhatsApp(); err != nil {
		logger.Error("failed to init WhatsApp", "error", err)
		os.Exit(1)
	}

	go app.autoResume()

	r := chi.NewRouter()
	r.Get("/", app.handleIndex)
	r.Get("/events", app.handleEvents)
	r.Get("/health", app.handleHealth)

	srv := &http.Server{
		Addr:    ":" + port,
		Handler: r,
	}

	go func() {
		sig := make(chan os.Signal, 1)
		signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
		<-sig
		logger.Info("shutting down...")
		cancel()
		if app.client != nil {
			app.client.Disconnect()
		}
		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer shutdownCancel()
		srv.Shutdown(shutdownCtx)
	}()

	logger.Info("starting", "port", port, "data_dir", dataDir)
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logger.Error("server error", "error", err)
		os.Exit(1)
	}
}

func (a *App) initDB() error {
	dbPath := filepath.Join(a.dataDir, "messages.db")
	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return fmt.Errorf("open db: %w", err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
		return fmt.Errorf("enable foreign keys: %w", err)
	}
	schema := `
	CREATE TABLE IF NOT EXISTS messages (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		phone TEXT NOT NULL,
		name TEXT NOT NULL DEFAULT '',
		message TEXT NOT NULL DEFAULT '',
		message_id TEXT NOT NULL,
		timestamp INTEGER NOT NULL
	);
	CREATE UNIQUE INDEX IF NOT EXISTS idx_message_id ON messages(message_id);
	`
	if _, err := db.Exec(schema); err != nil {
		return fmt.Errorf("exec schema: %w", err)
	}
	a.db = db
	return nil
}

func (a *App) initWhatsApp() error {
	sessionPath := filepath.Join(a.dataDir, "session.db")
	db, err := sql.Open("sqlite", sessionPath)
	if err != nil {
		return fmt.Errorf("open session db: %w", err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
		return fmt.Errorf("enable foreign keys: %w", err)
	}
	a.container = sqlstore.NewWithDB(db, "sqlite", waLog.Stdout("SQL", "DEBUG", true))
	if err := a.container.Upgrade(a.ctx); err != nil {
		return fmt.Errorf("upgrade session db: %w", err)
	}
	return nil
}

func (a *App) autoResume() {
	a.mu.Lock()
	stores, err := a.container.GetAllDevices(a.ctx)
	a.mu.Unlock()
	if err != nil {
		a.logger.Error("failed to get stores", "error", err)
		return
	}
	if len(stores) == 0 {
		a.logger.Info("no stored session, starting QR flow")
		a.startQRFlow()
		return
	}
	a.logger.Info("found stored session, resuming")
	a.startQRFlow()
}

func (a *App) startQRFlow() {
	device := a.container.NewDevice()
	client := whatsmeow.NewClient(device, waLog.Stdout("Client", "DEBUG", true))
	a.mu.Lock()
	a.client = client
	a.state = StateWaitingQR
	a.mu.Unlock()

	qrChan, err := client.GetQRChannel(a.ctx)
	if err != nil {
		a.logger.Error("failed to get QR channel", "error", err)
		a.setState(StateFailed, "failed to get QR channel")
		return
	}

	if err := client.Connect(); err != nil {
		a.logger.Error("failed to connect", "error", err)
		a.setState(StateFailed, "connection failed")
		return
	}

	a.addMessageHandler(client)

	go func() {
		for evt := range qrChan {
			switch evt.Event {
			case whatsmeow.QRChannelSuccess.Event:
				a.logger.Info("WhatsApp connected successfully")
				a.setState(StateConnected, "")
				return
			case whatsmeow.QRChannelEventError:
				a.logger.Error("QR error", "error", evt.Error)
				a.setState(StateFailed, evt.Error.Error())
				return
			case whatsmeow.QRChannelTimeout.Event:
				a.logger.Info("QR timeout, restarting")
				a.setState(StateDisconnected, "QR timed out")
				go a.startQRFlow()
				return
			default:
				if evt.Code != "" {
					a.mu.Lock()
					a.currentQR = evt.Code
					a.state = StateWaitingQR
					a.mu.Unlock()
					a.broadcastQR(evt.Code)
				}
			}
		}
	}()
}

func (a *App) setState(state ConnectionState, errMsg string) {
	a.mu.Lock()
	a.state = state
	if errMsg != "" {
		a.logger.Info("state changed", "state", state, "error", errMsg)
	}
	a.mu.Unlock()

	a.eventChan <- Event{
		Type:  "state",
		State: string(state),
		Error: errMsg,
	}
	a.eventChan <- Event{
		Type: "log",
		Log:  fmt.Sprintf("[%s] %s", time.Now().Format("15:04:05"), state),
	}

	if state == StateDisconnected {
		go a.startQRFlow()
	}
}

func (a *App) broadcastQR(code string) {
	a.eventChan <- Event{
		Type: "state",
		State: string(StateWaitingQR),
		QR: code,
	}
	a.eventChan <- Event{
		Type: "log",
		Log: fmt.Sprintf("[%s] QR code generated", time.Now().Format("15:04:05")),
	}
}

func (a *App) addMessageHandler(client *whatsmeow.Client) {
	client.AddEventHandler(func(raw interface{}) {
		evt, ok := raw.(*events.Message)
		if !ok {
			return
		}

		a.logger.Debug("raw event",
			"type", fmt.Sprintf("%T", raw),
			"is_from_me", evt.Info.IsFromMe,
			"sender", evt.Info.Sender.String(),
			"chat", evt.Info.Chat.String(),
			"message_id", evt.Info.ID,
			"timestamp", evt.Info.Timestamp,
			"push_name", evt.Info.PushName,
			"sender_server", evt.Info.Sender.Server,
			"chat_server", evt.Info.Chat.Server,
			"message_type", fmt.Sprintf("%T", evt.Message),
		)

		if evt.Info.IsFromMe {
			a.logger.Info("skipping own message", "message_id", evt.Info.ID)
			return
		}

		if evt.Info.Chat.Server == types.GroupServer {
			a.logger.Info("skipping group message", "chat", evt.Info.Chat.String())
			return
		}

		if evt.Info.ID == "" {
			a.logger.Info("skipping message with empty ID")
			return
		}

		sender := evt.Info.Sender
		phone := sender.User
		if phone == "" {
			phone = sender.String()
		}
		if !strings.HasPrefix(phone, "+") {
			phone = "+" + phone
		}
		if strings.HasSuffix(phone, "@s.whatsapp.net") {
			phone = strings.TrimSuffix(phone, "@s.whatsapp.net")
			phone = "+" + strings.TrimPrefix(phone, "+")
		}

		pushName := evt.Info.PushName

		var text string
		if evt.Message != nil {
			if c := evt.Message.GetConversation(); c != "" {
				text = c
			} else if ext := evt.Message.GetExtendedTextMessage(); ext != nil {
				text = ext.GetText()
			}
		}

		if text == "" {
			a.logger.Info("skipping non-text message", "message_id", evt.Info.ID)
			return
		}

		ts := evt.Info.Timestamp.Unix()
		if ts == 0 {
			ts = time.Now().Unix()
		}

		msg := &Message{
			Phone:     phone,
			Name:      pushName,
			Message:   text,
			MessageID: evt.Info.ID,
			Timestamp: ts,
		}

		if err := a.saveMessage(msg); err != nil {
			a.logger.Error("failed to save message", "error", err)
			return
		}

		a.logger.Info("message saved",
			"phone", phone,
			"name", pushName,
			"message", text,
			"message_id", evt.Info.ID,
		)

		output := map[string]interface{}{
			"phone":      phone,
			"name":       pushName,
			"message":    text,
			"message_id": evt.Info.ID,
			"timestamp":  ts,
		}
		if b, err := json.Marshal(output); err == nil {
			fmt.Println(string(b))
		}

		a.eventChan <- Event{
			Type:    "message",
			Message: msg,
		}
		a.eventChan <- Event{
			Type: "log",
			Log:  fmt.Sprintf("[%s] MSG from %s (%s): %s", time.Now().Format("15:04:05"), phone, pushName, text),
		}
	})
}

func (a *App) saveMessage(msg *Message) error {
	_, err := a.db.Exec(
		`INSERT OR IGNORE INTO messages (phone, name, message, message_id, timestamp) VALUES (?, ?, ?, ?, ?)`,
		msg.Phone, msg.Name, msg.Message, msg.MessageID, msg.Timestamp,
	)
	return err
}

func (a *App) getMessages() ([]Message, error) {
	rows, err := a.db.Query(
		`SELECT id, phone, name, message, message_id, timestamp FROM messages ORDER BY timestamp DESC LIMIT 100`,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var msgs []Message
	for rows.Next() {
		var m Message
		if err := rows.Scan(&m.ID, &m.Phone, &m.Name, &m.Message, &m.MessageID, &m.Timestamp); err != nil {
			return nil, err
		}
		msgs = append(msgs, m)
	}
	return msgs, nil
}

func (a *App) handleIndex(w http.ResponseWriter, r *http.Request) {
	a.mu.RLock()
	state := a.state
	currentQR := a.currentQR
	a.mu.RUnlock()

	msgs, err := a.getMessages()
	if err != nil {
		a.logger.Error("failed to get messages", "error", err)
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}

	tmpl, err := template.ParseFS(templateFS, "templates/index.gohtml")
	if err != nil {
		a.logger.Error("failed to parse template", "error", err)
		http.Error(w, "template error", http.StatusInternalServerError)
		return
	}

	data := struct {
		State   string
		QR      string
		Messages []Message
	}{
		State:    string(state),
		QR:       currentQR,
		Messages: msgs,
	}

	if err := tmpl.Execute(w, data); err != nil {
		a.logger.Error("failed to render template", "error", err)
	}
}

func (a *App) handleEvents(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming not supported", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	a.mu.RLock()
	state := a.state
	currentQR := a.currentQR
	a.mu.RUnlock()

	if state == StateWaitingQR && currentQR != "" {
		fmt.Fprintf(w, "event: state\ndata: {\"type\":\"state\",\"state\":\"%s\",\"qr\":\"%s\"}\n\n", state, escapeJSON(currentQR))
		flusher.Flush()
	} else {
		fmt.Fprintf(w, "event: state\ndata: {\"type\":\"state\",\"state\":\"%s\"}\n\n", state)
		flusher.Flush()
	}

	for evt := range a.eventChan {
		switch evt.Type {
		case "state":
			if evt.QR != "" {
				fmt.Fprintf(w, "event: state\ndata: {\"type\":\"state\",\"state\":\"%s\",\"qr\":\"%s\",\"error\":\"%s\"}\n\n", evt.State, escapeJSON(evt.QR), evt.Error)
			} else {
				fmt.Fprintf(w, "event: state\ndata: {\"type\":\"state\",\"state\":\"%s\",\"error\":\"%s\"}\n\n", evt.State, evt.Error)
			}
		case "message":
			if evt.Message != nil {
				msgJSON, _ := json.Marshal(evt.Message)
				fmt.Fprintf(w, "event: message\ndata: {\"type\":\"message\",\"message\":%s}\n\n", msgJSON)
			}
		case "log":
			fmt.Fprintf(w, "event: log\ndata: {\"type\":\"log\",\"log\":\"%s\"}\n\n", escapeJSON(evt.Log))
		}
		flusher.Flush()
	}
}

func (a *App) handleHealth(w http.ResponseWriter, r *http.Request) {
	a.mu.RLock()
	state := a.state
	a.mu.RUnlock()
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"status":      "ok",
		"connected":   state == StateConnected,
		"state":       state,
	})
}

func escapeJSON(s string) string {
	s = strings.ReplaceAll(s, "\n", "\\n")
	s = strings.ReplaceAll(s, "\r", "\\r")
	s = strings.ReplaceAll(s, "\t", "\\t")
	s = strings.ReplaceAll(s, `"`, `\"`)
	return s
}