import os
import shutil
from mega import Mega
import gdown
import re
import wget
import sys
import uuid
import zipfile


class InvalidDriveId(Exception):
    def __init__(self, message="Error de la url"):
        self.message = message
        super().__init__(self.message)


def model_downloader(url, zip_path, dest_path):
    """Download and unzip a file from Google Drive or Mega."""
    
    def drive_download(url, dest_folder):
        print(f"Descargando desde drive...")
        try:
            filename = gdown.download(url, os.path.join(dest_folder, f"{uuid.uuid4()}.zip"), fuzzy=True)
            return os.path.basename(filename)
        except:
            print("El intento de descargar con drive no funcionó")
            return None

    def mega_download(url, dest_folder):
        try:
            file_id = None
            if "#!" in url:
                file_id = url.split("#!")[1].split("!")[0]
            elif "file/" in url:
                file_id = url.split("file/")[1].split("/")[0]
            else:
                file_id = None

            print(f"Descargando desde mega...")
            if file_id:
                mega = Mega()
                m = mega.login()
                filename = m.download_url(url, dest_path=dest_folder, dest_filename=f"{uuid.uuid4()}.zip")

                return os.path.basename(filename)
            else:
                return None

        except Exception as e:
            print("Ocurrio un error**")
            print(e)
            return None

    def download(url, dest_folder):
        try:
            print(f"Descargando desde url generica...")
            dest_path = wget.download(url=url, out=os.path.join(dest_folder, f"{uuid.uuid4()}.zip"))

            return os.path.basename(dest_path)
        except Exception as e:
            print(f"Error al descargar el archivo: {str(e)}")

    filename = ""

    if not os.path.exists(zip_path):
        os.mkdir(zip_path)

    if url and 'drive.google.com' in url:
        # Descargar el elemento si la URL es de Google Drive
        filename = drive_download(url, zip_path)
    elif url and 'mega.nz' in url:
        filename = mega_download(url, zip_path)
    elif url and 'pixeldrain' in url:
        print("No se puede descargar de pixeldrain")
        sys.exit()
    else:
        filename = download(url, zip_path)

    if filename:
        print(f"Descomprimiendo {filename}...")
        modelname = str(filename).replace(".zip", "")
        zip_file_path = os.path.join(zip_path, filename)
        
        try:
            shutil.unpack_archive(zip_file_path, os.path.join(dest_path, modelname))
        except Exception as e:
            try:
                with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                    zip_ref.extractall(dest_path)
            except zipfile.BadZipFile as e:
                print(f"Error: El archivo ZIP no es válido - {e}")
            except Exception as e:
                print(f"Error inesperado: {e}")
        
        if os.path.exists(zip_file_path):
            os.remove(zip_file_path)

        return modelname
    else:
        return None

def get_model(weight_path, modelname):
    resources = {}
    for root, dirs, files in os.walk(os.path.join(weight_path, modelname)):
        for file in files:
            if file.endswith('.index'):
                resources['index'] =  os.path.relpath(os.path.join(root, file))
            if file.endswith('.pth') and not 'G_' in file and not 'D_' in file:
                resources['pth'] =  os.path.relpath(os.path.join(root, file), start=weight_path)
    return resources