from cryptography.fernet import Fernet
from BunnyCDN.Storage import Storage
from instagrapi import Client
from uuid import uuid4
from PIL import Image
import config
import funcs
import os

def insert_highlight_items(media_ids, highlight_id, title, user_id):
    try:
        db, cursor = config.gen_connection()

        query = "INSERT IGNORE INTO highlights (media_id, highlight_id, title, user_id) VALUES (%s, %s, %s, %s)"

        values = [(media_id, highlight_id, title, user_id) for media_id in media_ids]
        cursor.executemany(query, values)
        db.commit()
        if cursor.rowcount > 0:
            print(f"Added {cursor.rowcount} highlight items to the database.")
    except Exception as e:
        print(f"Failed to add highlight items to the database. Error: {e}")


def upload_to_storage(local_path, server_path):
    try:
        obj_storage = Storage("345697f9-d9aa-4a6b-a5ec8bffc16d-ceaf-453e", "storysave")
        obj_storage.PutFile(local_path, server_path)
        print(f"Uploaded to https://storysave.b-cdn.net/{server_path}")
    except Exception as e:
        print(f"Failed to upload {local_path} to {server_path}. Error: {e}")


def login():
    client = Client()

    if os.path.exists("session_data.json"):
        client.load_settings("session_data.json")
        return client
        
    with open("p.enc", "rb") as encrypted_file:
        encrypted_data = encrypted_file.read()

    fernet = Fernet(open("key.enc", "r").read())
    password = str(fernet.decrypt(encrypted_data), "utf-8")
    username = "olivercury"

    auth = input("Enter your 2FA code (leave blank if not enabled): ")
    client.login(username=username, password=password, verification_code=auth)
    client.dump_settings("session_data.json")

    print("Logged in successfully.")

    return client


def parse_media_data(media_item):
    mediaTypes = {1: 'image', 2: 'video', 8: 'album'}

    fileExtension = "jpg" if media_item.media_type == 1 else "mp4"

    mediaInfo = {}
    mediaInfo['taken_at'] = media_item.taken_at if hasattr(media_item, 'taken_at') else None
    mediaInfo['post_type'] = media_item.product_type
    mediaInfo['media_type'] = mediaTypes[media_item.media_type]
    mediaInfo['user_id'] = media_item.user.pk
    mediaInfo['media_id'] = int(media_item.pk)
    mediaInfo['filename'] = f"{media_item.pk}.{fileExtension}"
    mediaInfo['media_url'] = media_item.thumbnail_url if media_item.media_type == 1 else media_item.video_url

    return mediaInfo


def add_media_to_db(mediaInfo):
    media_id = mediaInfo["media_id"]
    user_id = mediaInfo["user_id"]
    username = mediaInfo["username"]
    date = mediaInfo["taken_at"] if "taken_at" in mediaInfo else None
    media_type = mediaInfo["media_type"]
    post_type = mediaInfo["post_type"]
    duration = mediaInfo.get("duration", 0)
    media_url = mediaInfo["media_url"]
    width = mediaInfo["width"]
    height = mediaInfo["height"]
    filehash = mediaInfo["hash"]

    try:
        db, cursor = config.gen_connection()

        query = """
        INSERT INTO media (user_id, username, date, media_type, post_type, media_url, duration, width, height, media_id, hash)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """

        data = (user_id, username, date, media_type, post_type, media_url, duration, width, height, media_id, filehash)

        cursor.execute(query, data)
        db.commit()
        print(f"Added media for {username} to the database.")
    except Exception as e:
        print(f"Failed to add media for {username} to the database. Error: {e}")


def get_profile_picture(client, user_id, username):
    mediaInfo = {}
    mediaInfo['mediaDir'] = 'profile'
    mediaInfo['username'] = username
    mediaInfo['user_id'] = user_id
    mediaInfo['media_id'] = None
    mediaInfo['media_type'] = 'image'
    mediaInfo['post_type'] = 'profile'
    mediaInfo['media_url'] = client.user_info(user_id).profile_pic_url_hd
    mediaInfo['duration'] = 0
    mediaInfo['filename'] = f"{uuid4()}.jpg"

    return mediaInfo


def get_all_stories(client, user_id, firstImport=False):
    stories = client.user_stories(user_id)
    highlights = client.user_highlights(user_id)

    for highlight in highlights:
        highlight_items = client.highlight_info_v1(highlight.pk).items
        stories.extend(highlight_items)

    parsedStories = []
    for story in stories:
        mediaInfo = parse_media_data(story)
        parsedStories.append(mediaInfo)

    return stories


def get_all_posts(client, user_id):
    posts = client.user_medias(user_id, 36)

    medias = []
    for post in posts:
        if post.media_type == 8:
            for item in post.resources:
                medias.append(item)
            continue
        medias.append(post)
        
    parsedPosts = []
    for media in medias:
        mediaInfo = parse_media_data(media)
        parsedPosts.append(mediaInfo)
        
    return posts

    
if __name__ == "__main__":
    client = login()

    db, cursor = config.gen_connection()

    cursor.execute("SELECT instagram_username, instagram_user_id, favorite FROM following ORDER BY favorite DESC;")
    following = cursor.fetchall()

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

    continueFromLast = input("Continue from the last user? (y/N): ").lower() == "y"
    if continueFromLast:
        cursor.execute("SELECT username FROM media ORDER BY id DESC LIMIT 1;")
        lastUser = cursor.fetchone()
        if lastUser:
            lastUser = lastUser[0]
            for idx, user in enumerate(following):
                if user[0] == lastUser:
                    following = following[idx:]
                    break


    for user in following:
        try:
            firstImport = False
            username, user_id, isFavorite = user

            if not user_id:
                firstImport = True
                user_id = client.user_id_from_username(username)
                cursor.execute("UPDATE following SET instagram_user_id = %s WHERE instagram_username = %s;", (user_id, username))
                db.commit()
                print(f"Updated user ID for {username} to {user_id}")

            profile = get_profile_picture(client, user_id, username)
            allStories = get_all_stories(client, user_id, firstImport)
            allPosts = get_all_posts(client, user_id)
            
            medias = allStories + allPosts

            medias = [parse_media_data(media) for media in medias]
            medias = [media for media in medias if media["media_id"] not in existing_files]
            
            for mediaInfo in medias:
                filePath = os.path.join('media', mediaInfo['post_type'], username, mediaInfo['filename'])
                
                funcs.download_file(mediaInfo['media_url'], filePath)

                mediaInfo["hash"] = funcs.calculate_file_hash(filePath)
                mediaInfo["username"] = username

                if mediaInfo["media_type"] == "image":
                    with Image.open(filePath) as img:
                        mediaInfo["width"], mediaInfo["height"] = img.size
                else:
                    mediaInfo["width"], mediaInfo["height"] = funcs.get_video_dimensions(filePath)
                    mediaInfo["duration"] = funcs.get_video_duration(filePath)

                upload_to_storage(filePath, filePath)
                add_media_to_db(mediaInfo)
                os.remove(filePath)

                existing_files.append(mediaInfo["media_id"])
        except Exception as e:
            if "login_required" in str(e):
                print("Please log in to your account again.")
                os.remove("session_data.json")
                client = login()
            else:
                print("An unexpected error occurred:", e)