message-bridge/main.go

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gorilla/mux"
	_ "github.com/mattn/go-sqlite3"
	"github.com/skip2/go-qrcode"
	"go.mau.fi/whatsmeow"
	"go.mau.fi/whatsmeow/proto/waE2E"
	"go.mau.fi/whatsmeow/store/sqlstore"
	"go.mau.fi/whatsmeow/types"
	"go.mau.fi/whatsmeow/types/events"
	waLog "go.mau.fi/whatsmeow/util/log"
)

type Message struct {
	ID            string    `json:"id"`
	Platform      string    `json:"platform"`
	From          string    `json:"from"`
	FromName      string    `json:"from_name,omitempty"`
	To            string    `json:"to,omitempty"`
	Body          string    `json:"body"`
	Timestamp     time.Time `json:"timestamp"`
	IsGroup       bool      `json:"is_group"`
	GroupName     string    `json:"group_name,omitempty"`
	MediaType     string    `json:"media_type,omitempty"`
	MediaURL      string    `json:"media_url,omitempty"`
	MediaPath     string    `json:"media_path,omitempty"`
	HasMedia      bool      `json:"has_media"`
	Transcription string    `json:"transcription,omitempty"`
}

type MessageBridge struct {
	waClient      *whatsmeow.Client
	messages      []Message
	msgMutex      sync.RWMutex
	dataDir       string
	webhookURL    string
	webhookToken  string
	qrChan        chan string
	connected     bool
	webhooksSent  map[string]bool // Track sent webhook message IDs
	webhooksMutex sync.Mutex
}

var bridge *MessageBridge

func NewMessageBridge(dataDir string) *MessageBridge {
	return &MessageBridge{
		dataDir:      dataDir,
		messages:     make([]Message, 0),
		qrChan:       make(chan string, 1),
		webhookURL:   os.Getenv("WEBHOOK_URL"),
		webhookToken: os.Getenv("WEBHOOK_TOKEN"),
		webhooksSent: make(map[string]bool),
	}
}

func (b *MessageBridge) InitWhatsApp() error {
	ctx := context.Background()
	dbPath := filepath.Join(b.dataDir, "whatsapp.db")
	container, err := sqlstore.New(ctx, "sqlite3", "file:"+dbPath+"?_foreign_keys=on", waLog.Noop)
	if err != nil {
		return fmt.Errorf("failed to create store: %w", err)
	}

	deviceStore, err := container.GetFirstDevice(ctx)
	if err != nil {
		return fmt.Errorf("failed to get device: %w", err)
	}

	b.waClient = whatsmeow.NewClient(deviceStore, waLog.Noop)
	b.waClient.AddEventHandler(b.handleWhatsAppEvent)

	if b.waClient.Store.ID == nil {
		// Need to pair
		qrChan, _ := b.waClient.GetQRChannel(ctx)
		err = b.waClient.Connect()
		if err != nil {
			return fmt.Errorf("failed to connect: %w", err)
		}

		go func() {
			for evt := range qrChan {
				if evt.Event == "code" {
					select {
					case b.qrChan <- evt.Code:
					default:
					}
				}
			}
		}()
	} else {
		err = b.waClient.Connect()
		if err != nil {
			return fmt.Errorf("failed to connect: %w", err)
		}
		b.connected = true
	}

	return nil
}

func (b *MessageBridge) handleWhatsAppEvent(evt interface{}) {
	switch v := evt.(type) {
	case *events.Connected:
		b.connected = true
		log.Println("WhatsApp connected")

	case *events.Message:
		msg := b.parseWhatsAppMessage(v)
		if msg != nil {
			b.addMessage(*msg)
			b.sendWebhook(*msg)
		}
	}
}

func (b *MessageBridge) parseWhatsAppMessage(evt *events.Message) *Message {
	info := evt.Info

	msg := &Message{
		ID:        info.ID,
		Platform:  "whatsapp",
		From:      info.Sender.User,
		Timestamp: info.Timestamp,
		IsGroup:   info.IsGroup,
	}

	// Get sender name
	if info.PushName != "" {
		msg.FromName = info.PushName
	}

	// Get group info
	if info.IsGroup {
		msg.To = info.Chat.User
		// Try to get group name from cache
		ctx := context.Background()
		groupInfo, err := b.waClient.GetGroupInfo(ctx, info.Chat)
		if err == nil {
			msg.GroupName = groupInfo.Name
		}
	}

	// Parse message content
	if evt.Message != nil {
		if evt.Message.GetConversation() != "" {
			msg.Body = evt.Message.GetConversation()
		} else if evt.Message.GetExtendedTextMessage() != nil {
			msg.Body = evt.Message.GetExtendedTextMessage().GetText()
		} else if img := evt.Message.GetImageMessage(); img != nil {
			msg.HasMedia = true
			msg.MediaType = "image"
			msg.Body = img.GetCaption()
			b.downloadMedia(evt, msg)
		} else if vid := evt.Message.GetVideoMessage(); vid != nil {
			msg.HasMedia = true
			msg.MediaType = "video"
			msg.Body = vid.GetCaption()
			b.downloadMedia(evt, msg)
		} else if audio := evt.Message.GetAudioMessage(); audio != nil {
			msg.HasMedia = true
			if audio.GetPTT() {
				msg.MediaType = "voice"
			} else {
				msg.MediaType = "audio"
			}
			b.downloadMedia(evt, msg)
		} else if doc := evt.Message.GetDocumentMessage(); doc != nil {
			msg.HasMedia = true
			msg.MediaType = "document"
			msg.Body = doc.GetFileName()
			b.downloadMedia(evt, msg)
		}
	}

	return msg
}

func (b *MessageBridge) downloadMedia(evt *events.Message, msg *Message) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	// Get the downloadable message
	var downloadable whatsmeow.DownloadableMessage
	var mimeType string
	if img := evt.Message.GetImageMessage(); img != nil {
		downloadable = img
		mimeType = img.GetMimetype()
	} else if vid := evt.Message.GetVideoMessage(); vid != nil {
		downloadable = vid
		mimeType = vid.GetMimetype()
	} else if audio := evt.Message.GetAudioMessage(); audio != nil {
		downloadable = audio
		mimeType = audio.GetMimetype()
	} else if doc := evt.Message.GetDocumentMessage(); doc != nil {
		downloadable = doc
		mimeType = doc.GetMimetype()
	}

	if downloadable == nil {
		return
	}

	data, err := b.waClient.Download(ctx, downloadable)
	if err != nil {
		log.Printf("Failed to download media: %v", err)
		return
	}

	// Save to file
	mediaDir := filepath.Join(b.dataDir, "media")
	os.MkdirAll(mediaDir, 0755)

	ext := getExtFromMime(mimeType, msg.MediaType, msg.Body)
	filename := fmt.Sprintf("%s_%s%s", msg.Platform, msg.ID, ext)
	filePath := filepath.Join(mediaDir, filename)

	if err := os.WriteFile(filePath, data, 0644); err != nil {
		log.Printf("Failed to save media: %v", err)
		return
	}

	msg.MediaURL = "/media/" + filename
	msg.MediaPath = filePath
	log.Printf("Saved media: %s (%d bytes, %s)", filePath, len(data), mimeType)

	// Transcribe voice messages
	if msg.MediaType == "voice" {
		b.transcribeAudio(msg)
	}
}

// transcribeAudio calls the transcription script for voice messages
func (b *MessageBridge) transcribeAudio(msg *Message) {
	transcriptScript := os.Getenv("TRANSCRIBE_SCRIPT")
	if transcriptScript == "" {
		transcriptScript = "/home/johan/clawd/scripts/transcribe-audio.sh"
	}

	// Check if script exists and GROQ_API_KEY is set
	if _, err := os.Stat(transcriptScript); os.IsNotExist(err) {
		log.Printf("Transcription script not found: %s", transcriptScript)
		return
	}
	if os.Getenv("GROQ_API_KEY") == "" {
		log.Printf("GROQ_API_KEY not set, skipping transcription")
		return
	}

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

	cmd := exec.CommandContext(ctx, transcriptScript, msg.MediaPath)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		log.Printf("Transcription failed: %v (stderr: %s)", err, stderr.String())
		return
	}

	transcription := strings.TrimSpace(stdout.String())
	if transcription != "" {
		msg.Transcription = transcription
		log.Printf("Transcribed voice message: %s", truncate(transcription, 100))
	}
}

func getExtFromMime(mimeType, mediaType, body string) string {
	// Try mime type first
	switch {
	case strings.HasPrefix(mimeType, "image/jpeg"):
		return ".jpg"
	case strings.HasPrefix(mimeType, "image/png"):
		return ".png"
	case strings.HasPrefix(mimeType, "image/webp"):
		return ".webp"
	case strings.HasPrefix(mimeType, "image/gif"):
		return ".gif"
	case strings.HasPrefix(mimeType, "video/mp4"):
		return ".mp4"
	case strings.HasPrefix(mimeType, "audio/ogg"):
		return ".ogg"
	case strings.HasPrefix(mimeType, "audio/mpeg"):
		return ".mp3"
	case strings.HasPrefix(mimeType, "audio/mp4"):
		return ".m4a"
	case strings.HasPrefix(mimeType, "application/pdf"):
		return ".pdf"
	}

	// Fallback to media type
	switch mediaType {
	case "image":
		return ".jpg"
	case "video":
		return ".mp4"
	case "audio", "voice":
		return ".ogg"
	case "document":
		if body != "" {
			if ext := filepath.Ext(body); ext != "" {
				return ext
			}
		}
		return ".bin"
	}
	return ".bin"
}

func (b *MessageBridge) addMessage(msg Message) {
	b.msgMutex.Lock()
	defer b.msgMutex.Unlock()

	b.messages = append(b.messages, msg)

	// Keep last 1000 messages in memory
	if len(b.messages) > 1000 {
		b.messages = b.messages[len(b.messages)-1000:]
	}

	log.Printf("[%s] %s (%s): %s", msg.Platform, msg.FromName, msg.From, truncate(msg.Body, 50))
}

func (b *MessageBridge) sendWebhook(msg Message) {
	if b.webhookURL == "" {
		return
	}

	// Check if we've already sent a webhook for this message
	b.webhooksMutex.Lock()
	if b.webhooksSent[msg.ID] {
		b.webhooksMutex.Unlock()
		return
	}
	b.webhooksSent[msg.ID] = true
	// Prune old entries to prevent memory leak (keep last 1000)
	if len(b.webhooksSent) > 1000 {
		b.webhooksSent = make(map[string]bool)
		b.webhooksSent[msg.ID] = true
	}
	b.webhooksMutex.Unlock()

	payload, _ := json.Marshal(map[string]interface{}{
		"event":   "new_message",
		"message": msg,
	})

	go func() {
		req, err := http.NewRequest("POST", b.webhookURL, strings.NewReader(string(payload)))
		if err != nil {
			log.Printf("Webhook request error: %v", err)
			return
		}
		req.Header.Set("Content-Type", "application/json")
		if b.webhookToken != "" {
			req.Header.Set("Authorization", "Bearer "+b.webhookToken)
		}

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Printf("Webhook send error: %v", err)
			return
		}
		defer resp.Body.Close()

		if resp.StatusCode >= 400 {
			log.Printf("Webhook error response: %d", resp.StatusCode)
		} else {
			log.Printf("Webhook sent: %s from %s", msg.Platform, msg.From)
		}
	}()
}

func truncate(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max] + "..."
}

// HTTP Handlers
func handleStatus(w http.ResponseWriter, r *http.Request) {
	status := map[string]interface{}{
		"whatsapp": map[string]interface{}{
			"connected": bridge.connected,
			"needs_qr":  bridge.waClient != nil && bridge.waClient.Store.ID == nil,
		},
		"message_count": len(bridge.messages),
	}
	json.NewEncoder(w).Encode(status)
}

func handleQR(w http.ResponseWriter, r *http.Request) {
	format := r.URL.Query().Get("format")
	
	select {
	case code := <-bridge.qrChan:
		// Put it back for other requests
		select {
		case bridge.qrChan <- code:
		default:
		}

		if format == "png" {
			png, err := qrcode.Encode(code, qrcode.Medium, 256)
			if err != nil {
				http.Error(w, err.Error(), 500)
				return
			}
			w.Header().Set("Content-Type", "image/png")
			w.Write(png)
		} else if format == "base64" {
			png, err := qrcode.Encode(code, qrcode.Medium, 256)
			if err != nil {
				http.Error(w, err.Error(), 500)
				return
			}
			json.NewEncoder(w).Encode(map[string]string{
				"qr_base64": base64.StdEncoding.EncodeToString(png),
				"qr_raw":    code,
			})
		} else {
			// Terminal QR
			qr, _ := qrcode.New(code, qrcode.Medium)
			w.Header().Set("Content-Type", "text/plain")
			w.Write([]byte(qr.ToSmallString(false)))
		}
	default:
		if bridge.connected {
			json.NewEncoder(w).Encode(map[string]string{"status": "already_connected"})
		} else {
			json.NewEncoder(w).Encode(map[string]string{"status": "no_qr_available"})
		}
	}
}

func handleMessages(w http.ResponseWriter, r *http.Request) {
	bridge.msgMutex.RLock()
	defer bridge.msgMutex.RUnlock()

	platform := r.URL.Query().Get("platform")
	limit := 50

	var filtered []Message
	for i := len(bridge.messages) - 1; i >= 0 && len(filtered) < limit; i-- {
		msg := bridge.messages[i]
		if platform == "" || msg.Platform == platform {
			filtered = append(filtered, msg)
		}
	}

	json.NewEncoder(w).Encode(filtered)
}

func handleMedia(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	filename := vars["filename"]
	
	mediaPath := filepath.Join(bridge.dataDir, "media", filename)
	http.ServeFile(w, r, mediaPath)
}

func handleSend(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Platform string `json:"platform"`
		To       string `json:"to"`
		Message  string `json:"message"`
	}
	
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), 400)
		return
	}

	if req.Platform == "whatsapp" && bridge.waClient != nil && bridge.connected {
		jid := types.NewJID(req.To, types.DefaultUserServer)
		_, err := bridge.waClient.SendMessage(context.Background(), jid, &waE2E.Message{
			Conversation: &req.Message,
		})
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		json.NewEncoder(w).Encode(map[string]string{"status": "sent"})
	} else {
		http.Error(w, "platform not connected", 400)
	}
}

func main() {
	dataDir := os.Getenv("DATA_DIR")
	if dataDir == "" {
		dataDir = filepath.Join(os.Getenv("HOME"), ".message-bridge")
	}
	os.MkdirAll(dataDir, 0755)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8030"
	}

	bridge = NewMessageBridge(dataDir)

	// Initialize WhatsApp
	log.Println("Initializing WhatsApp...")
	if err := bridge.InitWhatsApp(); err != nil {
		log.Printf("WhatsApp init error: %v", err)
	}

	// HTTP server
	r := mux.NewRouter()
	r.HandleFunc("/status", handleStatus).Methods("GET")
	r.HandleFunc("/qr", handleQR).Methods("GET")
	r.HandleFunc("/messages", handleMessages).Methods("GET")
	r.HandleFunc("/media/{filename}", handleMedia).Methods("GET")
	r.HandleFunc("/send", handleSend).Methods("POST")

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

	go func() {
		log.Printf("Message Bridge listening on :%s", port)
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatal(err)
		}
	}()

	// Wait for shutdown
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	<-c

	log.Println("Shutting down...")
	if bridge.waClient != nil {
		bridge.waClient.Disconnect()
	}
	srv.Shutdown(context.Background())
}