# filename: request_router.py

import gradio as gr
import os
import importlib.util
import json
import requests
from typing import Any, Dict

# Define the path to the tools folder
TOOLS_DIR = './tools'
os.makedirs(TOOLS_DIR, exist_ok=True)

# This is the only hardcoded Tool endpoint
def register_tool(tool_data):
    # Send a POST request to register the tool
    response = requests.post("https://huggingface.co/chat/tools/new", json=tool_data)
    return response.json()

def create_tool(tool_data: Dict[str, Any]):
    """Create a new tool script in the tools folder.
    
    Args:
        tool_data: Dictionary containing registration payload values and tool function script.
    """
    print("Processing create_tool")
    
    # Prepare the payload for registration
    tool_payload = {
        "displayName": tool_data.get("displayName"),
        "description": tool_data.get("description"),
        "color": tool_data.get("color"),
        "icon": tool_data.get("icon"),
        "baseUrl": "K00B404/toolshed",
        "endpoint": "/router",
        "name": tool_data.get("name"),
        "inputs": tool_data.get("inputs"),
        "outputComponent": tool_data.get("outputComponent"),
        "outputComponentIdx": tool_data.get("outputComponentIdx"),
        "showOutput": tool_data.get("showOutput"),
    }
    
    # Write the new tool script from the 'tfn' (tool function name) data
    with open(f'{TOOLS_DIR}/{tool_data.get("name").lower()}.py', 'w') as f:
        f.write(tool_data.get("tfn"))
    
    # Register the new tool
    return register_tool(tool_payload)

# Function to dynamically load methods from files in the ./tools folder
def load_methods_from_tools():
    method_mapping = {}

    # Manually map the create_tool function
    method_mapping["create_tool"] = create_tool

    # List all .py files in the tools directory
    for filename in os.listdir(TOOLS_DIR):
        if filename.endswith('.py'):
            method_name = filename[:-3].lower()  # Remove '.py' and convert to lowercase
            
            # Load the module
            module_path = os.path.join(TOOLS_DIR, filename)
            spec = importlib.util.spec_from_file_location(method_name, module_path)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)

            # Map all functions from the module to the method mapping
            for attr in dir(module):
                if callable(getattr(module, attr)) and not attr.startswith("__"):
                    method_mapping[method_name] = getattr(module, attr)

    return method_mapping

# Load the method mapping at the start
method_mapping = load_methods_from_tools()

# Main request router method
def request_routerOld(name: str, input_data: Dict[str, Any]) -> Dict[str, Any]:
    # Reload methods to include newly created tools
    global method_mapping
    method_mapping = load_methods_from_tools()

    # Retrieve the appropriate method based on the name
    method = method_mapping.get(name)
    if method is None:
        return {"error": "Method not found"}
    
    # Call the method with the provided input data
    try:
        output = method(input_data)
        return output
    except Exception as e:
        return {"error": str(e)}
        
def request_router(name: str, input_data: Dict[str, Any]) -> Dict[str, Any]:
    global method_mapping
    method_mapping = load_methods_from_tools()

    logging.info(f"Request to invoke method: {name} with input: {input_data}")
    method = method_mapping.get(name)
    if method is None:
        error_message = "Method not found"
        logging.error(error_message)
        return {"error": error_message}
    
    try:
        output = method(input_data)
        logging.info(f"Output: {output}")
        return output
    except Exception as e:
        logging.error(f"Error occurred: {str(e)}")
        return {"error": str(e)}
        
# Gradio Interface
def launch_gradio_app():
    with gr.Blocks() as demo:
        gr.Markdown("# Request Router API")
        with gr.Row():
            name_input = gr.Textbox(label="Method Name")
            input_data_input = gr.Textbox(label="Input Data (JSON format)")
            output_output = gr.JSON(label="Output")
        
        submit_button = gr.Button("Submit")
        submit_button.click(
            fn=lambda name, input_data: request_router(name, json.loads(input_data)),
            inputs=[name_input, input_data_input],
            outputs=output_output
        )

    demo.launch()

import logging

# Set up logging configuration
logging.basicConfig(level=logging.INFO, filename='request_logs.txt', 
                    format='%(asctime)s - %(message)s')

def log_payload(*args, **kwargs):
    # Log the incoming payload
    logging.info(f"Received payload: args={args}, kwargs={kwargs}")
    return main_function(*args, **kwargs)

def main_function(name):
    return "Hello " + name + "!!"




if __name__ == "__main__":
    launch_gradio_app()