import streamlit as st
import ezdxf
import tempfile
import pandas as pd
import os

# ---------- Function to extract data from DXF ----------
def extract_entities(uploaded_file):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".dxf") as tmp:
        tmp.write(uploaded_file.read())
        tmp_path = tmp.name

    doc = ezdxf.readfile(tmp_path)
    msp = doc.modelspace()

    plan_data = {"rooms": [], "doors": [], "windows": []}
    section_data = {"slabs": [], "beams": [], "columns": []}

    for e in msp:
        if e.dxftype() == "LWPOLYLINE" and "plan" in e.dxf.layer.lower():
            points = e.get_points()
            plan_data["rooms"].append({"points": points})
        elif e.dxftype() == "LINE" and "door" in e.dxf.layer.lower():
            plan_data["doors"].append({"start": e.dxf.start, "end": e.dxf.end})
        elif e.dxftype() == "LINE" and "window" in e.dxf.layer.lower():
            plan_data["windows"].append({"start": e.dxf.start, "end": e.dxf.end})
        elif "section" in e.dxf.layer.lower():
            if "slab" in e.dxf.layer.lower():
                section_data["slabs"].append(e)
            elif "beam" in e.dxf.layer.lower():
                section_data["beams"].append(e)
            elif "column" in e.dxf.layer.lower():
                section_data["columns"].append(e)

    return plan_data, section_data

# ---------- Function to estimate quantities ----------
def estimate_quantities(plan_data, section_data):
    items = []

    # Assumptions (typical values based on AU construction norms)
    wall_height_ft = 10
    wall_thickness_in = 9
    brick_volume_in3 = 9 * 4.5 * 3
    concrete_mix_ratio = {'cement': 1, 'sand': 1.5, 'crush': 3}
    dry_volume_factor = 1.54
    cement_bag_volume_cuft = 1.25  # 1 bag = 1.25 cuft of cement

    # === 1. Brick masonry estimation ===
    for room in plan_data["rooms"]:
        points = room["points"]
        if len(points) < 2:
            continue
        perimeter = 0
        for i in range(len(points)):
            x1, y1 = points[i][0:2]
            x2, y2 = points[(i+1) % len(points)][0:2]
            dist = ((x2 - x1)**2 + (y2 - y1)**2)**0.5
            perimeter += dist

        height_in = wall_height_ft * 12
        wall_volume_in3 = perimeter * height_in * wall_thickness_in
        wall_volume_cft = wall_volume_in3 / 1728
        num_bricks = wall_volume_in3 / brick_volume_in3

        items.append({
            "Item": "Brick Masonry",
            "Length": round(perimeter, 2),
            "Width": wall_thickness_in,
            "Height": height_in,
            "Quantity": round(wall_volume_cft, 2),
            "Unit": "CFT"
        })
        items.append({
            "Item": "Bricks",
            "Length": "-",
            "Width": "-",
            "Height": "-",
            "Quantity": round(num_bricks),
            "Unit": "No"
        })

    # === 2. Concrete works (slabs, beams, columns) ===
    concrete_items = section_data["slabs"] + section_data["beams"] + section_data["columns"]

    for component in concrete_items:
        # Dummy default dimensions (can later be detected from DXF entities)
        length = 12  # ft
        width = 1    # ft
        depth = 0.5  # ft

        volume = length * width * depth
        dry_volume = volume * dry_volume_factor
        total_parts = sum(concrete_mix_ratio.values())

        cement_cuft = dry_volume * (concrete_mix_ratio['cement'] / total_parts)
        cement_bags = cement_cuft / cement_bag_volume_cuft

        sand_cuft = dry_volume * (concrete_mix_ratio['sand'] / total_parts)
        crush_cuft = dry_volume * (concrete_mix_ratio['crush'] / total_parts)

        steel_kg = volume * 80  # approx 80 kg/m3 (AU rule of thumb)

        items.append({
            "Item": "Concrete",
            "Length": length,
            "Width": width,
            "Height": depth,
            "Quantity": round(volume, 2),
            "Unit": "CFT"
        })
        items.append({
            "Item": "Cement",
            "Length": "-",
            "Width": "-",
            "Height": "-",
            "Quantity": round(cement_bags, 1),
            "Unit": "Bags"
        })
        items.append({
            "Item": "Sand",
            "Length": "-",
            "Width": "-",
            "Height": "-",
            "Quantity": round(sand_cuft, 1),
            "Unit": "CFT"
        })
        items.append({
            "Item": "Crush",
            "Length": "-",
            "Width": "-",
            "Height": "-",
            "Quantity": round(crush_cuft, 1),
            "Unit": "CFT"
        })
        items.append({
            "Item": "Steel",
            "Length": "-",
            "Width": "-",
            "Height": "-",
            "Quantity": round(steel_kg, 1),
            "Unit": "kg"
        })

    df = pd.DataFrame(items)
    return df

# ---------- Streamlit App UI ----------
st.set_page_config(layout="centered")
st.title("🏗️ AutoCAD Drawing Estimator (AU Standards)")

uploaded_file = st.file_uploader("Upload AutoCAD DXF File", type=["dxf"])

if uploaded_file:
    with st.spinner("Reading and analyzing the drawing..."):
        try:
            plan_data, section_data = extract_entities(uploaded_file)
            df = estimate_quantities(plan_data, section_data)

            st.success("✅ Drawing processed and estimate generated.")
            st.dataframe(df)

            output_path = "Construction_Estimate.xlsx"
            df.to_excel(output_path, index=False)

            with open(output_path, "rb") as f:
                st.download_button("📥 Download Estimate (Excel)", f, file_name=output_path)

        except Exception as e:
            st.error(f"Error processing file: {e}")