app.py

import gradio as gr
from gradio_client import Client, GradioClientError
from PIL import Image
import base64
import io
import logging

# --- Configure Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ==============================================================================
#  CONFIGURATION: IDs of the remote Gradio Spaces
# ==============================================================================
# These should be the names of YOUR deployed spaces.
CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
# We also provide a link to the original client for reference.
CLIENT_SPACE_LINK = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Client"


# ==============================================================================
#  API CALL WRAPPER FUNCTIONS
# ==============================================================================

def create_image_via_api(secret_data: str, public_key: str):
    """
    Calls the Creator Space API to generate an encrypted image.
    This function uses 'yield' to provide real-time status updates to the UI.
    """
    if not all([secret_data, public_key]):
        raise gr.Error("Secret Data and Public Key are both required.")
    
    status = f"Initializing client for Creator: {CREATOR_SPACE_ID}..."
    yield None, None, status  # Yield initial status update

    try:
        client = Client(src=CREATOR_SPACE_ID)
        status = f"Calling API '/create_image' on {CREATOR_SPACE_ID}..."
        yield None, None, status
        
        # The 'predict' method calls the API endpoint.
        # The result from an Image output in the remote API is a filepath to a temporary file.
        temp_filepath = client.predict(
            secret_data_str=secret_data,
            public_key_pem=public_key,
            api_name="/create_image"  # The API name defined in the Creator space
        )
        
        if not temp_filepath:
            raise gr.Error("Creator API did not return a valid image file path.")
            
        # Load the image from the temporary file to display it in our dashboard's UI
        created_image = Image.open(temp_filepath)
        
        status = "✅ Success! Image created by the Creator service."
        yield created_image, temp_filepath, status

    except Exception as e:
        logger.error(f"Creator API call failed: {e}", exc_info=True)
        # Yield the error message back to the UI
        yield None, None, f"❌ Error calling Creator API: {e}"


def decrypt_image_via_api(image: Image.Image):
    """
    Calls the Server Space API to decrypt an image.
    Uses 'yield' for status updates.
    """
    if image is None:
        raise gr.Error("Please upload an image to decrypt.")

    status = f"Initializing client for Server: {SERVER_SPACE_ID}..."
    yield None, status

    try:
        client = Client(src=SERVER_SPACE_ID)
        
        # Convert the user's uploaded PIL image to a base64 string for the server's API
        with io.BytesIO() as buffer:
            image.save(buffer, format="PNG")
            b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
        
        status = f"Calling API '/keylock-auth-decoder' on {SERVER_SPACE_ID}..."
        yield None, status
        
        # Call the server's API endpoint. The result will be a dictionary.
        decrypted_json = client.predict(
            image_base64_string=b64_string,
            api_name="/keylock-auth-decoder"  # The API name defined in the Server space
        )
        
        status = "✅ Success! Data decrypted by the Server."
        yield decrypted_json, status

    except Exception as e:
        logger.error(f"Server API call failed: {e}", exc_info=True)
        yield None, f"❌ Error calling Server API: {e}"


# ==============================================================================
#  GRADIO DASHBOARD INTERFACE
# ==============================================================================
theme = gr.themes.Base(
    primary_hue="blue",
    secondary_hue="sky",
    neutral_hue="slate",
    font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
).set(
    body_background_fill="#F8FAFC",
    block_background_fill="white",
    block_border_width="1px",
    block_shadow="*shadow_drop_lg",
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    button_primary_text_color="white",
)

with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
    gr.Markdown("# 🔑 KeyLock Operations Dashboard")
    gr.Markdown("A centralized dashboard demonstrating the entire KeyLock ecosystem by making live API calls to the deployed Creator and Server spaces.")

    with gr.Tabs():
        with gr.TabItem("🏭 Image Creator", id=0):
            gr.Markdown("## Create an Encrypted Image via API Call")
            gr.Markdown(f"This interface calls the **`{CREATOR_SPACE_ID}`** Space to generate a new encrypted image.")
            with gr.Row(variant="panel"):
                with gr.Column(scale=2):
                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data", placeholder="API_KEY: sk-123...\nUSER: demo-user")
                    creator_pubkey_input = gr.Textbox(lines=7, label="Public Key of the Target Server", placeholder="Paste the Server's public key here...")
                    creator_button = gr.Button("Create Image via Creator API", variant="primary", icon="✨")
                with gr.Column(scale=1):
                    creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=True)
                    creator_download_output = gr.File(label="Download Image File")

        with gr.TabItem("💻 Client / Decoder", id=1):
            gr.Markdown("## Decrypt an Image via API Call")
            gr.Markdown(f"This interface acts as a client, calling the **`{SERVER_SPACE_ID}`** Space to decrypt an image.")
            with gr.Row(variant="panel"):
                with gr.Column(scale=1):
                    client_image_input = gr.Image(type="pil", label="Upload Encrypted Image", sources=["upload", "clipboard"])
                    client_button = gr.Button("Decrypt Image via Server API", variant="primary", icon="🔓")
                with gr.Column(scale=1):
                    client_status = gr.Textbox(label="Status", interactive=False, lines=2)
                    client_json_output = gr.JSON(label="Decrypted Data from Server")
        
        with gr.TabItem("ℹ️ Service Information", id=2):
            gr.Markdown("## Ecosystem Overview")
            gr.Markdown(
                f"""
                This dashboard coordinates three separate Hugging Face Spaces to demonstrate a complete workflow:

                1.  **Creator Service:**
                    -   **Space:** [{CREATOR_SPACE_ID}](https://huggingface.co/spaces/{CREATOR_SPACE_ID})
                    -   **Role:** Provides a UI and an API (`/run/create_image`) to encrypt data into PNG images.

                2.  **Server (Decoder) Service:**
                    -   **Space:** [{SERVER_SPACE_ID}](https://huggingface.co/spaces/{SERVER_SPACE_ID})
                    -   **Role:** Holds a secret private key and provides a secure API (`/run/keylock-auth-decoder`) to decrypt images.

                3.  **This Dashboard:**
                    -   **Role:** Acts as a master client and control panel, making live API calls to the other services to showcase the end-to-end process.
                
                *Note: The original `{CLIENT_SPACE_LINK.split('/')[-1]}` is now superseded by the 'Client / Decoder' tab in this dashboard.*
                """
            )

    # --- Wire up the component logic ---
    creator_button.click(
        fn=create_image_via_api,
        inputs=[creator_secret_input, creator_pubkey_input],
        outputs=[creator_image_output, creator_download_output, creator_status]
    )
    
    client_button.click(
        fn=decrypt_image_via_api,
        inputs=[client_image_input],
        outputs=[client_json_output, client_status]
    )

if __name__ == "__main__":
    demo.launch()