import json
import gradio as gr

def generate_files(title="Text Generation Tool", emoji="🌖", colorFrom="blue", colorTo="blue",
                   sdk="gradio", sdk_version="4.3.0", app_file="app.py", pinned=False,
                   tags=["tool"], tool_name="text_generator", tool_description="This is a tool that chats with a user. "
                   "It takes an input named `prompt` which contains a system_role, user_message, context and history. "
                   "It returns a text message."):
    # Generate readme content
    readme_content = f'''## readme
title: {title}
emoji: {emoji}
colorFrom: {colorFrom}
colorTo: {colorTo}
sdk: {sdk}
sdk_version: {sdk_version}
app_file: {app_file}
pinned: {pinned}
tags:
  - {tags[0]}
'''

    # Generate tool config JSON content
    tool_config = {
        "description": tool_description,
        "name": tool_name,
        "tool_class": f"{tool_name.capitalize()}Tool"
    }
    tool_config_json = json.dumps(tool_config, indent=4)

    # Generate app.py content
    app_py_content = f'''from transformers.tools.base import launch_gradio_demo
                    from {tool_name} import {tool_name.capitalize()}Tool
                    
                    launch_gradio_demo({tool_name.capitalize()}Tool)
                    '''

    # Generate requirements.txt content
    requirements_content = '''transformers>=4.29.0
            # diffusers
            accelerate
            torch
            '''

    # Generate text_generator.py content
    text_generator_py_content = f'''import os
    from transformers import pipeline
    from transformers import Tool
    
    class {tool_name.capitalize()}Tool(Tool):
        name = "{tool_name}"
        description = (
            "{tool_description}"
        )
    
        inputs = ["text"]
        outputs = ["text"]
    
        def __call__(self, prompt: str):
            token = os.environ['hf']
            text_generator = pipeline(model="microsoft/Orca-2-13b", token=token)
            generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
            print(generated_text)
            return generated_text
    '''

    # Write content to files
    with open("README.md", "w") as readme_file:
        readme_file.write(readme_content)

    with open("tool_config.json", "w") as tool_config_file:
        tool_config_file.write(tool_config_json)

    with open("app.py", "w") as app_py_file:
        app_py_file.write(app_py_content)

    with open("requirements.txt", "w") as requirements_file:
        requirements_file.write(requirements_content)

    with open(f"{tool_name}.py", "w") as text_generator_py_file:
        text_generator_py_file.write(text_generator_py_content)

    # Return the generated files for download
    return "README.md", "tool_config.json", "app.py", "requirements.txt", f"{tool_name}.py"


# Define the inputs for the Gradio interface
io = gr.Interface(generate_files, 
                  inputs=["text", "text", "text", "text", "text", "text", "text", "text", "checkbox", "text", "text"],
                  outputs=["text", "text", "text", "text", "text"])

# Launch the Gradio interface
io.launch()