import gradio as gr
from gradio_client import Client
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__)

# ==============================================================================
#  CLIENT-SIDE API CALL LOGIC (Using the Official Gradio Client)
# ==============================================================================

def call_api(server_space_id: str, image: Image.Image) -> dict:
    """
    Uses the official gradio_client to call the secure decoder API.
    """
    # --- Input Validation ---
    if not server_space_id or "/" not in server_space_id:
        raise gr.Error("Please provide the Server Space ID in the format 'username/space-name'.")
    if image is None:
        raise gr.Error("Please upload an encrypted image.")

    logger.info(f"Initializing client for server space: {server_space_id}")

    try:
        # 1. Initialize the Gradio Client
        # This points the client to your server space.
        client = Client(src=server_space_id)

        # 2. Convert the PIL Image to a base64 string
        with io.BytesIO() as buffer:
            image.save(buffer, format="PNG")
            image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')

        # 3. Call the specific API endpoint on the server
        # This matches the documentation you provided exactly.
        logger.info(f"Calling api_name='/keylock-auth-decoder' on {server_space_id}")
        result = client.predict(
            # The keyword argument MUST match the server function's parameter name.
            image_base64_string=image_base64,
            api_name="/keylock-auth-decoder"
        )
        
        logger.info("Successfully received result from API.")
        # The result from a successful predict call is the data itself.
        return result

    except Exception as e:
        # Handle other unexpected errors
        logger.error(f"An unexpected client-side error occurred: {e}", exc_info=True)
        raise gr.Error(f"An unexpected error occurred: {e}")

# ==============================================================================
#  GRADIO INTERFACE
# ==============================================================================

with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock API Client") as demo:
    gr.Markdown("# 🔑 KeyLock API Client")
    gr.Markdown(
        "This UI calls the **Secure Decoder API**. It sends an encrypted image to the server for decryption "
        "using the official `gradio_client` library."
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Configure Server Space ID")
            server_id_input = gr.Textbox(
                label="Decoder Server Space ID",
                value="broadfield-dev/KeyLock-Auth",
                info="This is the Hugging Face ID of your server, e.g., 'username/space-name'."
            )

            gr.Markdown("### 2. Upload Encrypted Image")
            image_input = gr.Image(type="pil", label="Image encrypted with the server's public key", sources=["upload", "clipboard"])
            submit_button = gr.Button("Decrypt via API", variant="primary")

        with gr.Column(scale=2):
            gr.Markdown("### 3. Decrypted Data")
            json_output = gr.JSON(label="Result from Server")

    submit_button.click(
        fn=call_api,
        inputs=[server_id_input, image_input],
        outputs=[json_output]
    )

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