import inspect
import json
import ast
import gradio as gr


def function_to_json(func_str, func_description, param_descriptions, required_params):
    # Create a new Module instance with the missing field
    module_ast = ast.Module(body=[ast.Pass()], type_ignores=[])

    # Parse the function string into the AST and replace the body
    func_ast = ast.parse(func_str)
    module_ast.body = func_ast.body

    # Extract the function definition node
    func_def = next(node for node in module_ast.body if isinstance(node, ast.FunctionDef))

    # Get function signature
    code_obj = compile(module_ast, '<string>', 'exec')
    func_globals = {}
    exec(code_obj, func_globals)
    signature = inspect.signature(func_globals[func_def.name])
    parameters = signature.parameters

    # Convert param_descriptions string to a dictionary
    param_desc_dict = json.loads(param_descriptions)

    # Create JSON structure
    function_json = {
        "name": func_def.name,
        "description": func_description,
        "parameters": {
            "type": "object",
            "properties": {}
        }
    }

    # Add parameter information to JSON structure
    for param_name, param in parameters.items():
        param_info = param_desc_dict.get(param_name, {})
        param_type = param_info.get("type", str(param.annotation))
        param_desc = param_info.get("description", param_name.replace('_', ' '))

        function_json["parameters"]["properties"][param_name] = {
            "type": param_type,
            "description": param_desc
        }

        # Add required parameters based on user input
        if param_name in required_params:
            if "required" not in function_json["parameters"]:
                function_json["parameters"]["required"] = []
            function_json["parameters"]["required"].append(param_name)

    return json.dumps(function_json, indent=4)
    

""" Example uasge:
# Example usage with user-provided function information
sample_function_str = '''
def generate_music(input_text, input_melody):
    ''' generate music based on an input text '''
    client = Client("https://ysharma-musicgendupe.hf.space/", hf_token="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")
    result = client.predict(
        "melody",
        input_text,
        input_melody,
        5,
        250,
        0,
        1,
        3,
        fn_index=1
    )
    return result
'''

sample_func_description = "generate music based on an input text and input melody"

sample_param_descriptions = '''
{
    "input_text": {
        "type": "str",
        "description": "Input text for music generation."
    },
    "input_melody": {
        "type": "str",
        "description": "File path of the input melody."
    }
}
'''

sample_required_params = ["input_text"]

# Convert the sample function information to JSON
json_str = function_to_json(sample_function_str, sample_func_description, sample_param_descriptions, sample_required_params)
print(json_str)

{
    "name": "generate_music",
    "description": "generate music based on an input text and input melody",
    "parameters": {
        "type": "object",
        "properties": {
            "input_text": {
                "type": "str",
                "description": "Input text for music generation."
            },
            "input_melody": {
                "type": "str",
                "description": "File path of the input melody."
            }
        },
        "required": [
            "input_text"
        ]
    }
}

"""


title = "<h1 align='center'>Convert any function to function definitions required for GPT</h1>"
demo = gr.Blocks()

with demo:
  gr.HTML(title)
  with gr.Row():
    input_function_str = gr.Code(label="Enter function definition", language='python', lines=10)
    #input_function_str = gr.Textbox(lines=10, label='Enter function definition')
    with gr.Column():
      input_func_description = gr.Textbox(placeholder='', label='Enter your function description:')
      input_param_description = gr.Textbox(
                                placeholder="""Enter description as a dictionary with keys as param_name and values as param type and description as shown, eg. -
                                {
                                  "param1": {
                                          "type": "str",
                                          "description": "description of param1"
                                      },
                                  "param2": {
                                          "type": "int/float/list/tuple/dict/set/bool etc..",
                                          "description": "description of param2"
                                      }
                                  }""",
                                label='Enter descriptions for parameters:')
      input_required_params = gr.Textbox(placeholder="""Enter a list of required parameters, eg. - ['param1', 'param2', ...]""",
                                         label='Enter required parameters for your function:')
  generate_json = gr.Button('Get JSON definition')
  gpt_function = gr.Code(label="GPT function definition", language='python', lines=7)

  generate_json.click(function_to_json,
                    [input_function_str, input_func_description, input_param_description, input_required_params],
                    [gpt_function])

  gr.Examples(
      [ ["""
        def generate_music(input_text, input_melody):
            "generate music based on an input text"
            client = Client("https://ysharma-musicgendupe.hf.space/", hf_token="hf_...")
            result = client.predict(
                "melody",
                input_text,
                input_melody,
                5,
                250,
                0,
                1,
                3,
                fn_index=1
            )
            return result
      """,
      """Generate music based on an input text.""",
      """{
            "input_text": {
                "type": "string",
                "description": "Input text for music generation."
            },
            "input_melody": {
                "type": "string",
                "description": "File path of the input melody."
            }
        }""",
      """["input_text"]""" ],

       ["""
       def generate_image(prompt):
            client = Client("https://jingyechen22-textdiffuser.hf.space/")
            result = client.predict(
                    prompt,
                    20,
                    7.5,
                    1,
                    "Stable Diffusion v2.1",
                    fn_index=1)
            return result[0]
        """,
        """generate image based on the input text prompt""",
        """{
              "prompt": {
                  "type": "string",
                  "description": "input text prompt for the image generation."
              }
          }""",
        """["prompt"]""" ,
         ],
       ],
        [input_function_str, input_func_description, input_param_description, input_required_params],
    )
  demo.launch() #(debug=True)