# wsd.py by mgx.me
# inspired by @wibus-wee's DownloadWeChatStickers_macOS.js

import os
import re
import shutil
import time
import uuid
import subprocess
from pathlib import Path
import requests

def sleep(seconds):
    time.sleep(seconds)

def copy_fav_archive():
    wechat_data_path = Path.home() / 'Library/Containers/com.tencent.xinWeChat/Data/Library/Application Support/com.tencent.xinWeChat/2.0b4.0.9'
    user_dirs = [d for d in wechat_data_path.iterdir() if d.is_dir() and len(d.name) > 30]
    user_dirs = [d for d in user_dirs if (d / 'Stickers/fav.archive').exists()]

    if not user_dirs:
        print('No user folder found')
        return None

    if len(user_dirs) > 1:
        for index, dir_path in enumerate(user_dirs, start=1):
            print(f"{index}. {dir_path}")
        choice = input("Please enter the user folder number: ")
        try:
            user_dir = user_dirs[int(choice) - 1]
        except (IndexError, ValueError):
            print("Invalid choice")
            return None
    else:
        user_dir = user_dirs[0]

    fav_archive_path = user_dir / 'Stickers/fav.archive'
    dest_path = Path.home() / 'Desktop/fav.archive.plist'
    shutil.copy(fav_archive_path, dest_path)
    print(f"Copied fav.archive file to desktop: {dest_path}")
    return dest_path

def convert_to_plist(plist_path):
    subprocess.run(['plutil', '-convert', 'xml1', str(plist_path)], check=True)

def extract_urls(plist_path):
    with open(plist_path, 'r', encoding='utf-8') as file:
        content = file.read()

    link_regex = re.compile(r'<string>(https?:\/\/.*?)<\/string>')
    links = [match.replace('&amp;', '&') for match in link_regex.findall(content)]
    print(f"Found {len(links)} sticker URLs")
    return links

def download_stickers(links, output_dir):
    output_dir.mkdir(parents=True, exist_ok=True)
    user_agent = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/94.0.4606.71 Safari/537.36'

    for index, url in enumerate(links, start=1):
        filename = f"{uuid.uuid4()}.gif"
        filepath = output_dir / filename
        print(f"Downloading sticker {index}: {url}")

        headers = {'User-Agent': user_agent}
        try:
            response = requests.get(url, headers=headers)
            response.raise_for_status()
            with open(filepath, 'wb') as f:
                f.write(response.content)
            print(f"Download completed: {filename}")
        except requests.RequestException as e:
            print(f"Failed to download: {url} Error: {e}")

def main():
    plist_path = copy_fav_archive()
    if not plist_path:
        return

    convert_to_plist(plist_path)
    links = extract_urls(plist_path)

    if links:
        output_dir = Path.home() / 'Desktop' / 'WX-Stickers'
        print('Starting sticker download in 5 seconds...')
        sleep(5)
        download_stickers(links, output_dir)

if __name__ == "__main__":
    main()