#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Internal STEP sheet metal analyser

- Input:  ./test.step
- Unfold with K-factor = 0.5
- Auto-detect thickness if not provided
- Export:
    - flat.dxf
    - result.json
    - debug_last.FCStd
"""

import os
import sys
import json
import argparse
import traceback
from datetime import datetime


# -----------------------------
# Configuration
# -----------------------------

K_FACTOR = 0.5
K_STANDARD = "ansi"

DENSITY_KG_M3 = {
    "alu": 2700.0,
    "aluminum": 2700.0,
    "stainless": 8000.0,
    "edelstahl": 8000.0,
    "copper": 8960.0,
    "kupfer": 8960.0,
}


# -----------------------------
# Helpers
# -----------------------------

def mm2_to_m2(x):
    return x / 1_000_000.0


def mm_to_m(x):
    return x / 1000.0


def pick_main_object(doc):
    candidates = []
    for obj in doc.Objects:
        if hasattr(obj, "Shape") and obj.Shape:
            try:
                if obj.Shape.Solids:
                    candidates.append((float(obj.Shape.Volume), obj))
            except Exception:
                pass
    if not candidates:
        raise RuntimeError("No solid object found in STEP.")
    candidates.sort(key=lambda x: x[0], reverse=True)
    return candidates[0][1]


def pick_root_face_index(shape):
    planar = []
    all_faces = []

    for i, face in enumerate(shape.Faces, start=1):
        try:
            area = float(face.Area)
        except Exception:
            area = 0.0

        all_faces.append((area, i, face))
        try:
            surf = face.Surface
            if surf and "plane" in surf.__class__.__name__.lower():
                planar.append((area, i, face))
        except Exception:
            pass

    if planar:
        planar.sort(key=lambda x: x[0], reverse=True)
        return planar[0][1], True

    all_faces.sort(key=lambda x: x[0], reverse=True)
    return all_faces[0][1], False


def compute_bbox_mm(shape):
    bb = shape.BoundBox
    return float(bb.XLength), float(bb.YLength)


def estimate_thickness_mm_from_solid(part_obj, max_faces=60):
    """
    Robust thickness estimation for sheet-metal solids:
    - Collect planar faces
    - Cluster by (normalized) face normal direction
    - For each cluster, compute plane offsets d = n·p for face points
    - Thickness candidates are small positive differences between distinct d values
    - Fallback: use distToShape on representative face pairs

    Returns thickness in mm or None.
    """
    import math
    import Part

    shape = part_obj.Shape

    # --- helpers ---
    def vec_to_tuple(v):
        return (float(v.x), float(v.y), float(v.z))

    def norm(v):
        l = math.sqrt(v.x*v.x + v.y*v.y + v.z*v.z)
        if l <= 1e-12:
            return None
        return v.multiply(1.0 / l)

    def canonical_normal(n):
        """
        Make the normal direction canonical so +n and -n map consistently:
        Flip so the first significant component is positive.
        """
        # convert to tuple for easier checks
        x, y, z = float(n.x), float(n.y), float(n.z)
        # find first component with magnitude > eps
        eps = 1e-9
        if abs(x) > eps:
            if x < 0: return n.multiply(-1)
        elif abs(y) > eps:
            if y < 0: return n.multiply(-1)
        elif abs(z) > eps:
            if z < 0: return n.multiply(-1)
        return n

    def angle_close(n1, n2, cos_tol):
        # cos(angle) = n1·n2
        return (n1.dot(n2) >= cos_tol)

    def face_midpoint(face):
        u0, u1, v0, v1 = face.ParameterRange
        return face.valueAt((u0+u1)*0.5, (v0+v1)*0.5), face.normalAt((u0+u1)*0.5, (v0+v1)*0.5)

    # --- collect planar faces ---
    planar = []
    for face in shape.Faces:
        try:
            surf = face.Surface
            if not (surf and "plane" in surf.__class__.__name__.lower()):
                continue
            area = float(face.Area)
            if area < 1.0:  # mm², ignore tiny faces
                continue
            p, n = face_midpoint(face)
            n = norm(n)
            if n is None:
                continue
            n = canonical_normal(n)
            d = float(n.dot(p))  # plane offset for n·x = d
            planar.append((area, face, n, d))
        except Exception:
            continue

    if not planar:
        return None

    planar.sort(key=lambda x: x[0], reverse=True)
    planar = planar[:max_faces]

    # --- cluster by normal direction ---
    # Tolerance: within ~2 degrees
    cos_tol = math.cos(math.radians(2.0))

    clusters = []  # each: {"n": normal, "faces": [(area, face, d), ...]}
    for area, face, n, d in planar:
        placed = False
        for c in clusters:
            if angle_close(n, c["n"], cos_tol):
                c["faces"].append((area, face, d))
                placed = True
                break
        if not placed:
            clusters.append({"n": n, "faces": [(area, face, d)]})

    # --- build thickness candidates from d-values inside each cluster ---
    # For a sheet, same-normal planes occur at (outer) and (inner) offsets,
    # so distinct d-values differ ~thickness.
    candidates = []

    def add_candidate(val):
        if 0.05 <= val <= 20.0:  # mm range guard (tune if needed)
            candidates.append(val)

    for c in clusters:
        ds = [d for _a, _f, d in c["faces"]]
        if len(ds) < 2:
            continue
        ds.sort()

        # unique d-values with binning (0.01 mm)
        uniq = []
        for d in ds:
            b = round(d / 0.01) * 0.01
            if not uniq or abs(b - uniq[-1]) > 1e-9:
                uniq.append(b)

        if len(uniq) < 2:
            continue

        # candidate: smallest positive gap between uniq planes
        # Often thickness is the smallest meaningful separation.
        for i in range(1, len(uniq)):
            gap = abs(uniq[i] - uniq[i-1])
            add_candidate(gap)

    # --- if candidates exist, pick most frequent bin (mode-ish) ---
    def pick_mode(vals, bin_size=0.01):
        bins = {}
        for x in vals:
            b = round(x / bin_size) * bin_size
            bins.setdefault(b, []).append(x)
        best_bin = max(bins.items(), key=lambda kv: len(kv[1]))[0]
        bucket = sorted(bins[best_bin])
        return bucket[len(bucket)//2]

    if candidates:
        return pick_mode(candidates, 0.01)

    # --- Fallback: distToShape between face pairs in same normal cluster ---
    # Slower but can rescue cases where d-values are too noisy.
    # We try only top clusters and top faces.
    try:
        for c in clusters[:6]:
            faces = sorted(c["faces"], key=lambda t: t[0], reverse=True)[:8]
            # compare each face to others in same cluster; minimal non-zero distance tends to thickness
            for i in range(len(faces)):
                fi = faces[i][1]
                for j in range(i+1, len(faces)):
                    fj = faces[j][1]
                    dist = fi.distToShape(fj)[0]  # returns (dist, pts, info)
                    if dist and dist > 0.05 and dist <= 20.0:
                        candidates.append(float(dist))
        if candidates:
            return pick_mode(candidates, 0.01)
    except Exception:
        pass

    return None



# -----------------------------
# Main
# -----------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--material", required=True, help="alu | stainless | copper")
    parser.add_argument("--thickness-mm", required=False, type=float, default=None,
                        help="Optional sheet thickness in mm (auto-detect if omitted)")
    args = parser.parse_args()

    material_key = args.material.strip().lower()
    if material_key not in DENSITY_KG_M3:
        raise SystemExit(f"Unknown material '{args.material}'")

    density = DENSITY_KG_M3[material_key]

    cwd = os.getcwd()
    step_path = os.path.join(cwd, "03341791-01_01.step")
    out_dxf = os.path.join(cwd, "03341791-01_01.dxf")
    out_json = os.path.join(cwd, "03341791-01_01-result.json")
    out_fcstd = os.path.join(cwd, "debug_last.FCStd")

    print("=== STEPANALYSER START ===", flush=True)
    print("Material:", material_key, "Density:", density, flush=True)

    if not os.path.exists(step_path):
        raise SystemExit("STEP File not found not found in working directory")

    import FreeCAD as App
    import Import
    import importDXF

    try:
        import SheetMetalNewUnfolder
        from SheetMetalNewUnfolder import BendAllowanceCalculator
        HAS_V2 = True
    except Exception:
        HAS_V2 = False

    try:
        import SheetMetalUnfolder
        HAS_V1 = True
    except Exception:
        HAS_V1 = False

    if not HAS_V1 and not HAS_V2:
        raise SystemExit("No SheetMetal unfolder available")

    doc = App.newDocument("StepAnalyser")
    warnings = []

    try:
        Import.insert(step_path, doc.Name)
        doc.recompute()

        part_obj = pick_main_object(doc)
        face_idx, planar = pick_root_face_index(part_obj.Shape)
        base_face = f"Face{face_idx}"

        thickness_mm = args.thickness_mm
        if thickness_mm is None:
            print("Auto-detecting thickness...", flush=True)
            thickness_mm = estimate_thickness_mm_from_solid(part_obj)
            if thickness_mm is None:
                raise RuntimeError("Could not auto-detect thickness")
            print(f"Detected thickness: {thickness_mm:.3f} mm", flush=True)

        unfolded_shape = None
        sketches = []

        if HAS_V2:
            try:
                bac = BendAllowanceCalculator.from_single_value(K_FACTOR, K_STANDARD)
                sel_face, unfolded_shape, bend_lines, root_normal = \
                    SheetMetalNewUnfolder.getUnfold(bac, part_obj, base_face)
                sketches = SheetMetalNewUnfolder.getUnfoldSketches(
                    "Unfold", sel_face, unfolded_shape, bend_lines,
                    root_normal, [], False, "#000080", "#c00000", "#ff5733"
                )
            except Exception:
                warnings.append("V2 unfold failed")

        if unfolded_shape is None and HAS_V1:
            ktable = {1: K_FACTOR}
            unfolded_shape, foldComp, norm, *_ = \
                SheetMetalUnfolder.getUnfold(ktable, part_obj, base_face, K_STANDARD)
            sketches = SheetMetalUnfolder.getUnfoldSketches(
                "Unfold", unfolded_shape, foldComp.Edges,
                norm, [], False, "#000080",
                bendSketchColor="#ff5733", internalSketchColor="#c00000"
            )

        if unfolded_shape is None or not sketches:
            raise RuntimeError("Unfold failed")

        importDXF.export(sketches, out_dxf)

        bbox_w, bbox_h = compute_bbox_mm(unfolded_shape)
        area_bbox_mm2 = bbox_w * bbox_h
        area_net_mm2 = float(unfolded_shape.Area)

        t_m = mm_to_m(thickness_mm)
        area_bbox_m2 = mm2_to_m2(area_bbox_mm2)
        area_net_m2 = mm2_to_m2(area_net_mm2)

        mass_bbox_kg = area_bbox_m2 * t_m * density
        mass_net_kg = area_net_m2 * t_m * density

        result = {
            "ok": True,
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "input": {
                "material": material_key,
                "density_kg_m3": density,
                "thickness_mm": thickness_mm,
                "k_factor": K_FACTOR
            },
            "flat": {
                "bbox_mm": {"width": bbox_w, "height": bbox_h},
                "area_bbox_mm2": area_bbox_mm2,
                "area_net_mm2": area_net_mm2
            },
            "weight": {
                "bbox_kg": mass_bbox_kg,
                "net_kg": mass_net_kg
            },
            "warnings": warnings
        }

        with open(out_json, "w") as f:
            json.dump(result, f, indent=2)

        doc.saveAs(out_fcstd)

        print("OK: flat.dxf + result.json written", flush=True)

    except Exception as e:
        try:
            doc.saveAs(out_fcstd)
        except Exception:
            pass

        err = {
            "ok": False,
            "error": str(e),
            "traceback": traceback.format_exc()
        }
        with open(out_json, "w") as f:
            json.dump(err, f, indent=2)

        print("ERROR:", e, flush=True)
        print(traceback.format_exc(), flush=True)
        os._exit(1)

    os._exit(0)


if __name__ == "__main__":
    main()
