import os
import gradio as gr
import shutil
from PIL import Image
import tempfile
import json
import zipfile

# Set the directory where files will be saved
SAVE_DIR = os.path.abspath("saved_media") #make sure its the as media_folder Above

# Ensure the save directory exists
os.makedirs(SAVE_DIR, exist_ok=True)

# Define supported file extensions
SUPPORTED_EXTENSIONS = {
    "image": [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"],
    "audio": [".mp3", ".wav", ".ogg"],
    "video": [".mp4", ".avi", ".mov", ".webm"],
    "model3d": [".glb", ".gltf", ".obj"]
}

def save_file(file):
    if file is None:
        return "No file uploaded.", gr.update()

    try:
        # Get the original filename and extension
        original_filename = os.path.basename(file.name)
        _, extension = os.path.splitext(original_filename)

        # Check if the file extension is supported
        if not any(extension.lower() in exts for exts in SUPPORTED_EXTENSIONS.values()):
            return f"Unsupported file type: {extension}", gr.update()

        # Create a unique filename to avoid overwriting
        base_name = os.path.splitext(original_filename)[0]
        counter = 1
        new_filename = f"{base_name}{extension}"
        while os.path.exists(os.path.join(SAVE_DIR, new_filename)):
            new_filename = f"{base_name}_{counter}{extension}"
            counter += 1

        # Copy the file from the temporary location to our save directory
        dest_path = os.path.join(SAVE_DIR, new_filename)
        shutil.copy2(file.name, dest_path)

        # Return success message and updated FileExplorer
        return f"File saved as {SAVE_DIR}/{new_filename}", gr.update(value=SAVE_DIR), gr.update(value=None)
    except Exception as e:
        return f"Error saving file: {str(e)}", gr.update(value=SAVE_DIR), gr.update()

def view_file(file_path):
    """View a file and return appropriate outputs for each media type.
    Returns: (image, audio, video, message, model3d)
    """
    if not file_path:
        return None, None, None, "No file selected.", None

    try:
        # FileExplorer returns relative path from root_dir, or could be full path
        # Handle both cases
        if os.path.isabs(file_path):
            full_path = file_path
        elif file_path.startswith(SAVE_DIR):
            full_path = file_path
        else:
            full_path = os.path.join(SAVE_DIR, file_path)

        # Normalize the path
        full_path = os.path.normpath(full_path)

        if not os.path.exists(full_path):
            return None, None, None, f"File not found: {full_path}", None

        _, extension = os.path.splitext(full_path)
        extension = extension.lower()

        if extension in SUPPORTED_EXTENSIONS["image"]:
            return Image.open(full_path), None, None, f"Viewing: {os.path.basename(full_path)}", None
        elif extension in SUPPORTED_EXTENSIONS["audio"]:
            return None, full_path, None, f"Viewing: {os.path.basename(full_path)}", None
        elif extension in SUPPORTED_EXTENSIONS["video"]:
            return None, None, full_path, f"Viewing: {os.path.basename(full_path)}", None
        elif extension in SUPPORTED_EXTENSIONS["model3d"]:
            return None, None, None, f"Viewing 3D: {os.path.basename(full_path)}", full_path
        else:
            return None, None, None, f"Unsupported file type: {extension}", None
    except Exception as e:
        return None, None, None, f"Error viewing file: {str(e)}", None


def get_all_media_files():
    """Get all media files from saved_media folder for use in dropdowns"""
    try:
        if not os.path.exists(SAVE_DIR):
            return []
        files = os.listdir(SAVE_DIR)
        # Filter to only supported media files
        media_files = []
        for f in files:
            ext = os.path.splitext(f)[1].lower()
            if any(ext in exts for exts in SUPPORTED_EXTENSIONS.values()):
                media_files.append(f)
        return sorted(media_files)
    except Exception as e:
        print(f"Error listing media files: {e}")
        return []

def refresh_file_explorer():
    files = os.listdir(SAVE_DIR)
    return gr.update(value=files)

def delete_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)
        fileexplorerupdate = refresh_file_explorer()
        return fileexplorerupdate, f"{file_path} has been deleted."
    else:
        return f"{file_path} does not exist."


def delete_file_and_refresh(filename):
    """Delete file and return updated dropdown choices"""
    if not filename:
        return gr.update(choices=get_all_media_files()), "No file selected"

    full_path = os.path.join(SAVE_DIR, filename)
    if os.path.exists(full_path):
        os.remove(full_path)
        return gr.update(choices=get_all_media_files(), value=None), f"Deleted: {filename}"
    return gr.update(choices=get_all_media_files()), f"File not found: {filename}"


def save_file_and_refresh(file):
    """Save file and return updated dropdown choices for Browse tab"""
    if file is None:
        return "No file uploaded.", gr.update(), gr.update(value=None)

    try:
        original_filename = os.path.basename(file.name)
        _, extension = os.path.splitext(original_filename)

        if not any(extension.lower() in exts for exts in SUPPORTED_EXTENSIONS.values()):
            return f"Unsupported file type: {extension}", gr.update(), gr.update()

        base_name = os.path.splitext(original_filename)[0]
        counter = 1
        new_filename = f"{base_name}{extension}"
        while os.path.exists(os.path.join(SAVE_DIR, new_filename)):
            new_filename = f"{base_name}_{counter}{extension}"
            counter += 1

        dest_path = os.path.join(SAVE_DIR, new_filename)
        shutil.copy2(file.name, dest_path)

        # Return updated dropdown choices
        return f"File saved: {new_filename}", gr.update(choices=get_all_media_files()), gr.update(value=None)
    except Exception as e:
        return f"Error saving file: {str(e)}", gr.update(), gr.update()

def upload_file_or_bundle(file):
    """
    Unified upload handler that handles both media files and zip bundles.
    Returns: (status, dropdown_update, file_input_clear, config_json, extracted_files)
    """
    if file is None:
        return "No file uploaded.", gr.update(), gr.update(value=None), None, ""

    try:
        original_filename = os.path.basename(file.name)
        _, extension = os.path.splitext(original_filename)
        ext_lower = extension.lower()

        # Check if it's a zip bundle
        if ext_lower == ".zip":
            try:
                config, extracted_files = import_config_with_media(file.name)
                config_json = json.dumps(config, indent=2)
                files_str = "\n".join(extracted_files) if extracted_files else "No new files extracted (all already existed)"
                status = f"Bundle imported! Config loaded, {len(extracted_files)} new media files extracted."
                return (
                    status,
                    gr.update(choices=get_all_media_files()),
                    gr.update(value=None),
                    config_json,
                    files_str
                )
            except FileNotFoundError as e:
                return f"Error: {str(e)}", gr.update(), gr.update(), None, ""
            except Exception as e:
                return f"Error importing bundle: {str(e)}", gr.update(), gr.update(), None, ""

        # Handle regular media files
        if not any(ext_lower in exts for exts in SUPPORTED_EXTENSIONS.values()):
            return f"Unsupported file type: {extension}", gr.update(), gr.update(), None, ""

        base_name = os.path.splitext(original_filename)[0]
        counter = 1
        new_filename = f"{base_name}{extension}"
        while os.path.exists(os.path.join(SAVE_DIR, new_filename)):
            new_filename = f"{base_name}_{counter}{extension}"
            counter += 1

        dest_path = os.path.join(SAVE_DIR, new_filename)
        shutil.copy2(file.name, dest_path)

        return (
            f"File saved: {new_filename}",
            gr.update(choices=get_all_media_files()),
            gr.update(value=None),
            None,
            ""
        )
    except Exception as e:
        return f"Error: {str(e)}", gr.update(), gr.update(), None, ""


def export_config_with_media_wrapper(config_json):
    """Wrapper for export that provides status message along with the file"""
    from my_text_game_engine_attempt import export_config_with_media
    if not config_json or not config_json.strip():
        return None, "Please paste a config JSON to export"
    try:
        zip_path = export_config_with_media(config_json)
        return zip_path, f"Export successful! Bundle ready for download."
    except json.JSONDecodeError as e:
        return None, f"Invalid JSON: {str(e)}"
    except Exception as e:
        return None, f"Export failed: {str(e)}"


def import_config_with_media(zip_path):
    global SAVE_DIR
    target_folder = SAVE_DIR
    """
    Import the config JSON and media files from a zip file.
    Extract only the media files that don't already exist in the target folder.
    
    :param zip_path: Path to the zip file containing the config and media files
    :param target_folder: Path to the folder where media files should be extracted
    :return: Tuple containing the loaded config (as a dictionary) and a list of newly extracted files
    """
    config = None
    extracted_files = []

    with tempfile.TemporaryDirectory() as temp_dir:
        # Extract all files to a temporary directory
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Load the config (check both config.json and game_config.json for compatibility)
        config_path = os.path.join(temp_dir, 'config.json')
        alt_config_path = os.path.join(temp_dir, 'game_config.json')

        if os.path.exists(config_path):
            with open(config_path, 'r') as f:
                config = json.load(f)
        elif os.path.exists(alt_config_path):
            with open(alt_config_path, 'r') as f:
                config = json.load(f)
        else:
            raise FileNotFoundError("config.json (or game_config.json) not found in the zip file")
        
        # Create the target folder if it doesn't exist
        os.makedirs(target_folder, exist_ok=True)
        
        # Copy media files that don't already exist in the target folder
        config_files = {'config.json', 'game_config.json'}
        for root, _, files in os.walk(temp_dir):
            for file in files:
                if file not in config_files:
                    src_path = os.path.join(root, file)
                    dst_path = os.path.join(target_folder, file)
                    if not os.path.exists(dst_path):
                        shutil.copy2(src_path, dst_path)
                        extracted_files.append(file)
    
    return config, extracted_files