import gradio as gr
import json

def char(char_name, char_age, char_features, char_personality):
    char_info = {
        "name": char_name,
        "age": char_age,
        "features": char_features,
        "personality": char_personality
    }
    return char_info

def user(user_name, user_age, user_features, user_related, char_name):
    user_info = {
        "name": user_name,
        "age": user_age,
        "features": user_features,
        "relation": f"{user_related} of {char_name}"
    }
    return user_info

def combine_info(user_info, char_info, scenario):
    combined_info = f"""You are an expert actor, capable of fully immersing yourself in any role.
Once you assume a role, you do not break character under any circumstances, even if addressed as an AI or language model.
Your current role is {char_info['name']}, described in detail below.
Engage in the exchange as {char_info['name']}, treating the other participant as a {user_info['relation']} in this role-play scenario.

Role-play scenario:\n{scenario}

User Info:
Name: {user_info['name']}
Age: {user_info['age']}
Features: {user_info['features']}

Character Info:
Name: {char_info['name']}
Age: {char_info['age']}
Features: {char_info['features']}
Personality: {char_info['personality']}"""

    return combined_info

def save_character(char_info, user_info, combined_info):
    save_data = {
        "user": user_info,
        "character": char_info,
        "combined_info": combined_info
    }
    with open("saved_character.json", "w") as outfile:
        json.dump(save_data, outfile)
    return "Character and User information saved successfully!"

def load_character():
    try:
        with open("saved_character.json", "r") as infile:
            save_data = json.load(infile)
        user_info = save_data["user"]
        char_info = save_data["character"]
        combined_info = save_data["combined_info"]
        return (
            user_info["name"], user_info["age"], user_info["features"], user_info["relation"].split(" ")[0],
            char_info["name"], char_info["age"], char_info["features"], char_info["personality"],
            combined_info, "Data loaded successfully!"
        )
    except FileNotFoundError:
        return "", 0, "", "", "", 0, "", "", "", "No saved data found."

# Advanced UI with features to dynamically add traits and load/save capabilities
with gr.Blocks() as demo:
    saved_character_state = gr.State()
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### User Information")
            user_name = gr.Textbox(label="User Name", placeholder="e.g. John")
            user_age = gr.Number(label="User Age", placeholder="e.g. 30")
            user_features = gr.Textbox(label="User Features", placeholder="e.g. tall, brown hair, blue eyes")
            user_related = gr.Textbox(label="Relation to Character", placeholder="e.g. friend, sibling, colleague")
            user_output = gr.Textbox(label="User Information", interactive=False)
        with gr.Column():
            gr.Markdown("### Character Information")
            char_name = gr.Textbox(label="Character Name", placeholder="e.g. Emily")
            char_age = gr.Number(label="Character Age", placeholder="e.g. 25")
            char_features = gr.Textbox(label="Character Features", placeholder="e.g. short, blonde hair, green eyes")
            char_personality = gr.Textbox(label="Character Personality", placeholder="e.g. outgoing, ambitious, creative")
            char_output = gr.Textbox(label="Character Information", interactive=False)

    scenario = gr.Textbox(label="Scenario", placeholder="e.g. {char} and {user} are meeting for coffee to discuss a potential business partnership.")
    combined_output = gr.Textbox(label="Combined User and Character Information", interactive=False)
    
    # Buttons
    submit_button = gr.Button("Submit User and Character Info")
    save_button = gr.Button("Save Character and User Info")
    load_button = gr.Button("Load Previous Info")
    save_message = gr.Textbox(label="Save/Load Status", interactive=False)

    # Define the actions for the buttons
    submit_button.click(
        lambda user_name, user_age, user_features, user_related, char_name, char_age, char_features, char_personality, scenario: (
            user(user_name, user_age, user_features, user_related, char_name),
            char(char_name, char_age, char_features, char_personality),
            combine_info(
                user(user_name, user_age, user_features, user_related, char_name),
                char(char_name, char_age, char_features, char_personality),
                scenario
            )
        ),
        inputs=[user_name, user_age, user_features, user_related, char_name, char_age, char_features, char_personality, scenario],
        outputs=[user_output, char_output, combined_output]
    )

    save_button.click(
        lambda user_name, user_age, user_features, user_related, char_name, char_age, char_features, char_personality, combined_info: save_character(
            char(char_name, char_age, char_features, char_personality),
            user(user_name, user_age, user_features, user_related, char_name),
            combined_info
        ),
        inputs=[user_name, user_age, user_features, user_related, char_name, char_age, char_features, char_personality, combined_output],
        outputs=[save_message]
    )
    
    load_button.click(
        load_character,
        inputs=[],
        outputs=[user_name, user_age, user_features, user_related, char_name, char_age, char_features, char_personality, combined_output, save_message]
    )

demo.launch()