import streamlit as st
from PIL import Image
import cv2
import numpy as np
import ezdxf
import io

def preprocess_image(image_path):
    try:
        img_pil = Image.open(image_path)
        img = np.array(img_pil.convert('L'))  # Convert to grayscale
    except Exception as e:
        st.error(f"Error loading image: {e}")
        return None

    if img is None:
        raise ValueError("Failed to load image. Check the image path.")
    
    # Apply GaussianBlur to smooth the image
    img_blurred = cv2.GaussianBlur(img, (5, 5), 0)
    
    # Edge detection using Canny
    edges = cv2.Canny(img_blurred, threshold1=100, threshold2=200)
    
    # Invert the edge image
    edges_inverted = cv2.bitwise_not(edges)
    
    return edges_inverted

def convert_sketch_to_dxf(image_path, output_dxf):
    edges_inverted = preprocess_image(image_path)
    if edges_inverted is None:
        st.error("Edge detection failed!")
        return None
    
    edges_inverted = cv2.threshold(edges_inverted, 127, 255, cv2.THRESH_BINARY)[1]
    contours, _ = cv2.findContours(edges_inverted, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    doc = ezdxf.new()
    msp = doc.modelspace()
    
    for contour in contours:
        for point in contour:
            x, y = point[0]
            msp.add_line((x, y), (x+1, y+1))
    
    # Create a buffer to write the DXF data
    dxf_buffer = io.BytesIO()
    
    # Write the DXF to the buffer (ensure that the buffer is in binary write mode)
    doc.write(dxf_buffer)
    dxf_buffer.seek(0)  # Reset buffer pointer to the start
    
    return dxf_buffer

# Streamlit app interface
st.title('Sketch to DXF Converter')

image_file = st.file_uploader("Upload your sketch", type=["jpeg", "jpg", "png"])

if image_file:
    # Save the uploaded file temporarily
    image_path = "uploaded_image.jpeg"
    with open(image_path, "wb") as f:
        f.write(image_file.getbuffer())
    
    output_dxf = "output_sketch.dxf"
    dxf_buffer = convert_sketch_to_dxf(image_path, output_dxf)
    
    if dxf_buffer:
        # Provide download button for the generated DXF file
        st.download_button(
            label="Download DXF file",
            data=dxf_buffer,
            file_name="output_sketch.dxf",
            mime="application/dxf"
        )