import os
import streamlit as st
from groq import Groq
from pyvista import read
from cadquery import exporters
from datasets import load_dataset
import json

# Initialize Groq API client
client = Groq(api_key="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")

# Case 1: File Upload and DFM Analysis
def process_uploaded_file(uploaded_file):
    file_path = f"uploaded_{uploaded_file.name}"
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    st.write(f"File saved as: {file_path}")
    return file_path

def perform_dfm_analysis(file_path):
    st.write("Performing DFM Analysis...")
    # Placeholder for actual DFM analysis logic using LLMs
    result = {"CNC_Machinable": True, "Complexity_Score": 0.75}  # Mock result
    return result

# Case 2: Template Selection and CAD Generation
def get_templates_from_dataset():
    st.write("Loading templates from CADBench dataset...")
    try:
        dataset = load_dataset("FreedomIntelligence/CADBench", split="train")
        templates = [entry["name"] for entry in dataset if "name" in entry]
        return templates
    except Exception as e:
        st.error(f"Error loading dataset: {e}")
        return []

def generate_cad_from_template(template_name, parameters):
    st.write(f"Generating CAD for template: {template_name}")
    # Placeholder for template CAD generation logic
    generated_file = f"{template_name}_output.step"
    return generated_file

# Streamlit App Interface
st.title("Quick CAD Model Generator")
st.sidebar.header("Select Case")
case = st.sidebar.radio("Choose a Case", ["Upload File for DFM", "Select Predefined Template"])

if case == "Upload File for DFM":
    uploaded_file = st.file_uploader("Upload your CAD file", type=["step", "stl", "obj", "dwg"])
    if uploaded_file:
        file_path = process_uploaded_file(uploaded_file)
        if st.button("Analyze for CNC"):
            try:
                dfm_result = perform_dfm_analysis(file_path)
                st.success("DFM Analysis Completed!")
                st.json(dfm_result)
            except Exception as e:
                st.error(f"Error in DFM Analysis: {e}")

elif case == "Select Predefined Template":
    templates = get_templates_from_dataset()
    if templates:
        template_name = st.selectbox("Choose a Template", templates)
        parameters = st.text_area("Enter Parameters (JSON format)", placeholder='{"length": 100, "width": 50}')
        if st.button("Generate CAD"):
            try:
                params = json.loads(parameters) if parameters else {}
                generated_file = generate_cad_from_template(template_name, params)
                st.success("CAD Model Generated!")
                st.download_button("Download CAD File", open(generated_file, "rb"), file_name=generated_file)
            except json.JSONDecodeError:
                st.error("Invalid JSON format for parameters.")
            except Exception as e:
                st.error(f"Error in CAD Generation: {e}")