feat: handle incoming media (images, audio, videos, documents) and show names in group messages (#34)

* added support for recieving media + improved group printing

* fix formatting

whatsapp-bridge/main.go

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"path/filepath"
 	"reflect"
 	"strings"
 	"syscall"
@@ -32,10 +33,12 @@ import (
 
 // Message represents a chat message for our client
 type Message struct {
-	Time     time.Time
-	Sender   string
-	Content  string
-	IsFromMe bool
+	Time      time.Time
+	Sender    string
+	Content   string
+	IsFromMe  bool
+	MediaType string
+	Filename  string
 }
 
 // Database handler for storing message history
@@ -71,6 +74,13 @@ func NewMessageStore() (*MessageStore, error) {
 			content TEXT,
 			timestamp TIMESTAMP,
 			is_from_me BOOLEAN,
+			media_type TEXT,
+			filename TEXT,
+			url TEXT,
+			media_key BLOB,
+			file_sha256 BLOB,
+			file_enc_sha256 BLOB,
+			file_length INTEGER,
 			PRIMARY KEY (id, chat_jid),
 			FOREIGN KEY (chat_jid) REFERENCES chats(jid)
 		);
@@ -98,15 +108,18 @@ func (store *MessageStore) StoreChat(jid, name string, lastMessageTime time.Time
 }
 
 // Store a message in the database
-func (store *MessageStore) StoreMessage(id, chatJID, sender, content string, timestamp time.Time, isFromMe bool) error {
-	// Only store if there's actual content
-	if content == "" {
+func (store *MessageStore) StoreMessage(id, chatJID, sender, content string, timestamp time.Time, isFromMe bool,
+	mediaType, filename, url string, mediaKey, fileSHA256, fileEncSHA256 []byte, fileLength uint64) error {
+	// Only store if there's actual content or media
+	if content == "" && mediaType == "" {
 		return nil
 	}
 
 	_, err := store.db.Exec(
-		"INSERT OR REPLACE INTO messages (id, chat_jid, sender, content, timestamp, is_from_me) VALUES (?, ?, ?, ?, ?, ?)",
-		id, chatJID, sender, content, timestamp, isFromMe,
+		`INSERT OR REPLACE INTO messages 
+		(id, chat_jid, sender, content, timestamp, is_from_me, media_type, filename, url, media_key, file_sha256, file_enc_sha256, file_length) 
+		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+		id, chatJID, sender, content, timestamp, isFromMe, mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength,
 	)
 	return err
 }
@@ -114,7 +127,7 @@ func (store *MessageStore) StoreMessage(id, chatJID, sender, content string, tim
 // Get messages from a chat
 func (store *MessageStore) GetMessages(chatJID string, limit int) ([]Message, error) {
 	rows, err := store.db.Query(
-		"SELECT sender, content, timestamp, is_from_me FROM messages WHERE chat_jid = ? ORDER BY timestamp DESC LIMIT ?",
+		"SELECT sender, content, timestamp, is_from_me, media_type, filename FROM messages WHERE chat_jid = ? ORDER BY timestamp DESC LIMIT ?",
 		chatJID, limit,
 	)
 	if err != nil {
@@ -126,7 +139,7 @@ func (store *MessageStore) GetMessages(chatJID string, limit int) ([]Message, er
 	for rows.Next() {
 		var msg Message
 		var timestamp time.Time
-		err := rows.Scan(&msg.Sender, &msg.Content, &timestamp, &msg.IsFromMe)
+		err := rows.Scan(&msg.Sender, &msg.Content, &timestamp, &msg.IsFromMe, &msg.MediaType, &msg.Filename)
 		if err != nil {
 			return nil, err
 		}
@@ -358,8 +371,312 @@ func sendWhatsAppMessage(client *whatsmeow.Client, recipient string, message str
 	return true, fmt.Sprintf("Message sent to %s", recipient)
 }
 
+// Extract media info from a message
+func extractMediaInfo(msg *waProto.Message) (mediaType string, filename string, url string, mediaKey []byte, fileSHA256 []byte, fileEncSHA256 []byte, fileLength uint64) {
+	if msg == nil {
+		return "", "", "", nil, nil, nil, 0
+	}
+
+	// Check for image message
+	if img := msg.GetImageMessage(); img != nil {
+		return "image", "image_" + time.Now().Format("20060102_150405") + ".jpg",
+			img.GetURL(), img.GetMediaKey(), img.GetFileSHA256(), img.GetFileEncSHA256(), img.GetFileLength()
+	}
+
+	// Check for video message
+	if vid := msg.GetVideoMessage(); vid != nil {
+		return "video", "video_" + time.Now().Format("20060102_150405") + ".mp4",
+			vid.GetURL(), vid.GetMediaKey(), vid.GetFileSHA256(), vid.GetFileEncSHA256(), vid.GetFileLength()
+	}
+
+	// Check for audio message
+	if aud := msg.GetAudioMessage(); aud != nil {
+		return "audio", "audio_" + time.Now().Format("20060102_150405") + ".ogg",
+			aud.GetURL(), aud.GetMediaKey(), aud.GetFileSHA256(), aud.GetFileEncSHA256(), aud.GetFileLength()
+	}
+
+	// Check for document message
+	if doc := msg.GetDocumentMessage(); doc != nil {
+		filename := doc.GetFileName()
+		if filename == "" {
+			filename = "document_" + time.Now().Format("20060102_150405")
+		}
+		return "document", filename,
+			doc.GetURL(), doc.GetMediaKey(), doc.GetFileSHA256(), doc.GetFileEncSHA256(), doc.GetFileLength()
+	}
+
+	return "", "", "", nil, nil, nil, 0
+}
+
+// Handle regular incoming messages with media support
+func handleMessage(client *whatsmeow.Client, messageStore *MessageStore, msg *events.Message, logger waLog.Logger) {
+	// Save message to database
+	chatJID := msg.Info.Chat.String()
+	sender := msg.Info.Sender.User
+
+	// Get appropriate chat name (pass nil for conversation since we don't have one for regular messages)
+	name := GetChatName(client, messageStore, msg.Info.Chat, chatJID, nil, sender, logger)
+
+	// Update chat in database with the message timestamp (keeps last message time updated)
+	err := messageStore.StoreChat(chatJID, name, msg.Info.Timestamp)
+	if err != nil {
+		logger.Warnf("Failed to store chat: %v", err)
+	}
+
+	// Extract text content
+	content := extractTextContent(msg.Message)
+
+	// Extract media info
+	mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength := extractMediaInfo(msg.Message)
+
+	// Skip if there's no content and no media
+	if content == "" && mediaType == "" {
+		return
+	}
+
+	// Store message in database
+	err = messageStore.StoreMessage(
+		msg.Info.ID,
+		chatJID,
+		sender,
+		content,
+		msg.Info.Timestamp,
+		msg.Info.IsFromMe,
+		mediaType,
+		filename,
+		url,
+		mediaKey,
+		fileSHA256,
+		fileEncSHA256,
+		fileLength,
+	)
+
+	if err != nil {
+		logger.Warnf("Failed to store message: %v", err)
+	} else {
+		// Log message reception
+		timestamp := msg.Info.Timestamp.Format("2006-01-02 15:04:05")
+		direction := "←"
+		if msg.Info.IsFromMe {
+			direction = "→"
+		}
+
+		// Log based on message type
+		if mediaType != "" {
+			fmt.Printf("[%s] %s %s: [%s: %s] %s\n", timestamp, direction, sender, mediaType, filename, content)
+		} else if content != "" {
+			fmt.Printf("[%s] %s %s: %s\n", timestamp, direction, sender, content)
+		}
+	}
+}
+
+// DownloadMediaRequest represents the request body for the download media API
+type DownloadMediaRequest struct {
+	MessageID string `json:"message_id"`
+	ChatJID   string `json:"chat_jid"`
+}
+
+// DownloadMediaResponse represents the response for the download media API
+type DownloadMediaResponse struct {
+	Success  bool   `json:"success"`
+	Message  string `json:"message"`
+	Filename string `json:"filename,omitempty"`
+	Path     string `json:"path,omitempty"`
+}
+
+// Store additional media info in the database
+func (store *MessageStore) StoreMediaInfo(id, chatJID, url string, mediaKey, fileSHA256, fileEncSHA256 []byte, fileLength uint64) error {
+	_, err := store.db.Exec(
+		"UPDATE messages SET url = ?, media_key = ?, file_sha256 = ?, file_enc_sha256 = ?, file_length = ? WHERE id = ? AND chat_jid = ?",
+		url, mediaKey, fileSHA256, fileEncSHA256, fileLength, id, chatJID,
+	)
+	return err
+}
+
+// Get media info from the database
+func (store *MessageStore) GetMediaInfo(id, chatJID string) (string, string, string, []byte, []byte, []byte, uint64, error) {
+	var mediaType, filename, url string
+	var mediaKey, fileSHA256, fileEncSHA256 []byte
+	var fileLength uint64
+
+	err := store.db.QueryRow(
+		"SELECT media_type, filename, url, media_key, file_sha256, file_enc_sha256, file_length FROM messages WHERE id = ? AND chat_jid = ?",
+		id, chatJID,
+	).Scan(&mediaType, &filename, &url, &mediaKey, &fileSHA256, &fileEncSHA256, &fileLength)
+
+	return mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength, err
+}
+
+// MediaDownloader implements the whatsmeow.DownloadableMessage interface
+type MediaDownloader struct {
+	URL           string
+	DirectPath    string
+	MediaKey      []byte
+	FileLength    uint64
+	FileSHA256    []byte
+	FileEncSHA256 []byte
+	MediaType     whatsmeow.MediaType
+}
+
+// GetDirectPath implements the DownloadableMessage interface
+func (d *MediaDownloader) GetDirectPath() string {
+	return d.DirectPath
+}
+
+// GetURL implements the DownloadableMessage interface
+func (d *MediaDownloader) GetURL() string {
+	return d.URL
+}
+
+// GetMediaKey implements the DownloadableMessage interface
+func (d *MediaDownloader) GetMediaKey() []byte {
+	return d.MediaKey
+}
+
+// GetFileLength implements the DownloadableMessage interface
+func (d *MediaDownloader) GetFileLength() uint64 {
+	return d.FileLength
+}
+
+// GetFileSHA256 implements the DownloadableMessage interface
+func (d *MediaDownloader) GetFileSHA256() []byte {
+	return d.FileSHA256
+}
+
+// GetFileEncSHA256 implements the DownloadableMessage interface
+func (d *MediaDownloader) GetFileEncSHA256() []byte {
+	return d.FileEncSHA256
+}
+
+// GetMediaType implements the DownloadableMessage interface
+func (d *MediaDownloader) GetMediaType() whatsmeow.MediaType {
+	return d.MediaType
+}
+
+// Function to download media from a message
+func downloadMedia(client *whatsmeow.Client, messageStore *MessageStore, messageID, chatJID string) (bool, string, string, string, error) {
+	// Query the database for the message
+	var mediaType, filename, url string
+	var mediaKey, fileSHA256, fileEncSHA256 []byte
+	var fileLength uint64
+	var err error
+
+	// First, check if we already have this file
+	chatDir := fmt.Sprintf("store/%s", strings.ReplaceAll(chatJID, ":", "_"))
+	localPath := ""
+
+	// Get media info from the database
+	mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength, err = messageStore.GetMediaInfo(messageID, chatJID)
+
+	if err != nil {
+		// Try to get basic info if extended info isn't available
+		err = messageStore.db.QueryRow(
+			"SELECT media_type, filename FROM messages WHERE id = ? AND chat_jid = ?",
+			messageID, chatJID,
+		).Scan(&mediaType, &filename)
+
+		if err != nil {
+			return false, "", "", "", fmt.Errorf("failed to find message: %v", err)
+		}
+	}
+
+	// Check if this is a media message
+	if mediaType == "" {
+		return false, "", "", "", fmt.Errorf("not a media message")
+	}
+
+	// Create directory for the chat if it doesn't exist
+	if err := os.MkdirAll(chatDir, 0755); err != nil {
+		return false, "", "", "", fmt.Errorf("failed to create chat directory: %v", err)
+	}
+
+	// Generate a local path for the file
+	localPath = fmt.Sprintf("%s/%s", chatDir, filename)
+
+	// Get absolute path
+	absPath, err := filepath.Abs(localPath)
+	if err != nil {
+		return false, "", "", "", fmt.Errorf("failed to get absolute path: %v", err)
+	}
+
+	// Check if file already exists
+	if _, err := os.Stat(localPath); err == nil {
+		// File exists, return it
+		return true, mediaType, filename, absPath, nil
+	}
+
+	// If we don't have all the media info we need, we can't download
+	if url == "" || len(mediaKey) == 0 || len(fileSHA256) == 0 || len(fileEncSHA256) == 0 || fileLength == 0 {
+		return false, "", "", "", fmt.Errorf("incomplete media information for download")
+	}
+
+	fmt.Printf("Attempting to download media for message %s in chat %s...\n", messageID, chatJID)
+
+	// Extract direct path from URL
+	directPath := extractDirectPathFromURL(url)
+
+	// Create a downloader that implements DownloadableMessage
+	var waMediaType whatsmeow.MediaType
+	switch mediaType {
+	case "image":
+		waMediaType = whatsmeow.MediaImage
+	case "video":
+		waMediaType = whatsmeow.MediaVideo
+	case "audio":
+		waMediaType = whatsmeow.MediaAudio
+	case "document":
+		waMediaType = whatsmeow.MediaDocument
+	default:
+		return false, "", "", "", fmt.Errorf("unsupported media type: %s", mediaType)
+	}
+
+	downloader := &MediaDownloader{
+		URL:           url,
+		DirectPath:    directPath,
+		MediaKey:      mediaKey,
+		FileLength:    fileLength,
+		FileSHA256:    fileSHA256,
+		FileEncSHA256: fileEncSHA256,
+		MediaType:     waMediaType,
+	}
+
+	// Download the media using whatsmeow client
+	mediaData, err := client.Download(downloader)
+	if err != nil {
+		return false, "", "", "", fmt.Errorf("failed to download media: %v", err)
+	}
+
+	// Save the downloaded media to file
+	if err := os.WriteFile(localPath, mediaData, 0644); err != nil {
+		return false, "", "", "", fmt.Errorf("failed to save media file: %v", err)
+	}
+
+	fmt.Printf("Successfully downloaded %s media to %s (%d bytes)\n", mediaType, absPath, len(mediaData))
+	return true, mediaType, filename, absPath, nil
+}
+
+// Extract direct path from a WhatsApp media URL
+func extractDirectPathFromURL(url string) string {
+	// The direct path is typically in the URL, we need to extract it
+	// Example URL: https://mmg.whatsapp.net/v/t62.7118-24/13812002_698058036224062_3424455886509161511_n.enc?ccb=11-4&oh=...
+
+	// Find the path part after the domain
+	parts := strings.SplitN(url, ".net/", 2)
+	if len(parts) < 2 {
+		return url // Return original URL if parsing fails
+	}
+
+	pathPart := parts[1]
+
+	// Remove query parameters
+	pathPart = strings.SplitN(pathPart, "?", 2)[0]
+
+	// Create proper direct path format
+	return "/" + pathPart
+}
+
 // Start a REST API server to expose the WhatsApp client functionality
-func startRESTServer(client *whatsmeow.Client, port int) {
+func startRESTServer(client *whatsmeow.Client, messageStore *MessageStore, port int) {
 	// Handler for sending messages
 	http.HandleFunc("/api/send", func(w http.ResponseWriter, r *http.Request) {
 		// Only allow POST requests
@@ -406,6 +723,57 @@ func startRESTServer(client *whatsmeow.Client, port int) {
 		})
 	})
 
+	// Handler for downloading media
+	http.HandleFunc("/api/download", func(w http.ResponseWriter, r *http.Request) {
+		// Only allow POST requests
+		if r.Method != http.MethodPost {
+			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+			return
+		}
+
+		// Parse the request body
+		var req DownloadMediaRequest
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			http.Error(w, "Invalid request format", http.StatusBadRequest)
+			return
+		}
+
+		// Validate request
+		if req.MessageID == "" || req.ChatJID == "" {
+			http.Error(w, "Message ID and Chat JID are required", http.StatusBadRequest)
+			return
+		}
+
+		// Download the media
+		success, mediaType, filename, path, err := downloadMedia(client, messageStore, req.MessageID, req.ChatJID)
+
+		// Set response headers
+		w.Header().Set("Content-Type", "application/json")
+
+		// Handle download result
+		if !success || err != nil {
+			errMsg := "Unknown error"
+			if err != nil {
+				errMsg = err.Error()
+			}
+
+			w.WriteHeader(http.StatusInternalServerError)
+			json.NewEncoder(w).Encode(DownloadMediaResponse{
+				Success: false,
+				Message: fmt.Sprintf("Failed to download media: %s", errMsg),
+			})
+			return
+		}
+
+		// Send successful response
+		json.NewEncoder(w).Encode(DownloadMediaResponse{
+			Success:  true,
+			Message:  fmt.Sprintf("Successfully downloaded %s media", mediaType),
+			Filename: filename,
+			Path:     path,
+		})
+	})
+
 	// Start the server
 	serverAddr := fmt.Sprintf(":%d", port)
 	fmt.Printf("Starting REST API server on %s...\n", serverAddr)
@@ -538,7 +906,7 @@ func main() {
 	fmt.Println("\n✓ Connected to WhatsApp! Type 'help' for commands.")
 
 	// Start REST API server
-	startRESTServer(client, 8080)
+	startRESTServer(client, messageStore, 8080)
 
 	// Create a channel to keep the main goroutine alive
 	exitChan := make(chan os.Signal, 1)
@@ -637,49 +1005,6 @@ func GetChatName(client *whatsmeow.Client, messageStore *MessageStore, jid types
 	return name
 }
 
-// Handle regular incoming messages
-func handleMessage(client *whatsmeow.Client, messageStore *MessageStore, msg *events.Message, logger waLog.Logger) {
-	// Extract text content
-	content := extractTextContent(msg.Message)
-	if content == "" {
-		return // Skip non-text messages
-	}
-
-	// Save message to database
-	chatJID := msg.Info.Chat.String()
-	sender := msg.Info.Sender.User
-
-	// Get appropriate chat name (pass nil for conversation since we don't have one for regular messages)
-	name := GetChatName(client, messageStore, msg.Info.Chat, chatJID, nil, sender, logger)
-
-	// Update chat in database with the message timestamp (keeps last message time updated)
-	err := messageStore.StoreChat(chatJID, name, msg.Info.Timestamp)
-	if err != nil {
-		logger.Warnf("Failed to store chat: %v", err)
-	}
-
-	// Store message in database
-	err = messageStore.StoreMessage(
-		msg.Info.ID,
-		chatJID,
-		sender,
-		content,
-		msg.Info.Timestamp,
-		msg.Info.IsFromMe,
-	)
-	if err != nil {
-		logger.Warnf("Failed to store message: %v", err)
-	} else {
-		// Log message reception
-		timestamp := msg.Info.Timestamp.Format("2006-01-02 15:04:05")
-		direction := "←"
-		if msg.Info.IsFromMe {
-			direction = "→"
-		}
-		fmt.Printf("[%s] %s %s: %s\n", timestamp, direction, sender, content)
-	}
-}
-
 // Handle history sync events
 func handleHistorySync(client *whatsmeow.Client, messageStore *MessageStore, historySync *events.HistorySync, logger waLog.Logger) {
 	fmt.Printf("Received history sync event with %d conversations\n", len(historySync.Data.Conversations))
@@ -738,11 +1063,20 @@ func handleHistorySync(client *whatsmeow.Client, messageStore *MessageStore, his
 					}
 				}
 
-				// Log the message content for debugging
-				logger.Infof("Message content: %v", content)
+				// Extract media info
+				var mediaType, filename, url string
+				var mediaKey, fileSHA256, fileEncSHA256 []byte
+				var fileLength uint64
 
-				// Skip non-text messages
-				if content == "" {
+				if msg.Message.Message != nil {
+					mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength = extractMediaInfo(msg.Message.Message)
+				}
+
+				// Log the message content for debugging
+				logger.Infof("Message content: %v, Media Type: %v", content, mediaType)
+
+				// Skip messages with no content and no media
+				if content == "" && mediaType == "" {
 					continue
 				}
 
@@ -785,19 +1119,32 @@ func handleHistorySync(client *whatsmeow.Client, messageStore *MessageStore, his
 					content,
 					timestamp,
 					isFromMe,
+					mediaType,
+					filename,
+					url,
+					mediaKey,
+					fileSHA256,
+					fileEncSHA256,
+					fileLength,
 				)
 				if err != nil {
 					logger.Warnf("Failed to store history message: %v", err)
 				} else {
 					syncedCount++
 					// Log successful message storage
-					logger.Infof("Stored message: [%s] %s -> %s: %s", timestamp.Format("2006-01-02 15:04:05"), sender, chatJID, content)
+					if mediaType != "" {
+						logger.Infof("Stored message: [%s] %s -> %s: [%s: %s] %s",
+							timestamp.Format("2006-01-02 15:04:05"), sender, chatJID, mediaType, filename, content)
+					} else {
+						logger.Infof("Stored message: [%s] %s -> %s: %s",
+							timestamp.Format("2006-01-02 15:04:05"), sender, chatJID, content)
+					}
 				}
 			}
 		}
 	}
 
-	fmt.Printf("History sync complete. Stored %d text messages.\n", syncedCount)
+	fmt.Printf("History sync complete. Stored %d messages.\n", syncedCount)
 }
 
 // Request history sync from the server

whatsapp-mcp-server/main.py

@@ -11,7 +11,8 @@ from whatsapp import (
     get_message_context as whatsapp_get_message_context,
     send_message as whatsapp_send_message,
     send_file as whatsapp_send_file,
-    send_audio_message as whatsapp_audio_voice_message
+    send_audio_message as whatsapp_audio_voice_message,
+    download_media as whatsapp_download_media
 )
 
 # Initialize FastMCP server
@@ -128,7 +129,7 @@ def get_contact_chats(jid: str, limit: int = 20, page: int = 0) -> List[Dict[str
     return chats
 
 @mcp.tool()
-def get_last_interaction(jid: str) -> Dict[str, Any]:
+def get_last_interaction(jid: str) -> str:
     """Get most recent WhatsApp message involving the contact.
     
     Args:
@@ -220,6 +221,31 @@ def send_audio_message(recipient: str, media_path: str) -> Dict[str, Any]:
         "message": status_message
     }
 
+@mcp.tool()
+def download_media(message_id: str, chat_jid: str) -> Dict[str, Any]:
+    """Download media from a WhatsApp message and get the local file path.
+    
+    Args:
+        message_id: The ID of the message containing the media
+        chat_jid: The JID of the chat containing the message
+    
+    Returns:
+        A dictionary containing success status, a status message, and the file path if successful
+    """
+    file_path = whatsapp_download_media(message_id, chat_jid)
+    
+    if file_path:
+        return {
+            "success": True,
+            "message": "Media downloaded successfully",
+            "file_path": file_path
+        }
+    else:
+        return {
+            "success": False,
+            "message": "Failed to download media"
+        }
+
 if __name__ == "__main__":
     # Initialize and run the server
     mcp.run(transport='stdio')

whatsapp-mcp-server/whatsapp.py

@@ -19,6 +19,7 @@ class Message:
     chat_jid: str
     id: str
     chat_name: Optional[str] = None
+    media_type: Optional[str] = None
 
 @dataclass
 class Chat:
@@ -46,146 +47,80 @@ class MessageContext:
     before: List[Message]
     after: List[Message]
 
-def print_message(message: Message, show_chat_info: bool = True) -> None:
-    """Print a single message with consistent formatting."""
-    direction = "→" if message.is_from_me else "←"
-    
-    if show_chat_info and message.chat_name:
-        print(f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] {direction} Chat: {message.chat_name} ({message.chat_jid})")
-    else:
-        print(f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] {direction}")
-        
-    print(f"From: {'Me' if message.is_from_me else message.sender}")
-    print(f"Message: {message.content}")
-    print("-" * 100)
-
-def print_messages_list(messages: List[Message], title: str = "", show_chat_info: bool = True) -> None:
-    """Print a list of messages with a title and consistent formatting."""
-    if not messages:
-        print("No messages to display.")
-        return
-        
-    if title:
-        print(f"\n{title}")
-        print("-" * 100)
-    
-    for message in messages:
-        print_message(message, show_chat_info)
-
-def print_chat(chat: Chat) -> None:
-    """Print a single chat with consistent formatting."""
-    print(f"Chat: {chat.name} ({chat.jid})")
-    if chat.last_message_time:
-        print(f"Last active: {chat.last_message_time:%Y-%m-%d %H:%M:%S}")
-        direction = "→" if chat.last_is_from_me else "←"
-        sender = "Me" if chat.last_is_from_me else chat.last_sender
-        print(f"Last message: {direction} {sender}: {chat.last_message}")
-    print("-" * 100)
-
-def print_chats_list(chats: List[Chat], title: str = "") -> None:
-    """Print a list of chats with a title and consistent formatting."""
-    if not chats:
-        print("No chats to display.")
-        return
-        
-    if title:
-        print(f"\n{title}")
-        print("-" * 100)
-    
-    for chat in chats:
-        print_chat(chat)
-
-def print_paginated_messages(messages: List[Message], page: int, total_pages: int, chat_name: str) -> None:
-    """Print a paginated list of messages with navigation hints."""
-    print(f"\nMessages for chat: {chat_name}")
-    print(f"Page {page} of {total_pages}")
-    print("-" * 100)
-    
-    print_messages_list(messages, show_chat_info=False)
-    
-    # Print pagination info
-    if page > 1:
-        print(f"Use page={page-1} to see newer messages")
-    if page < total_pages:
-        print(f"Use page={page+1} to see older messages")
-
-"""
-CREATE TABLE messages (
-			id TEXT,
-			chat_jid TEXT,
-			sender TEXT,
-			content TEXT,
-			timestamp TIMESTAMP,
-			is_from_me BOOLEAN,
-			PRIMARY KEY (id, chat_jid),
-			FOREIGN KEY (chat_jid) REFERENCES chats(jid)
-		)
-
-CREATE TABLE chats (
-			jid TEXT PRIMARY KEY,
-			name TEXT,
-			last_message_time TIMESTAMP
-		)
-"""
-
-def print_recent_messages(limit=10) -> List[Message]:
+def get_sender_name(sender_jid: str) -> str:
     try:
-        # Connect to the SQLite database
         conn = sqlite3.connect(MESSAGES_DB_PATH)
         cursor = conn.cursor()
         
-        # Query recent messages with chat info
-        query = """
-        SELECT 
-            m.timestamp,
-            m.sender,
-            c.name,
-            m.content,
-            m.is_from_me,
-            c.jid,
-            m.id
-        FROM messages m
-        JOIN chats c ON m.chat_jid = c.jid
-        ORDER BY m.timestamp DESC
-        LIMIT ?
-        """
+        # First try matching by exact JID
+        cursor.execute("""
+            SELECT name
+            FROM chats
+            WHERE jid = ?
+            LIMIT 1
+        """, (sender_jid,))
         
-        cursor.execute(query, (limit,))
-        messages = cursor.fetchall()
+        result = cursor.fetchone()
         
-        if not messages:
-            print("No messages found in the database.")
-            return []
+        # If no result, try looking for the number within JIDs
+        if not result:
+            # Extract the phone number part if it's a JID
+            if '@' in sender_jid:
+                phone_part = sender_jid.split('@')[0]
+            else:
+                phone_part = sender_jid
+                
+            cursor.execute("""
+                SELECT name
+                FROM chats
+                WHERE jid LIKE ?
+                LIMIT 1
+            """, (f"%{phone_part}%",))
             
-        result = []
+            result = cursor.fetchone()
         
-        # Convert to Message objects
-        for msg in messages:
-            message = Message(
-                timestamp=datetime.fromisoformat(msg[0]),
-                sender=msg[1],
-                chat_name=msg[2] or "Unknown Chat",
-                content=msg[3],
-                is_from_me=msg[4],
-                chat_jid=msg[5],
-                id=msg[6]
-            )
-            result.append(message)
+        if result and result[0]:
+            return result[0]
+        else:
+            return sender_jid
         
-        # Print messages using helper function
-        print_messages_list(result, title=f"Last {limit} messages:")
-        return result
-            
     except sqlite3.Error as e:
-        print(f"Error accessing database: {e}")
-        return []
-    except Exception as e:
-        print(f"Unexpected error: {e}")
-        return []
+        print(f"Database error while getting sender name: {e}")
+        return sender_jid
     finally:
         if 'conn' in locals():
             conn.close()
 
+def format_message(message: Message, show_chat_info: bool = True) -> None:
+    """Print a single message with consistent formatting."""
+    output = ""
+    
+    if show_chat_info and message.chat_name:
+        output += f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] Chat: {message.chat_name} "
+    else:
+        output += f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] "
+        
+    content_prefix = ""
+    if hasattr(message, 'media_type') and message.media_type:
+        content_prefix = f"[{message.media_type} - Message ID: {message.id} - Chat JID: {message.chat_jid}] "
+    
+    try:
+        sender_name = get_sender_name(message.sender) if not message.is_from_me else "Me"
+        output += f"From: {sender_name}: {content_prefix}{message.content}\n"
+    except Exception as e:
+        print(f"Error formatting message: {e}")
+    return output
+
+def format_messages_list(messages: List[Message], show_chat_info: bool = True) -> None:
+    output = ""
+    if not messages:
+        output += "No messages to display."
+        return output
+    
+    for message in messages:
+        output += format_message(message, show_chat_info)
+    return output
+
 def list_messages(
     after: Optional[str] = None,
     before: Optional[str] = None,
@@ -204,7 +139,7 @@ def list_messages(
         cursor = conn.cursor()
         
         # Build base query
-        query_parts = ["SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id FROM messages"]
+        query_parts = ["SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.media_type FROM messages"]
         query_parts.append("JOIN chats ON messages.chat_jid = chats.jid")
         where_clauses = []
         params = []
@@ -261,7 +196,8 @@ def list_messages(
                 content=msg[3],
                 is_from_me=msg[4],
                 chat_jid=msg[5],
-                id=msg[6]
+                id=msg[6],
+                media_type=msg[7]
             )
             result.append(message)
             
@@ -273,9 +209,11 @@ def list_messages(
                 messages_with_context.extend(context.before)
                 messages_with_context.append(context.message)
                 messages_with_context.extend(context.after)
-            return messages_with_context
             
-        return result
+            return format_messages_list(messages_with_context, show_chat_info=True)
+            
+        # Format and display messages without context
+        return format_messages_list(result, show_chat_info=True)    
         
     except sqlite3.Error as e:
         print(f"Database error: {e}")
@@ -297,7 +235,7 @@ def get_message_context(
         
         # Get the target message first
         cursor.execute("""
-            SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.chat_jid
+            SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.chat_jid, messages.media_type
             FROM messages
             JOIN chats ON messages.chat_jid = chats.jid
             WHERE messages.id = ?
@@ -314,12 +252,13 @@ def get_message_context(
             content=msg_data[3],
             is_from_me=msg_data[4],
             chat_jid=msg_data[5],
-            id=msg_data[6]
+            id=msg_data[6],
+            media_type=msg_data[8]
         )
         
         # Get messages before
         cursor.execute("""
-            SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id
+            SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.media_type
             FROM messages
             JOIN chats ON messages.chat_jid = chats.jid
             WHERE messages.chat_jid = ? AND messages.timestamp < ?
@@ -336,12 +275,13 @@ def get_message_context(
                 content=msg[3],
                 is_from_me=msg[4],
                 chat_jid=msg[5],
-                id=msg[6]
+                id=msg[6],
+                media_type=msg[7]
             ))
         
         # Get messages after
         cursor.execute("""
-            SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id
+            SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.media_type
             FROM messages
             JOIN chats ON messages.chat_jid = chats.jid
             WHERE messages.chat_jid = ? AND messages.timestamp > ?
@@ -358,7 +298,8 @@ def get_message_context(
                 content=msg[3],
                 is_from_me=msg[4],
                 chat_jid=msg[5],
-                id=msg[6]
+                id=msg[6],
+                media_type=msg[7]
             ))
         
         return MessageContext(
@@ -542,7 +483,7 @@ def get_contact_chats(jid: str, limit: int = 20, page: int = 0) -> List[Chat]:
             conn.close()
 
 
-def get_last_interaction(jid: str) -> Optional[Message]:
+def get_last_interaction(jid: str) -> str:
     """Get most recent message involving the contact."""
     try:
         conn = sqlite3.connect(MESSAGES_DB_PATH)
@@ -556,7 +497,8 @@ def get_last_interaction(jid: str) -> Optional[Message]:
                 m.content,
                 m.is_from_me,
                 c.jid,
-                m.id
+                m.id,
+                m.media_type
             FROM messages m
             JOIN chats c ON m.chat_jid = c.jid
             WHERE m.sender = ? OR c.jid = ?
@@ -569,16 +511,19 @@ def get_last_interaction(jid: str) -> Optional[Message]:
         if not msg_data:
             return None
             
-        return Message(
+        message = Message(
             timestamp=datetime.fromisoformat(msg_data[0]),
             sender=msg_data[1],
             chat_name=msg_data[2],
             content=msg_data[3],
             is_from_me=msg_data[4],
             chat_jid=msg_data[5],
-            id=msg_data[6]
+            id=msg_data[6],
+            media_type=msg_data[7]
         )
         
+        return format_message(message)
+        
     except sqlite3.Error as e:
         print(f"Database error: {e}")
         return None
@@ -778,3 +723,45 @@ def send_audio_message(recipient: str, media_path: str) -> Tuple[bool, str]:
         return False, f"Error parsing response: {response.text}"
     except Exception as e:
         return False, f"Unexpected error: {str(e)}"
+
+def download_media(message_id: str, chat_jid: str) -> Optional[str]:
+    """Download media from a message and return the local file path.
+    
+    Args:
+        message_id: The ID of the message containing the media
+        chat_jid: The JID of the chat containing the message
+    
+    Returns:
+        The local file path if download was successful, None otherwise
+    """
+    try:
+        url = f"{WHATSAPP_API_BASE_URL}/download"
+        payload = {
+            "message_id": message_id,
+            "chat_jid": chat_jid
+        }
+        
+        response = requests.post(url, json=payload)
+        
+        if response.status_code == 200:
+            result = response.json()
+            if result.get("success", False):
+                path = result.get("path")
+                print(f"Media downloaded successfully: {path}")
+                return path
+            else:
+                print(f"Download failed: {result.get('message', 'Unknown error')}")
+                return None
+        else:
+            print(f"Error: HTTP {response.status_code} - {response.text}")
+            return None
+            
+    except requests.RequestException as e:
+        print(f"Request error: {str(e)}")
+        return None
+    except json.JSONDecodeError:
+        print(f"Error parsing response: {response.text}")
+        return None
+    except Exception as e:
+        print(f"Unexpected error: {str(e)}")
+        return None