import gradio as gr
import core # Import all our local logic

# ==============================================================================
#  GRADIO UI HELPER FUNCTIONS
# ==============================================================================
def creator_wrapper(public_key: str, secret_data: str):
    """UI wrapper for the image creation logic."""
    try:
        created_image = core.create_encrypted_image(secret_data, public_key)
        return created_image, f"✅ Success! Image created. You can now go to the 'Send KeyLock' tab to test it."
    except Exception as e:
        return None, f"❌ Error: {e}"

def decoder_wrapper(image: gr.Image, private_key: str):
    """UI wrapper for the image decryption logic."""
    try:
        decrypted_data = core.decode_data_from_image(image, private_key)
        return decrypted_data, "✅ Success! Data decrypted from image."
    except Exception as e:
        return None, f"❌ Error: {e}"

# ==============================================================================
#  GRADIO DASHBOARD INTERFACE
# ==============================================================================
theme = gr.themes.Base(
    primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
    font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
).set(
    body_background_fill="#F1F5F9", panel_background_fill="white", 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",
)

with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
    gr.Markdown("# 🔑 KeyLock Operations Dashboard")
    gr.Markdown("A self-contained demonstration of the entire KeyLock workflow: Key Generation, Image Creation, and Decryption.")

    with gr.Tabs() as tabs:
        with gr.TabItem("① Generate Keys", id=0):
            gr.Markdown("## Step 1: Create a Secure Key Pair")
            gr.Markdown(
                """
                Every secure system starts with a key pair. This consists of a **Public Key** and a **Private Key**.
                - **Public Key 🔑:** You can share this with anyone. It's used only for *encrypting* data. In a real system, you would add this key to a server's configuration file (like an `endpoints.json`) so that clients know how to encrypt data for it.
                - **Private Key 🔐:** This must be kept **absolutely secret**. It is the only key that can *decrypt* data encrypted with its corresponding public key. In a real system, this would be stored as a secure environment variable or secret on the server.
                """
            )
            with gr.Row(variant="panel"):
                with gr.Column(scale=1):
                    gr.Markdown("### Your New Keys")
                    gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", variant="secondary")
                    gr.Markdown("⬇️ **Copy these keys for the next steps!**")
                with gr.Column(scale=2):
                     with gr.Row():
                        output_public_key = gr.Textbox(lines=11, label="Generated Public Key (For Creator)", interactive=False, show_copy_button=True)
                        output_private_key = gr.Textbox(lines=11, label="Generated Private Key (For Decoder)", interactive=False, show_copy_button=True)

        with gr.TabItem("② Create KeyLock", id=1):
            gr.Markdown("## Step 2: Create an Encrypted Auth Image")
            gr.Markdown("This tool acts as the **Auth Creator**. It takes your secret data and uses the **Public Key** from Step 1 to encrypt it into a new PNG image. This simulates a user preparing their credentials to send to a secure service.")
            with gr.Row(variant="panel"):
                with gr.Column(scale=1):
                    gr.Markdown("### Configuration")
                    creator_pubkey_input = gr.Textbox(lines=8, label="Paste the Public Key Here", placeholder="Copy the public key generated in Step 1...")
                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user")
                    creator_button = gr.Button("✨ Create Auth Image", variant="primary")
                with gr.Column(scale=1):
                    gr.Markdown("### Output")
                    creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
                    # format="png" is the key to ensuring the download button creates a proper PNG file.
                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png", show_share_button=False)

        with gr.TabItem("③ Send KeyLock", id=2):
            gr.Markdown("## Step 3: Decrypt the Image (Client Simulation)")
            gr.Markdown("This tool acts as the **Client** sending the image to a secure **Server**. To prove it can decrypt the data, the 'Server' needs the corresponding **Private Key** from Step 1. In a real application, you would only send the image; the server would already have its private key.")
            with gr.Row(variant="panel"):
                with gr.Column(scale=1):
                    gr.Markdown("### Input")
                    client_image_input = gr.Image(type="pil", label="Upload or Drag Encrypted Image Here", sources=["upload", "clipboard"])
                    client_private_key_input = gr.Textbox(lines=8, label="Paste the Private Key Here", placeholder="Copy the private key generated in Step 1...")
                    client_button = gr.Button("🔓 Decrypt Image", variant="primary")
                with gr.Column(scale=1):
                    gr.Markdown("### Decrypted Data")
                    client_status = gr.Textbox(label="Status", interactive=False, lines=2)
                    client_json_output = gr.JSON(label="Result from 'Server'")

    # --- Wire up the component logic ---
    gen_keys_button.click(fn=core.generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
    
    creator_button.click(
        fn=creator_wrapper,
        inputs=[creator_pubkey_input, creator_secret_input],
        outputs=[creator_image_output, creator_status]
    )
    
    client_button.click(
        fn=decoder_wrapper,
        inputs=[client_image_input, client_private_key_input],
        outputs=[client_json_output, client_status]
    )

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