import os
import csv
import freecad
import FreeCAD
import Part
import Import

DATA_ROOT = "../data/a1.0.0_00"
OUTPUT_STEP_DIR = "../data/s5_step"

def open_doc(doc_path):
    Import.open(doc_path)
    doc = FreeCAD.ActiveDocument
    return doc

def close_doc():
    docs = FreeCAD.listDocuments()
    for name in list(docs.keys()):
        FreeCAD.closeDocument(name)

def get_all_objs(doc):
    return [obj for obj in doc.Objects]

def merge_steps(step_paths, output_path):
    for name in list(FreeCAD.listDocuments().keys()):
        FreeCAD.closeDocument(name)
    FreeCAD.newDocument("MergeDoc")
    objs = []
    for path in step_paths:
        if not os.path.exists(path):
            print(f"[WARN] STEP not found: {path}")
            continue
        try:
            Import.open(path)
            doc = FreeCAD.ActiveDocument
            doc_objs = [obj for obj in doc.Objects]
            print(f"[DEBUG] {path} -> {len(doc_objs)} objects")
            objs.extend(doc_objs)
        except Exception as e:
            print(f"[ERROR] Failed to open {path}: {e}")
    if objs:
        try:
            Part.export(objs, output_path)
            print(f"[OK] Exported merged STEP: {output_path}")
        except Exception as e:
            print(f"[FAIL] Export failed: {e}")
    else:
        print("[FAIL] No objects to export.")
    for name in list(FreeCAD.listDocuments().keys()):
        FreeCAD.closeDocument(name)



def main():
    input_csv = "../result/s4_output.csv"
    os.makedirs(OUTPUT_STEP_DIR, exist_ok=True)
    with open(input_csv, "r", encoding="utf-8", newline='') as fin:
        reader = csv.DictReader(fin)
        for row in reader:
            assembly = row["Assembly"].strip()
            matched_names = [x.strip() for x in row.get("Matched_Part_Names", "").split(";") if x.strip()]
            step_paths = [
                os.path.join(DATA_ROOT, assembly, f"{name}.step") for name in matched_names
            ]
            output_step = os.path.join(OUTPUT_STEP_DIR, f"{assembly}.step")
            print(f"[INFO] Merging {step_paths} -> {output_step}")
            merge_steps(step_paths, output_step)

if __name__ == "__main__":
    main()
