import os
import gradio as gr
from ansys.mapdl.core import launch_mapdl


def generate_apdl_script(step_file_path, material_e, material_nu, mesh_size, load_value):
    """
    Generate a detailed APDL script for simulation.
    Parameters:
    - step_file_path: Path to the uploaded STEP file.
    - material_e: Young's Modulus of the material (Pa).
    - material_nu: Poisson's Ratio of the material.
    - mesh_size: Element size for meshing.
    - load_value: Load to apply (N).
    Returns:
    - Path to the generated APDL script.
    """
    file_name = os.path.basename(step_file_path)
    part_name = os.path.splitext(file_name)[0]

    # APDL script content
    apdl_script = f"""
/prep7
! Import STEP geometry
/import, '{part_name}', step
! Define material properties
mp, ex, 1, {material_e}
mp, nuxy, 1, {material_nu}
! Set element type
et, 1, solid186
! Mesh the model
esize, {mesh_size}
amesh, all
! Apply boundary conditions
nsel, s, loc, x, 0
d, all, ux, 0
nsel, s, loc, y, 0
d, all, uy, 0
nsel, s, loc, z, 0
d, all, uz, 0
! Apply loads
nsel, s, loc, z, max
f, all, fz, {load_value}
! Solve the model
/solu
antype, static
solve
! Post-processing
/post1
set, last
! Write results
prnsol, s, comp
prnsol, u, comp
/exit
"""

    # Save APDL script
    apdl_file_path = f"/tmp/{part_name}_simulation_apdl.txt"
    with open(apdl_file_path, "w") as f:
        f.write(apdl_script)

    return apdl_file_path


def run_simulation(step_file, material_e, material_nu, mesh_size, load_value):
    """
    Generate the APDL script, run the simulation, and fetch results.
    """
    # Generate the APDL script
    apdl_file_path = generate_apdl_script(step_file.name, material_e, material_nu, mesh_size, load_value)

    # Launch ANSYS MAPDL (ensure it's installed and licensed)
    mapdl = launch_mapdl()

    # Upload STEP file and APDL script to MAPDL
    mapdl.upload(step_file.name)
    mapdl.input(apdl_file_path)

    # Run the APDL script
    mapdl.run("/SOLU")
    mapdl.solve()

    # Post-process results
    stress_output = mapdl.result.principal_stress()
    displacement_output = mapdl.result.nodal_displacement()

    # Save outputs
    stress_file = "/tmp/stress_results.txt"
    with open(stress_file, "w") as f:
        f.write("Principal Stress Results:\n")
        f.write(str(stress_output))

    displacement_file = "/tmp/displacement_results.txt"
    with open(displacement_file, "w") as f:
        f.write("Nodal Displacement Results:\n")
        f.write(str(displacement_output))

    # Stop MAPDL
    mapdl.exit()

    return stress_file, displacement_file


# Gradio interface
def process_step_file(step_file, material_e, material_nu, mesh_size, load_value):
    stress_file, displacement_file = run_simulation(step_file, material_e, material_nu, mesh_size, load_value)
    return stress_file, displacement_file


# Define Gradio Interface
with gr.Blocks() as app:
    with gr.Row():
        step_file = gr.File(label="Upload STEP File")
        material_e = gr.Number(label="Young's Modulus (Pa)", value=210e9)
        material_nu = gr.Number(label="Poisson's Ratio", value=0.3)
        mesh_size = gr.Number(label="Mesh Size (mm)", value=10)
        load_value = gr.Number(label="Load Value (N)", value=1000)

    submit_button = gr.Button("Submit")
    stress_output = gr.File(label="Download Stress Results")
    displacement_output = gr.File(label="Download Displacement Results")

    submit_button.click(
        process_step_file,
        inputs=[step_file, material_e, material_nu, mesh_size, load_value],
        outputs=[stress_output, displacement_output],
    )

if __name__ == "__main__":
    app.queue()  # Ensure proper queuing for the submit button
    app.launch()