import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import base64
import io
import random
import os
import uuid
import json # <-- Import the json library

# --- Configuration & Setup ---
# Create a directory to store temporary output files
TEMP_DIR = "temp_outputs"
os.makedirs(TEMP_DIR, exist_ok=True)


# --- Core Image Processing Logic (Unchanged) ---
# ... (All the functions like scramble_image, unscramble_image, etc., are the same)

def process_image_for_grid(image_pil, grid_size):
    if image_pil is None: return None, 0, 0
    img_width, img_height = image_pil.size
    tile_w, tile_h = img_width // grid_size, img_height // grid_size
    if tile_w == 0 or tile_h == 0:
        raise gr.Error(f"Image is too small for a {grid_size}x{grid_size} grid.")
    cropped_image = image_pil.crop((0, 0, tile_w * grid_size, tile_h * grid_size))
    return cropped_image, tile_w, tile_h

def scramble_image(image_pil, grid_size, seed):
    cropped_image, tile_w, tile_h = process_image_for_grid(image_pil, grid_size)
    tiles = [cropped_image.crop((j * tile_w, i * tile_h, (j + 1) * tile_w, (i + 1) * tile_h)) for i in range(grid_size) for j in range(grid_size)]
    rng = np.random.default_rng(seed=int(seed))
    scramble_map = rng.permutation(len(tiles))
    scrambled_image = Image.new('RGB', cropped_image.size)
    for new_pos_index, original_tile_index in enumerate(scramble_map):
        i, j = divmod(new_pos_index, grid_size)
        scrambled_image.paste(tiles[original_tile_index], (j * tile_w, i * tile_h))
    return scrambled_image, scramble_map

def unscramble_image(scrambled_pil, scramble_map, grid_size):
    cropped_image, tile_w, tile_h = process_image_for_grid(scrambled_pil, grid_size)
    scrambled_tiles = [cropped_image.crop((j * tile_w, i * tile_h, (j + 1) * tile_w, (i + 1) * tile_h)) for i in range(grid_size) for j in range(grid_size)]
    unscrambled_image = Image.new('RGB', cropped_image.size)
    for new_pos_index, original_pos_index in enumerate(scramble_map):
        i, j = divmod(original_pos_index, grid_size)
        unscrambled_image.paste(scrambled_tiles[new_pos_index], (j * tile_w, i * tile_h))
    return unscrambled_image

def create_mapping_visualization(scramble_map, grid_size):
    map_size=(512, 512)
    vis_image = Image.new('RGB', map_size, color='lightgray')
    draw = ImageDraw.Draw(vis_image)
    try: font = ImageFont.truetype("arial.ttf", size=max(10, 32 - grid_size * 2))
    except IOError: font = ImageFont.load_default()
    tile_w, tile_h = map_size[0] // grid_size, map_size[1] // grid_size
    for new_pos_index, original_pos_index in enumerate(scramble_map):
        i, j = divmod(new_pos_index, grid_size)
        x0, y0 = j * tile_w, i * tile_h
        draw.rectangle([x0, y0, x0 + tile_w, y0 + tile_h], outline='black')
        text = str(original_pos_index)
        text_bbox = draw.textbbox((0, 0), text, font=font)
        text_w, text_h = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
        draw.text((x0 + (tile_w - text_w) / 2, y0 + (tile_h - text_h) / 2), text, fill='black', font=font)
    return vis_image

def pil_to_base64(pil_image):
    buffered = io.BytesIO()
    pil_image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def create_canvas_html(base64_string, width, height):
    return f"""
    <div style="display: flex; justify-content: center; align-items: center; background: #f0f0f0;">
        <canvas id="unscrambled-canvas" width="{width}" height="{height}" style="max-width: 100%; max-height: 512px; object-fit: contain;"></canvas>
    </div>
    <script>
        function drawImageOnCanvas() {{
            const canvas = document.getElementById('unscrambled-canvas');
            if (!canvas) return;
            const ctx = canvas.getContext('2d');
            const img = new Image();
            img.src = "{base64_string}";
            img.onload = () => {{ ctx.clearRect(0, 0, canvas.width, canvas.height); ctx.drawImage(img, 0, 0, canvas.width, canvas.height); }};
            canvas.oncontextmenu = (e) => {{ e.preventDefault(); return false; }};
        }}
        setTimeout(drawImageOnCanvas, 100);
    </script>
    """

# --- Main Gradio Function ---

def process_and_display(input_image, grid_size, seed):
    """
    Main orchestrator function. Saves PNGs and JSON and returns all file paths.
    """
    if input_image is None:
        # Return empty placeholders for all outputs
        return None, "<div>Please upload an image to begin.</div>", None, None, None

    # 1. Scramble the image
    scrambled_img, scramble_map = scramble_image(input_image, grid_size, seed)
    
    # 2. Unscramble for canvas preview
    unscrambled_img = unscramble_image(scrambled_img, scramble_map, grid_size)
    base64_unscrambled = pil_to_base64(unscrambled_img)
    canvas_html = create_canvas_html(base64_unscrambled, unscrambled_img.width, unscrambled_img.height)

    # 3. Create map visualization
    map_viz_img = create_mapping_visualization(scramble_map, grid_size)

    # --- Save all files and get their paths ---
    unique_id = uuid.uuid4()
    
    # Save the scrambled image PNG
    scrambled_filepath = os.path.join(TEMP_DIR, f"{unique_id}_scrambled.png")
    scrambled_img.save(scrambled_filepath)
    
    # Save the map visualization PNG
    map_viz_filepath = os.path.join(TEMP_DIR, f"{unique_id}_map.png")
    map_viz_img.save(map_viz_filepath)

    # NEW: Create and save the JSON map file
    map_json_filepath = os.path.join(TEMP_DIR, f"{unique_id}_map.json")
    map_data = {
        "gridSize": grid_size,
        "seed": seed,
        "width": scrambled_img.width,
        "height": scrambled_img.height,
        "scrambleMap": scramble_map.tolist() # Convert numpy array for JSON
    }
    with open(map_json_filepath, 'w') as f:
        json.dump(map_data, f, indent=2)
    
    # 4. Return all necessary file paths and the canvas HTML
    return scrambled_filepath, canvas_html, map_viz_filepath, scrambled_filepath, map_json_filepath


# --- Gradio UI Definition ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🖼️ Secure Image Scrambler & Viewer (v3)
        Upload an image to create a scrambled `.png` file and a `.json` map file.
        The unscrambled original can be viewed in a protected preview window that prevents easy downloading.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1, min_width=350):
            input_image = gr.Image(
                type="pil", 
                label="Upload Image or Paste from Clipboard/URL",
                sources=["upload", "clipboard"]
            )
            
            with gr.Accordion("Settings", open=True):
                grid_size_slider = gr.Slider(minimum=2, maximum=32, value=8, step=1, label="Grid Size (NxN)")
                seed_input = gr.Number(value=lambda: random.randint(0, 99999), label="Scramble Seed")
            
            submit_btn = gr.Button("Scramble & Process", variant="primary")

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Downloads"):
                    gr.Markdown("### Scrambled Image")
                    scrambled_output = gr.Image(
                        label="Scrambled Image Preview",
                        type="filepath",
                        interactive=False,
                    )
                    downloadable_png = gr.File(
                        label="Download Scrambled PNG File",
                        interactive=False
                    )
                    gr.Markdown("---")
                    gr.Markdown("### Scrambling Map")
                    downloadable_json = gr.File(
                        label="Download Map JSON File",
                        interactive=False
                    )

                with gr.TabItem("Unscrambled Preview (Protected)"):
                    unscrambled_canvas = gr.HTML(
                        label="Unscrambled Preview (Not directly downloadable)"
                    )
                    
                with gr.TabItem("Map Visualization"):
                    mapping_output = gr.Image(
                        label="Mapping Key Visualization",
                        type="filepath",
                        interactive=False
                    )

    # Connect the button to the main function, mapping all outputs correctly
    submit_btn.click(
        fn=process_and_display, 
        inputs=[input_image, grid_size_slider, seed_input], 
        outputs=[
            scrambled_output,      # Receives scrambled_filepath for display
            unscrambled_canvas,    # Receives canvas_html
            mapping_output,        # Receives map_viz_filepath
            downloadable_png,      # Receives scrambled_filepath for download
            downloadable_json      # Receives map_json_filepath for download
        ]
    )

demo.launch()