import bonsai.tool as tool
import ifcopenshell
import ifcopenshell.api
import bpy
import numpy as np

# Load the IFC File in Bonsai
ifc_file = tool.Ifc.get()

# Load the Blender (.blend) File
blend_filepath = "G:\\My Drive\\Buildtrade\\Kiri\\LR3.blend"

# Get IFC Context & Storey
contexts = ifc_file.by_type("IfcGeometricRepresentationContext")
if not contexts:
    raise ValueError("No IfcGeometricRepresentationContext found in IFC file.")
context = contexts[0]

storeys = ifc_file.by_type("IfcBuildingStorey")
if not storeys:
    raise ValueError("No IfcBuildingStorey found in the IFC file.")
ground_storey = storeys[0]

# IFC Class Mapping for Future Extensions
ifc_class_map = {
    "wall": "IfcWall",
    "furniture": "IfcFurniture",
    "appliance": "IfcFlowTerminal"
}

def create_ifc_object(blender_obj, ifc_class):
    # Extract Geometry
    if not blender_obj.data or not hasattr(blender_obj.data, "vertices"):
        return None

    vertices = np.array([blender_obj.matrix_world @ v.co for v in blender_obj.data.vertices])
    faces = [list(polygon.vertices) for polygon in blender_obj.data.polygons]

    # Create IFC Object
    ifc_entity = ifcopenshell.api.run("root.create_entity", ifc_file, ifc_class=ifc_class)
    ifc_entity.Name = f"Traced_{blender_obj.name}"
    ifc_entity.Description = "IFC object created from Blender mesh"

    # Assign to storey
    ifcopenshell.api.run("spatial.assign_container", ifc_file, relating_structure=ground_storey, products=[ifc_entity])

    # Create IfcCartesianPoints
    ifc_points = [ifc_file.create_entity("IfcCartesianPoint", Coordinates=[round(float(c), 6) for c in v]) for v in vertices]

    # Create IfcFaceBound for each face
    ifc_face_bounds = []
    for face in faces:
        poly_loop = ifc_file.create_entity("IfcPolyLoop", Polygon=[ifc_points[i] for i in face])
        face_bound = ifc_file.create_entity("IfcFaceBound", Bound=poly_loop, Orientation=True)
        ifc_face_bounds.append(ifc_file.create_entity("IfcFace", Bounds=[face_bound]))

    # Create IfcClosedShell and IfcFacetedBrep
    ifc_closed_shell = ifc_file.create_entity("IfcClosedShell", CfsFaces=ifc_face_bounds)
    ifc_brep = ifc_file.create_entity("IfcFacetedBrep", Outer=ifc_closed_shell)

    # Create IfcShapeRepresentation
    shape_representation = ifc_file.create_entity(
        "IfcShapeRepresentation",
        ContextOfItems=context,
        RepresentationIdentifier="Body",
        RepresentationType="Brep",
        Items=[ifc_brep]
    )

    # Create IfcProductDefinitionShape
    product_shape = ifc_file.create_entity("IfcProductDefinitionShape", Representations=[shape_representation])
    ifc_entity.Representation = product_shape
    
    return ifc_entity

# Loop through all objects with "wall_" prefix
bpy.ops.wm.open_mainfile(filepath=blend_filepath)
for obj in bpy.data.objects:
    if obj.name.startswith("wall_"):
        create_ifc_object(obj, ifc_class_map.get("wall", "IfcWall"))

# Save the Modified IFC File
ifc_filepath = "G:\\My Drive\\Buildtrade\\Kiri\\LR10C_mod.ifc"
ifc_file.write(ifc_filepath)

print(f"✅ IFC file saved to: {ifc_filepath}")