msgmerge/src/services/whatsapp/whatsapp.rs

use std::{collections::{HashMap, HashSet}, fs, io, path::PathBuf, process::exit};

use chrono::NaiveDateTime;
use dialoguer::Input;
use async_trait::async_trait;
use regex::Regex;
use tokio::{fs::File, io::AsyncReadExt};

use crate::{services::service::Service, types::{participant::Participant, project::ProjectRaw, user::User}};

use super::types::{WhatsAppAudio, WhatsAppMessage, WhatsAppParticipant, WhatsAppPhoto, WhatsAppVideo};

pub fn parse_messages(input: &str) -> Vec<WhatsAppMessage> {
    let mut messages = Vec::new();
    let re = Regex::new(r"^(\d{2}/\d{2}/\d{2}), (\d{2}:\d{2}) - (.*?): (.*)").unwrap();
    let mut current_message: Option<WhatsAppMessage> = None;

    for line in input.lines() {
        if let Some(captures) = re.captures(line) {
            let date = captures.get(1).unwrap().as_str();
            let time = captures.get(2).unwrap().as_str();
            let sender = captures.get(3).unwrap().as_str().to_string();
            let content = captures.get(4).unwrap().as_str();

            // Parse date as DD/MM/YY and convert to epoch
            let datetime_str = format!("{} {}", date, time); // DD/MM/YY HH:MM
            let datetime = NaiveDateTime::parse_from_str(&datetime_str, "%d/%m/%y %H:%M").unwrap();
            let epoch = datetime.and_utc().timestamp() as usize;

            // Save the previous message if there is one
            if let Some(msg) = current_message.take() {
                messages.push(msg);
            }

            // Handle media messages
            if content.starts_with("IMG") || content.starts_with("STK") {
                current_message = Some(WhatsAppMessage {
                    sender,
                    epoch,
                    content: None,
                    photos: vec![WhatsAppPhoto {
                        uri: content.split(' ').next().unwrap().to_string(),
                        text_content: String::new(),
                    }],
                    videos: vec![],
                    audio: None,
                });
            } else if content.starts_with("VID") {
                current_message = Some(WhatsAppMessage {
                    sender,
                    epoch,
                    content: None,
                    photos: vec![],
                    videos: vec![WhatsAppVideo {
                        uri: content.split(' ').next().unwrap().to_string(),
                        text_content: String::new(),
                    }],
                    audio: None,
                });
            } else if content.starts_with("PTT") {
                messages.push(WhatsAppMessage {
                    sender,
                    epoch,
                    content: None,
                    photos: vec![],
                    videos: vec![],
                    audio: Some(WhatsAppAudio {
                        uri: content.split(' ').next().unwrap().to_string(),
                        text_content: String::new(),
                    }),
                });
            } else {
                // Regular message
                current_message = Some(WhatsAppMessage {
                    sender,
                    epoch,
                    content: Some(content.to_string()),
                    photos: vec![],
                    videos: vec![],
                    audio: None,
                });
            }
        } else if let Some(msg) = current_message.as_mut() {
            // Handle continuations or appending to media
            if let Some(photo) = msg.photos.last_mut() {
                photo.text_content.push('\n');
                photo.text_content.push_str(line);
            } else if let Some(video) = msg.videos.last_mut() {
                video.text_content.push('\n');
                video.text_content.push_str(line);
            } else if let Some(content) = msg.content.as_mut() {
                content.push('\n');
                content.push_str(line);
            }
        }
    }

    // Add the last message if any
    if let Some(msg) = current_message {
        messages.push(msg);
    }

    messages
}


pub struct WhatsAppData {
    pub project_path: String,
    pub participants: Vec<WhatsAppParticipant>,
    pub messages: Vec<WhatsAppMessage>,
}

pub struct WhatsApp {
    pub data: Option<WhatsAppData>,
}

#[async_trait]
impl Service for WhatsApp {
    fn get_name(&self) -> &'static str {
        "WhatsApp"
    }

    fn setup(&self) -> Box<dyn Service> {
        let export_path: String = Input::new()
            .with_prompt("Enter path for your WhatsApp export. (extracted zip file)")
            .interact()
            .unwrap();

        Box::new(WhatsApp {
            data: Some(WhatsAppData {
                project_path: export_path,
                participants: vec![],
                messages: vec![],
            }),
        })
    }

    async fn load(&mut self) -> io::Result<()> {
        let data = self.data.as_mut().unwrap();
        let project_path = &data.project_path;

        let paths: Vec<PathBuf> = fs::read_dir(project_path).unwrap_or_else(|_| {
            eprintln!("Folder not found!");
            exit(1);
        })
            .filter_map(|entry| entry.ok())
            .filter(|entry| {
                if let Some(extension) = entry.path().extension() {
                    entry.path().is_file() && extension == "txt"
                } else {
                    false
                }
            })
            .map(|entry| entry.path()).collect();

        if paths.len() != 1 {
            eprintln!("The chat file was not found in the WhatsApp export");
            exit(1);
        }

        let file_path = paths.get(0).unwrap();

        let mut file = File::open(file_path).await?;
        let mut contents = String::new();
        file.read_to_string(&mut contents).await?;

        let messages = parse_messages(&contents);

        let mut participants: HashSet<String> = HashSet::new();
        
        for message in &messages {
            participants.insert(message.sender.clone());
        }
        
        data.participants = participants.into_iter().map(|p| {
            WhatsAppParticipant {
                name: p
            }
        }).collect();

        data.messages = messages;

        Ok(())
    }

    fn get_participants(&self) -> Vec<Participant> {
        self.data.as_ref().unwrap().participants.iter().map(|part| part.conv()).collect()
    }

    async fn merge_messages(&self, raw_project: &mut ProjectRaw, user_map: &HashMap<Participant, User>) {
        let data = self.data.as_ref().unwrap();
        let path = PathBuf::from(data.project_path.clone());
        let mut media_files: Vec<(String, String)> = vec![];

        for message in &data.messages {
            let mut media_index_clone = raw_project.project.media_index.clone();
            let msg = message.conv(&mut media_index_clone, path.clone(), user_map);

            let moved = raw_project.project.push_msg(&msg);

            if moved {
                // Move media files to the media folder
                for photo in msg.photos.clone() {
                    let original_path = photo.original_img_path.unwrap().clone();
                    let path = format!("media/{}", photo.img_path);
                    let photovalues = (path, original_path);
                    media_files.push(photovalues);
                }

                for video in msg.videos.clone() {
                    let original_path = video.original_vid_path.unwrap().clone();
                    let path = format!("media/{}", video.vid_path);
                    let photovalues = (path, original_path);
                    media_files.push(photovalues);
                }

                for audio in msg.audio.clone() {
                    let original_path = audio.original_audio_path.unwrap().clone();
                    let path = format!("media/{}", audio.audio_path);
                    let photovalues = (path, original_path);
                    media_files.push(photovalues);
                }

                // Update the highest media index of the project if the message was merged
                raw_project.project.media_index = media_index_clone;
            }

        }

        let _ = raw_project.save_media_files(media_files).await;
    }
}