import streamlit as st
import json
from groq import Groq

def load_config():
    with open('config.json', 'r') as f:
        return json.load(f)

def checkbox_with_default(label, value, key):
    return st.checkbox(label, value=value, key=key)

def text_input_with_default(label, value, key):
    return st.text_input(label, value=value, key=key)

def selectbox_with_default(label, options, value, key):
    return st.selectbox(label, options, index=options.index(value), key=key)

def get_user_input(config):
    st.title("Advanced Roblox Lua Script Generator")

    default_params = config["default_parameters"]

    # Game Object Properties
    color_enabled = checkbox_with_default("Enable Color", True, "color_enabled")
    size_enabled = checkbox_with_default("Enable Size", True, "size_enabled")
    position_enabled = checkbox_with_default("Enable Position", True, "position_enabled")
    shape_enabled = checkbox_with_default("Enable Shape", True, "shape_enabled")
    material_enabled = checkbox_with_default("Enable Material", True, "material_enabled")
    transparency_enabled = checkbox_with_default("Enable Transparency", True, "transparency_enabled")
    reflectance_enabled = checkbox_with_default("Enable Reflectance", True, "reflectance_enabled")
    rotation_speed_enabled = checkbox_with_default("Enable Rotation Speed", True, "rotation_speed_enabled")
    bounce_height_enabled = checkbox_with_default("Enable Bounce Height", True, "bounce_height_enabled")
    collision_enabled = checkbox_with_default("Enable Collision", True, "collision_enabled")
    surface_enabled = checkbox_with_default("Enable Surface Type", True, "surface_enabled")
    anchor_enabled = checkbox_with_default("Enable Anchor", True, "anchor_enabled")
    collision_group_enabled = checkbox_with_default("Enable Collision Group", True, "collision_group_enabled")
    custom_physics_enabled = checkbox_with_default("Enable Custom Physics", True, "custom_physics_enabled")
    part_cframe_enabled = checkbox_with_default("Enable Part CFrame", True, "part_cframe_enabled")

    # Camera Properties
    camera_enabled = checkbox_with_default("Enable Camera Configuration", True, "camera_enabled")
    camera_interpolation_enabled = checkbox_with_default("Enable Camera Interpolation", default_params["camera_interpolation"], "camera_interpolation_enabled")

    # Lighting Properties
    lighting_enabled = checkbox_with_default("Enable Lighting Configuration", True, "lighting_enabled")
    fog_enabled = checkbox_with_default("Enable Environmental Fog", True, "fog_enabled")
    shadow_enabled = checkbox_with_default("Enable Shadow Properties", True, "shadow_enabled")

    # Player Properties
    player_settings_enabled = checkbox_with_default("Enable Player Settings", True, "player_settings_enabled")
    respawn_time_enabled = checkbox_with_default("Enable Respawn Time", True, "respawn_time_enabled")

    # Sound Properties
    sound_enabled = checkbox_with_default("Enable Sound", True, "sound_enabled")
    sound_effects_enabled = checkbox_with_default("Enable Sound Effects", True, "sound_effects_enabled")
    environmental_sounds_enabled = checkbox_with_default("Enable Environmental Sounds", True, "environmental_sounds_enabled")

    # UI Properties
    ui_enabled = checkbox_with_default("Enable UI Configuration", True, "ui_enabled")

    # Decals
    decals_enabled = checkbox_with_default("Enable Decals", True, "decals_enabled")

    # Particle Effects
    particles_enabled = checkbox_with_default("Enable Particle Effects", True, "particles_enabled")

    # Advanced Game Object Properties Inputs
    if color_enabled:
        color = text_input_with_default("Enter Color (RGB, e.g., 255,0,0)", default_params["color"], "color")
    else:
        color = None

    if size_enabled:
        size = text_input_with_default("Enter Size (X,Y,Z, e.g., 4,1,2)", default_params["size"], "size")
    else:
        size = None

    if position_enabled:
        position = text_input_with_default("Enter Position (X,Y,Z, e.g., 0,10,0)", default_params["position"], "position")
    else:
        position = None

    if shape_enabled:
        shape = selectbox_with_default("Select Shape", config["shapes"], default_params["shape"], "shape")
    else:
        shape = None

    if material_enabled:
        material = selectbox_with_default("Select Material", config["materials"], default_params["material"], "material")
    else:
        material = None

    if transparency_enabled:
        transparency = text_input_with_default("Enter Transparency (0-1, e.g., 0.5)", default_params["transparency"], "transparency")
    else:
        transparency = None

    if reflectance_enabled:
        reflectance = text_input_with_default("Enter Reflectance (0-1, e.g., 0.5)", default_params["reflectance"], "reflectance")
    else:
        reflectance = None

    if rotation_speed_enabled:
        rotation_speed = text_input_with_default("Enter Rotation Speed (e.g., 10)", default_params["rotation_speed"], "rotation_speed")
    else:
        rotation_speed = None

    if bounce_height_enabled:
        bounce_height = text_input_with_default("Enter Bounce Height (e.g., 10)", default_params["bounce_height"], "bounce_height")
    else:
        bounce_height = None

    if collision_enabled:
        collision = selectbox_with_default("Enable Collision", config["collisions"], default_params["collision"], "collision")
    else:
        collision = None

    if surface_enabled:
        surface = selectbox_with_default("Select Surface Type", config["surfaces"], default_params["surface"], "surface")
    else:
        surface = None

    if anchor_enabled:
        anchor = selectbox_with_default("Enable Anchor", ["true", "false"], default_params["anchor"], "anchor")
    else:
        anchor = None

    if collision_group_enabled:
        collision_group = selectbox_with_default("Select Collision Group", config["collision_groups"], default_params["collision_group"], "collision_group")
    else:
        collision_group = None

    if custom_physics_enabled:
        angular_velocity = text_input_with_default("Enter Angular Velocity (X,Y,Z, e.g., 0,0,0)", default_params["angular_velocity"], "angular_velocity")
        linear_velocity = text_input_with_default("Enter Linear Velocity (X,Y,Z, e.g., 0,0,0)", default_params["linear_velocity"], "linear_velocity")
        density = text_input_with_default("Enter Density", default_params["density"], "density")
    else:
        angular_velocity = None
        linear_velocity = None
        density = None

    if part_cframe_enabled:
        part_cframe = text_input_with_default("Enter Part CFrame (e.g., CFrame.new(0, 50, -100))", default_params["part_cframe"], "part_cframe")
    else:
        part_cframe = None

    # Camera Properties Inputs
    if camera_enabled:
        camera_type = selectbox_with_default("Select Camera Type", config["camera_types"], default_params["camera_type"], "camera_type")
        camera_position = text_input_with_default("Enter Camera Position (X,Y,Z, e.g., 0,50,-100)", default_params["camera_position"], "camera_position")
        field_of_view = text_input_with_default("Enter Field of View", default_params["field_of_view"], "field_of_view")
    else:
        camera_type = None
        camera_position = None
        field_of_view = None

    if camera_interpolation_enabled:
        camera_interpolation = checkbox_with_default("Enable Camera Interpolation", default_params["camera_interpolation"], "camera_interpolation")
    else:
        camera_interpolation = None

    # Lighting Properties Inputs
    if lighting_enabled:
        ambient_light = text_input_with_default("Enter Ambient Light (RGB, e.g., 128,128,128)", default_params["ambient_light"], "ambient_light")
        brightness = text_input_with_default("Enter Brightness", default_params["brightness"], "brightness")
        time_of_day = selectbox_with_default("Select Time of Day", config["times_of_day"], default_params["time_of_day"], "time_of_day")
        shadow_softness = text_input_with_default("Enter Shadow Softness (0-1, e.g., 0.5)", default_params["shadow_softness"], "shadow_softness")
        shadow_color = text_input_with_default("Enter Shadow Color (RGB, e.g., 0,0,0)", default_params["shadow_color"], "shadow_color")
    else:
        ambient_light = None
        brightness = None
        time_of_day = None
        shadow_softness = None
        shadow_color = None

    if fog_enabled:
        fog_start = text_input_with_default("Enter Fog Start Distance", default_params["fog_start"], "fog_start")
        fog_end = text_input_with_default("Enter Fog End Distance", default_params["fog_end"], "fog_end")
        fog_color = text_input_with_default("Enter Fog Color (RGB, e.g., 128,128,128)", default_params["fog_color"], "fog_color")
    else:
        fog_start = None
        fog_end = None
        fog_color = None

    # Player Properties Inputs
    if player_settings_enabled:
        walk_speed = text_input_with_default("Enter Walk Speed", default_params["walk_speed"], "walk_speed")
        jump_power = text_input_with_default("Enter Jump Power", default_params["jump_power"], "jump_power")
        health = text_input_with_default("Enter Health", default_params["health"], "health")
    else:
        walk_speed = None
        jump_power = None
        health = None

    if respawn_time_enabled:
        respawn_time = text_input_with_default("Enter Respawn Time", default_params["respawn_time"], "respawn_time")
    else:
        respawn_time = None

    # Sound Properties Inputs
    if sound_enabled:
        background_music = text_input_with_default("Enter Background Music ID (e.g., rbxassetid://123456789)", default_params["background_music"], "background_music")
        sound_volume = text_input_with_default("Enter Sound Volume (e.g., 1)", default_params["sound_volume"], "sound_volume")
    else:
        background_music = None
        sound_volume = None

    if sound_effects_enabled:
        sound_effects = text_input_with_default("Enter Sound Effects ID (e.g., rbxassetid://987654321)", default_params["sound_effects"], "sound_effects")
    else:
        sound_effects = None

    if environmental_sounds_enabled:
        environmental_sounds = text_input_with_default("Enter Environmental Sounds ID (e.g., rbxassetid://1122334455)", default_params["environmental_sounds"], "environmental_sounds")
    else:
        environmental_sounds = None

    # UI Properties Inputs
    if ui_enabled:
        ui_text = text_input_with_default("Enter UI Text", default_params["ui_text"], "ui_text")
        ui_position = text_input_with_default("Enter UI Position (X,Y,Z, e.g., 0.5,-100,0.5,-25)", default_params["ui_position"], "ui_position")
        ui_size = text_input_with_default("Enter UI Size (X,Y,Z, e.g., 0,200,0,50)", default_params["ui_size"], "ui_size")
        ui_color = text_input_with_default("Enter UI Color (RGB, e.g., 255,255,255)", default_params["ui_color"], "ui_color")
        ui_font = selectbox_with_default("Select UI Font", config["fonts"], default_params["ui_font"], "ui_font")
        ui_text_size = text_input_with_default("Enter UI Text Size", default_params["ui_text_size"], "ui_text_size")
    else:
        ui_text = None
        ui_position = None
        ui_size = None
        ui_color = None
        ui_font = None
        ui_text_size = None

    # Decals Inputs
    if decals_enabled:
        texture_id = text_input_with_default("Enter Texture ID (e.g., rbxassetid://123456789)", default_params["texture_id"], "texture_id")
        decal_position = selectbox_with_default("Select Decal Position", config["decal_positions"], default_params["decal_position"], "decal_position")
    else:
        texture_id = None
        decal_position = None

    # Particle Effects Inputs
    if particles_enabled:
        particle_texture = text_input_with_default("Enter Particle Texture ID (e.g., rbxassetid://123456789)", default_params["particle_texture"], "particle_texture")
        particle_lifetime = text_input_with_default("Enter Particle Lifetime (e.g., 1,2)", default_params["particle_lifetime"], "particle_lifetime")
        particle_rate = text_input_with_default("Enter Particle Rate (e.g., 100)", default_params["particle_rate"], "particle_rate")
        particle_speed = text_input_with_default("Enter Particle Speed (e.g., 10,20)", default_params["particle_speed"], "particle_speed")
        particle_size = text_input_with_default("Enter Particle Size (e.g., 1,2)", default_params["particle_size"], "particle_size")
        particle_acceleration = text_input_with_default("Enter Particle Acceleration (X,Y,Z, e.g., 0,-10,0)", default_params["particle_acceleration"], "particle_acceleration")
        particle_rot_speed = text_input_with_default("Enter Particle Rotation Speed (e.g., 100,200)", default_params["particle_rot_speed"], "particle_rot_speed")
        particle_color = text_input_with_default("Enter Particle Color (RGB, e.g., 255,255,255)", default_params["particle_color"], "particle_color")
    else:
        particle_texture = None
        particle_lifetime = None
        particle_rate = None
        particle_speed = None
        particle_size = None
        particle_acceleration = None
        particle_rot_speed = None
        particle_color = None

    custom_prompt_enabled = checkbox_with_default("Enable Custom Prompt", False, "custom_prompt_enabled")

    custom_prompt = st.text_area("Enter Custom Prompt", "Create a Lua script for Roblox:", key="custom_prompt")

    return {
        "color": color, "size": size, "position": position, "shape": shape, "material": material,
        "transparency": transparency, "reflectance": reflectance, "rotation_speed": rotation_speed,
        "bounce_height": bounce_height, "collision": collision, "surface": surface, "anchor": anchor,
        "collision_group": collision_group, "angular_velocity": angular_velocity, "linear_velocity": linear_velocity,
        "density": density, "part_cframe": part_cframe, "camera_type": camera_type, "camera_position": camera_position,
        "field_of_view": field_of_view, "camera_interpolation": camera_interpolation, "ambient_light": ambient_light,
        "brightness": brightness, "time_of_day": time_of_day, "shadow_softness": shadow_softness, "shadow_color": shadow_color,
        "fog_start": fog_start, "fog_end": fog_end, "fog_color": fog_color, "walk_speed": walk_speed, "jump_power": jump_power,
        "health": health, "respawn_time": respawn_time, "background_music": background_music, "sound_volume": sound_volume,
        "sound_effects": sound_effects, "environmental_sounds": environmental_sounds, "ui_text": ui_text, "ui_position": ui_position,
        "ui_size": ui_size, "ui_color": ui_color, "ui_font": ui_font, "ui_text_size": ui_text_size, "texture_id": texture_id,
        "decal_position": decal_position, "particle_texture": particle_texture, "particle_lifetime": particle_lifetime,
        "particle_rate": particle_rate, "particle_speed": particle_speed, "particle_size": particle_size, 
        "particle_acceleration": particle_acceleration, "particle_rot_speed": particle_rot_speed, "particle_color": particle_color,
        "custom_prompt": custom_prompt
    }

def generate_lua_script(params, api_key):
    client = Groq(api_key=api_key)
    prompt = params["custom_prompt"]
    
    completion = client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=True,
        stop=None,
    )

    lua_script = ""
    for chunk in completion:
        lua_script += chunk.choices[0].delta.content or ""

    return lua_script

def main():
    config = load_config()
    user_input = get_user_input(config)

    if st.button("Generate Script"):
        api_key = config["api_key"]
        lua_script = generate_lua_script(user_input, api_key)
        st.code(lua_script, language='lua')
        st.download_button(label="Copy to Clipboard", data=lua_script, file_name="lua_script.txt", mime="text/plain")

if __name__ == "__main__":
    main()