import streamlit as st
import cadquery as cq
from datasets import load_dataset
import json

# Load dataset
dataset = load_dataset("FreedomIntelligence/CADBench")

# Helper to generate Cargo Crane CAD
def generate_cargo_crane(inputs):
    # Parse inputs
    boom_length = inputs.get("Boom Length (in mm)", 3000)
    lifting_capacity = inputs.get("Lifting Capacity (in tons)", 10)
    base_height = inputs.get("Base Height (in mm)", 500)
    material_type = inputs.get("Material Type", "Steel")

    # Create base
    base = cq.Workplane("XY").box(200, 200, base_height)

    # Create boom
    boom = cq.Workplane("XY").box(50, boom_length, 50).translate((0, base_height, 0))

    # Combine parts
    crane = base.union(boom)

    return crane

# Streamlit Interface
st.title("Quick CAD Model Generator")
case = st.radio("Select Case", ["DFM Analysis", "Predefined Template"])

if case == "Predefined Template":
    st.subheader("Predefined Template Selection")
    template_names = dataset["train"]["name"]
    selected_template = st.selectbox("Select a Template", template_names)

    # Retrieve criteria
    if selected_template:
        template_data = dataset["train"].filter(lambda x: x["name"] == selected_template).to_pandas()
        criteria = template_data.iloc[0]["criteria"]

        # Dynamic input collection
        inputs = {}
        for question in criteria.keys():
            inputs[question] = st.text_input(f"{question}: ")

        if st.button("Generate CAD Model"):
            if selected_template == "Cargo Crane":
                cad_model = generate_cargo_crane(inputs)
                st.success("CAD Model Generated!")
                cq.exporters.export(cad_model, "cargo_crane.step")
                cq.exporters.export(cad_model, "cargo_crane.stl")
                st.download_button("Download STEP File", open("cargo_crane.step", "rb").read(), "cargo_crane.step")
                st.download_button("Download STL File", open("cargo_crane.stl", "rb").read(), "cargo_crane.stl")


if case == "DFM Analysis":
    st.subheader("Case 1: DFM Analysis")
    cad_file = st.file_uploader("Upload a CAD file (.stl, .step, .dwg)", type=["stl", "step", "dwg"])
    if cad_file:
        file_format = os.path.splitext(cad_file.name)[1][1:]
        analysis_result = analyze_dfm(cad_file, file_format)
        st.success(analysis_result)

elif case == "Predefined Template Selection":
    st.subheader("Case 2: Predefined Template Selection")
    template_names = dataset["name"]
    selected_template = st.selectbox("Select a template", template_names)

    if selected_template:
        st.write(f"Selected Template: {selected_template}")
        template_data = dataset.filter(lambda x: x["name"] == selected_template).to_pandas()
        criteria = template_data.iloc[0]["criteria"]  # Already a dict

        responses = {}
        for criterion, prompt in criteria.items():
            prefilled_text = f"Enter {criterion} ({prompt}):"
            response = st.text_input(prefilled_text, key=criterion)
            if response:
                responses[criterion] = response

        if st.button("Generate CAD Model"):
            try:
                model = generate_cad_model(responses)
                file_format = st.selectbox("Select file format for download", ["stl", "step", "dwg"])
                file_path = export_cad(model, file_format)
                st.success("CAD Model Generated Successfully!")
                with open(file_path, "rb") as f:
                    st.download_button("Download CAD File", data=f.read(), file_name=f"model.{file_format}")
            except Exception as e:
                st.error(f"Error generating CAD model: {e}")