import os, requests, config
from snapchat import get_data, get_stories, get_highlight_stories

headers = {"user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36"}
directory = "snapchat"

def get_existing_media_ids(directory):
	# get all files and their their base filename without extension, split the filename by ~ and get the 3rd element
	existing_media_ids = set()
	for root, _, files in os.walk(directory):
		for file in files:
			if '~' not in file:
				continue
			
			filename, _ = os.path.splitext(file)
			media_id = filename.split('~')[2]
			existing_media_ids.add(media_id)
	return existing_media_ids

def get_media_id(url):
	return url.split('/')[-1].split('?')[0].split('.')[0]

def find_duplicate_snap(existing_snaps, snap_id):
	for snap in existing_snaps:
		if snap_id in snap[1]:
			return snap
	return False
	
def main():
	if not os.path.exists(directory):
		os.makedirs(directory)
    
	db, cursor = config.gen_connection()

	cursor.execute("SELECT username FROM following WHERE platform = 'snapchat'")
	usernames = [row[0] for row in cursor.fetchall()]
	
	cursor.execute("SELECT id, filename FROM media WHERE filename IS NOT NULL AND platform = 'snapchat' AND snap_id IS NULL")
	existing_medias = cursor.fetchall()
	
	existing_media_ids = get_existing_media_ids(directory)
	
	for username in usernames:
		print(f"Getting stories for {username}...")
		data = get_data(username)
		if not data:
			continue

		print("Getting stories...")
		stories = get_stories(data)

		print("Getting highlights...")
		stories.extend(get_highlight_stories(data))

		for story in stories:
			media_id = story['media_id']
			url = story['url']
			timestamp = story['timestamp']

			snap_id = get_media_id(url)
			duplicate_snap = find_duplicate_snap(existing_medias, snap_id)
			if duplicate_snap:
				snap_id = get_media_id(url)
				cursor.execute("UPDATE media SET snap_id = %s WHERE id = %s", (snap_id, duplicate_snap[0]))
				db.commit()
				print(f"{cursor.rowcount} Media {snap_id} updated.")
				continue
			
if __name__ == "__main__":
    main()