from datetime import datetime
import config
import funcs
import cv2
import os

directory = 'storysaver'

def UploadMedia(media):
    media_id = media['media_id']
    username = media['username']
    post_date = media['timestamp']
    user_id = media['user_id']
    filepath = media['filepath']
    highlight_id = media['highlight_id']
    post_type = media['post_type']
    thumbnail_url = None
    phash = None
    
    if media_id and int(media_id) in existing_files:
        print('Duplicate file detected. Removing...')
        os.remove(filepath)
        return True
        
    filename = os.path.basename(filepath)
    file_extension = os.path.splitext(filename)[1].lower()

    media_type = funcs.get_media_type(filename)

    file_hash = funcs.calculate_file_hash(filepath)

    width, height = funcs.get_media_dimensions(filepath)
    
    duration = funcs.get_video_duration(filepath)

    if media_type == 'video':
        try:
            thumbPath = f'temp/{media_id}.jpg'
            cap = cv2.VideoCapture(filepath)
            ret, frame = cap.read()
            cv2.imwrite(thumbPath, frame)
            cap.release()
            obj_storage.PutFile(thumbPath, f'thumbnails/{media_id}.jpg') # slower
            thumbnail_url = f"https://storysave.b-cdn.net/thumbnails/{media_id}.jpg"
            phash = funcs.generate_phash(thumbPath)
            os.remove(thumbPath)
        except:
            print('Error generating thumbnail. Skipping...')
            return False
    elif media_type == 'image':
        phash = funcs.generate_phash(filepath)

    if media_id:
        newFilename = f'{media_id}{file_extension}'
    else:
        newFilename = f'{file_hash}{file_extension}'

    server_path = f'media/{post_type}/{username}/{newFilename}'

    file_url = f"https://storysave.b-cdn.net/{server_path}"

    obj_storage.PutFile(filepath, server_path) # slow as fuck

    if highlight_id:
        newCursor.execute("INSERT IGNORE INTO highlights (highlight_id, user_id, media_id) VALUES (%s, %s, %s)", (highlight_id, user_id, media_id))
        newDB.commit()
        print(f'[{newCursor.rowcount}] added highlight {highlight_id} to user {user_id}')

    query = "INSERT IGNORE INTO media (username, media_type, media_url, width, height, media_id, post_type, date, user_id, hash, filename, duration, thumbnail, phash, platform) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
    values = (username, media_type, file_url, width, height, media_id, post_type, post_date, user_id, file_hash, filename, duration, thumbnail_url, phash, 'instagram')

    newCursor.execute(query, values) # slower
    newDB.commit()
    print(f'[{newCursor.rowcount}] records updated. File {filename} uploaded to {file_url}')

    os.remove(filepath)

    return True

def get_user_id(username):
    username = username.lower()
    if username in existing_users:
        return existing_users[username]
    
    return None

def get_media():
    medias = []
    post_types = {
        'posts': 'post',
        'stories': 'story',
        'profile': 'profile',
    }
    
    for post_type in os.listdir('media'):
        users = os.listdir(f'media/{post_type}')
        for user in users:
            user_path = f'media/{post_type}/{user}'
            for filename in os.listdir(user_path):
                data = {}
                filepath = os.path.join(user_path, filename)
                
                data['post_type'] = post_types[post_type]
                data['username'] = user
                data['timestamp'] = filename.split('__')[-1].split('.')[0] if 'com.instagram.android__' in filename else datetime.now()
                if 'com.instagram.android__' in filename:
                    data['timestamp'] = datetime.strptime(data, '%Y%m%d%H%M%S%f')
                data['filepath'] = filepath
                data['media_id'] = None
                data['user_id'] = get_user_id(data['username'])
                data['highlight_id'] = None
                medias.append(data)
    
    return medias

def dump_instagram():
    medias = get_media()

    for media in medias:
        UploadMedia(media)
        existing_files.append(media['media_id'])

if __name__ == '__main__':
    print('Starting processing...')

    if not os.listdir(directory):
        print('No files to process. Exiting...')
        exit()

    newDB, newCursor = config.gen_connection()

    obj_storage = config.get_storage()

    newCursor.execute("SELECT media_id FROM media WHERE media_id IS NOT NULL")
    existing_files = [image[0] for image in newCursor.fetchall()]

    newCursor.execute("SELECT DISTINCT username, user_id FROM media WHERE user_id IS NOT NULL")
    existing_users = {user[0].lower(): user[1].lower() for user in newCursor.fetchall()}
    
    dump_instagram(directory)

    print("Processing completed.")