import gradio as gr
from transformers import pipeline
import tempfile
import os

# Load zero-shot classification model
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Define candidate screw sizes and head types
screw_sizes = ["M3 screw", "M4 screw", "M6 screw"]
head_types = ["hex head", "flat head", "pan head", "round head", "socket head"]

def parse_prompt_to_scad(prompt: str):
    # Detect screw size
    screw_result = classifier(prompt, screw_sizes)
    screw_type = screw_result["labels"][0]

    # Detect head type
    head_result = classifier(prompt, head_types)
    head_type = head_result["labels"][0]

    # Define dimensions (approximate values)
    if screw_type == "M3 screw":
        d = 3
        h = 15
        pitch = 0.5
    elif screw_type == "M4 screw":
        d = 4
        h = 20
        pitch = 0.7
    elif screw_type == "M6 screw":
        d = 6
        h = 30
        pitch = 1.0
    else:
        d = 5
        h = 25
        pitch = 0.8

    # Head geometry (basic representation)
    head_code = ""
    if head_type == "hex head":
        head_code = f"cylinder(d=10, h=4); translate([0,0,4])"
    elif head_type == "flat head":
        head_code = f"cylinder(d1=10, d2={d}, h=4); translate([0,0,4])"
    elif head_type == "pan head":
        head_code = f"sphere(r=5); translate([0,0,5])"
    elif head_type == "round head":
        head_code = f"sphere(r=5); translate([0,0,5])"
    elif head_type == "socket head":
        head_code = f"cylinder(d=8, h=5); translate([0,0,5])"

    # Thread simulation (simplified)
    threads = f"""
// Simulated threads using rotate_extrude
module thread() {{
    translate([d/2, 0, 0])
        rotate_extrude(angle=360, $fn=100)
            polygon(points=[[0,0],[pitch/2,0],[pitch/2,1],[0,1]]);
}}

translate([0,0,0])
for(i = [0:{int(h/pitch)-1}])
    translate([0,0,i*pitch])
        thread();
"""

    # Final OpenSCAD script
    scad_script = f"""
// Auto-generated screw with thread and {head_type}
d = {d};
h = {h};
pitch = {pitch};

// Head
{head_code}

// Shaft
cylinder(d=d, h=h, center=false);

// Threads
{threads}
"""

    return scad_script.strip(), screw_type.replace(" ", "_"), head_type.replace(" ", "_"), d, h

def generate_scad_file(prompt: str):
    scad_script, screw_type, head_type, d, h = parse_prompt_to_scad(prompt)
    filename = f"{screw_type}_{head_type}_d{d}_h{h}.scad"
    filepath = os.path.join(tempfile.gettempdir(), filename)

    with open(filepath, "w") as f:
        f.write(scad_script)
    return filepath

iface = gr.Interface(
    fn=generate_scad_file,
    inputs=gr.Textbox(label="Enter prompt (e.g. 'Make an M6 hex head screw with threads')"),
    outputs=gr.File(label="Download .scad file"),
    title="Threaded Screw Generator",
    description="Generates OpenSCAD scripts for threaded screws with customizable head types."
)

if __name__ == "__main__":
    iface.launch()