app.py

import gradio as gr
import json
import logging
import tempfile
import core  # Import our core module

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

# ==============================================================================
#  CONFIGURATION & DATA HANDLING
# ==============================================================================

DEFAULT_CONFIG_FILE = "./endpoints.json"

def load_config_from_file(filepath: str) -> list:
    """Loads and validates the JSON configuration from a file."""
    try:
        with open(filepath, "r") as f:
            data = json.load(f)
        if not isinstance(data, list):
            raise TypeError("JSON root must be a list of service objects.")
        for service in data:
            if not all(k in service for k in ("name", "link", "public_key")):
                raise ValueError("Each service object must contain 'name', 'link', and 'public_key' keys.")
        logger.info(f"Successfully loaded {len(data)} services from {filepath}.")
        return data
    except FileNotFoundError:
        logger.warning(f"{filepath} not found. Returning empty list.")
        return []
    except Exception as e:
        logger.error(f"Error loading config file {filepath}: {e}")
        return []

def save_config_to_file(filepath: str, config_data: list) -> str:
    """Saves the current configuration data to a JSON file."""
    try:
        with open(filepath, "w") as f:
            json.dump(config_data, f, indent=2)
        status_msg = f"✅ Success! Configuration saved to {filepath}."
        logger.info(status_msg)
        return status_msg
    except Exception as e:
        status_msg = f"❌ Error saving configuration: {e}"
        logger.error(status_msg)
        return status_msg

initial_config = load_config_from_file(DEFAULT_CONFIG_FILE)

# ==============================================================================
#  GRADIO UI HELPER FUNCTIONS
# ==============================================================================

def add_new_service(current_config: list, name: str, link: str, public_key: str):
    """Adds a new service to the current configuration state."""
    if not all([name, link, public_key]):
        raise gr.Error("All fields (Name, Link, Public Key) are required to add a new service.")
    
    new_service = {"name": name, "link": link, "public_key": public_key}
    
    if any(service['name'] == name for service in current_config):
        raise gr.Error(f"A service with the name '{name}' already exists. Please use a unique name.")
        
    updated_config = current_config + [new_service]
    updated_choices = [service['name'] for service in updated_config]
    
    status_update = f"✅ '{name}' added. You can now use it in the 'Create Image' tab. Click 'Save to File' to make it permanent."
    return updated_config, gr.Dropdown(choices=updated_choices), status_update, "", "", ""

def generate_image(selected_service_name: str, secret_data_str: str, current_config: list):
    """The main image creation function for the UI."""
    if not all([selected_service_name, secret_data_str]):
        raise gr.Error("A selected service and secret data are both required.")

    try:
        public_key = next((s['public_key'] for s in current_config if s.get('name') == selected_service_name), None)
        if not public_key:
            raise gr.Error(f"Could not find the service '{selected_service_name}'. Please check the configuration.")

        # This call is correct and should not cause an error
        encrypted_image = core.create_encrypted_image(secret_data_str, public_key)
        
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
            encrypted_image.save(tmp_file.name)
            logger.info(f"Generated image saved for '{selected_service_name}'.")
            return encrypted_image, tmp_file.name, f"✅ Success! Image created for '{selected_service_name}'."

    except Exception as e:
        logger.error(f"Image creation failed: {e}", exc_info=True)
        return None, None, f"❌ Error: {e}"

def get_endpoints_json() -> str:
    """Reads and returns the content of endpoints.json as a string."""
    try:
        with open(DEFAULT_CONFIG_FILE, "r") as f:
            return f.read()
    except Exception as e:
        logger.error(f"Could not read {DEFAULT_CONFIG_FILE}: {e}")
        raise gr.Error(f"Server could not read its {DEFAULT_CONFIG_FILE} configuration file.")

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

with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock Image Creator") as demo:
    config_state = gr.State(initial_config)
    
    gr.Markdown("# 🏭 KeyLock Image Creator")
    gr.Markdown("Create secure, encrypted PNG images for various decoder services.")

    with gr.Tabs() as tabs:
        with gr.TabItem("Create Image", id=0):
            # ... (Your existing UI code is fine, no changes needed here) ...
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("### 1. Select Decoder Service")
                    service_dropdown = gr.Dropdown(
                        label="Target Service",
                        choices=[s['name'] for s in initial_config],
                        value=[s['name'] for s in initial_config][0] if initial_config else None,
                        interactive=True
                    )
                    gr.Markdown("### 2. Enter Your Secret Data")
                    secret_data = gr.Textbox(
                        lines=7,
                        label="Secret Data (Key-Value Pairs)",
                        placeholder="USERNAME: user@example.com\nAPI_KEY: sk-12345..."
                    )
                with gr.Column(scale=1):
                    gr.Markdown("### 3. Generate")
                    create_button = gr.Button("Create Encrypted Image", variant="primary")
                    status_output = gr.Textbox(label="Status", interactive=False, lines=2)
                    image_output = gr.Image(label="Generated Image", type="pil")
                    download_output = gr.File(label="Download Image")

        with gr.TabItem("Manage Configuration", id=1):
            # ... (Your existing config management code is fine, no changes needed here) ...
            gr.Markdown("## Decoder Service Configuration")
            gr.Markdown("View the current configuration, add new services, and save your changes.")
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("### Current Configuration")
                    config_display = gr.JSON(value=initial_config, label=f"Loaded from {DEFAULT_CONFIG_FILE}")
                with gr.Column(scale=1):
                    gr.Markdown("### Add a New Service")
                    new_service_name = gr.Textbox(label="Service Name", placeholder="e.g., My Production API")
                    new_service_link = gr.Textbox(label="Service Link", placeholder="https://huggingface.co/...")
                    new_service_key = gr.Textbox(label="Public Key (PEM Format)", lines=5, placeholder="-----BEGIN PUBLIC KEY-----...")
                    add_service_button = gr.Button("Add Service to Current Session")
            gr.Markdown("---")
            save_config_button = gr.Button("Save Current Configuration to File", variant="secondary", visible=False)
            config_status_output = gr.Textbox(label="Configuration Status", interactive=False)
            
        # FIX: Define API endpoints clearly instead of using a hidden row
        with gr.TabItem("API", id=2):
            gr.Markdown("## API Endpoints")
            gr.Markdown("Use these endpoints for programmatic access.")
            
            # API for creating an image
            gr.Interface(
                # By passing core.create_encrypted_image directly, we ensure it runs in the correct scope
                fn=core.create_encrypted_image,
                inputs=[
                    gr.Textbox(label="Secret Data", info="Key-value pairs, one per line."),
                    gr.Textbox(label="Public Key PEM", info="The full PEM-formatted public key.")
                ],
                outputs=gr.Image(type="pil", label="Encrypted Image"),
                title="Image Creation API",
                description="Creates an encrypted image from secret data and a public key. Use via the `/api/create_image/` route.",
                api_name="create_image" 
            )

            # API to get the endpoint list
            gr.Interface(
                fn=get_endpoints_json,
                inputs=[],
                outputs=gr.Textbox(label="Endpoint Configuration"),
                title="Get Endpoints API",
                description=f"Retrieves the contents of `{DEFAULT_CONFIG_FILE}`. Use via the `/api/get_endpoints/` route.",
                api_name="get_endpoints"
            )

    # --- Interactivity ---
    
    create_button.click(
        fn=generate_image,
        inputs=[service_dropdown, secret_data, config_state],
        outputs=[image_output, download_output, status_output]
    )

    add_service_button.click(
        fn=add_new_service,
        inputs=[config_state, new_service_name, new_service_link, new_service_key],
        outputs=[config_state, service_dropdown, config_status_output, new_service_name, new_service_link, new_service_key]
    ).then(
        lambda state: state,
        inputs=[config_state],
        outputs=[config_display]
    )
    
    save_config_button.click(
        fn=save_config_to_file,
        inputs=[gr.State(DEFAULT_CONFIG_FILE), config_state],
        outputs=[config_status_output]
    )

    config_state.change(
        fn=lambda state: state,
        inputs=[config_state],
        outputs=[config_display]
    )

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