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

"""
Internal STEP sheet metal analyser

- Input:  --input <file.step|file.stp> (relative to cwd or absolute)
- Unfold with K-factor = 0.5
- Auto-detect thickness if not provided
- Export:
    - <basename>_flat.dxf
    - <basename>_result.json
    - <basename>_debug.FCStd
  Additionally (compat):
    - flat.dxf
    - result.json
    - debug_last.FCStd
"""

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

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,
}


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


def mm_to_m(x: float) -> float:
    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:
    - cluster planar faces by normal direction
    - estimate thickness from plane offsets
    - fallback to distToShape on face pairs
    """
    import math

    shape = part_obj.Shape

    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):
        eps = 1e-9
        x, y, z = float(n.x), float(n.y), float(n.z)
        if abs(x) > eps:
            return n if x > 0 else n.multiply(-1)
        if abs(y) > eps:
            return n if y > 0 else n.multiply(-1)
        if abs(z) > eps:
            return n if z > 0 else n.multiply(-1)
        return n

    def angle_close(n1, n2, cos_tol):
        return (n1.dot(n2) >= cos_tol)

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

    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:
                continue
            p, n = face_midpoint(face)
            n = norm(n)
            if n is None:
                continue
            n = canonical_normal(n)
            d = float(n.dot(p))
            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]

    cos_tol = math.cos(math.radians(2.0))

    clusters = []
    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)]})

    candidates = []

    def add_candidate(val):
        if 0.05 <= val <= 20.0:
            candidates.append(val)

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

        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

        for i in range(1, len(uniq)):
            gap = abs(uniq[i] - uniq[i - 1])
            add_candidate(gap)

    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 (slower)
    try:
        for c in clusters[:6]:
            faces = sorted(c["faces"], key=lambda t: t[0], reverse=True)[:8]
            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]
                    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


def write_json(path, payload):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)


def try_copy(src, dst):
    try:
        if src != dst and os.path.exists(src):
            # overwrite
            with open(src, "rb") as fsrc:
                data = fsrc.read()
            with open(dst, "wb") as fdst:
                fdst.write(data)
    except Exception:
        pass


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="STEP file path (.step/.stp), absolute or relative to cwd")
    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 = args.input
    if not os.path.isabs(step_path):
        step_path = os.path.join(cwd, step_path)

    if not os.path.exists(step_path):
        raise SystemExit(f"STEP file not found in working directory: {step_path}")

    base = os.path.splitext(os.path.basename(step_path))[0]

    out_dxf_named = os.path.join(cwd, f"{base}_flat.dxf")
    out_json_named = os.path.join(cwd, f"{base}_result.json")
    out_fcstd_named = os.path.join(cwd, f"{base}_debug.FCStd")

    # compat outputs for the web UI
    out_dxf = os.path.join(cwd, "flat.dxf")
    out_json = os.path.join(cwd, "result.json")
    out_fcstd = os.path.join(cwd, "debug_last.FCStd")

    print("=== STEPANALYSER START ===", flush=True)
    print("Input:", os.path.basename(step_path), flush=True)
    print("Material:", material_key, "Density:", density, flush=True)

    import FreeCAD as App
    import Import
    import importDXF

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

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

    if not HAS_V1 and not HAS_V2:
        raise SystemExit("No SheetMetal unfolder available (V1/V2). Check SheetMetal installation.")

    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)

        if thickness_mm <= 0:
            raise RuntimeError("Invalid thickness (<= 0)")

        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; trying V1 fallback.")

        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 (no unfolded shape or sketches).")

        # Export DXF (named)
        importDXF.export(sketches, out_dxf_named)

        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": {
                "step_file": os.path.abspath(step_path),
                "step_filename": os.path.basename(step_path),
                "material": material_key,
                "density_kg_m3": density,
                "thickness_mm": thickness_mm,
                "k_factor": K_FACTOR,
                "k_standard": K_STANDARD,
            },
            "flat": {
                "bbox_mm": {"width_mm": bbox_w, "height_mm": bbox_h},
                "area_bbox_mm2": area_bbox_mm2,
                "area_bbox_m2": area_bbox_m2,
                "area_net_mm2": area_net_mm2,
                "area_net_m2": area_net_m2,
            },
            "weight": {
                "bbox_kg": mass_bbox_kg,
                "net_kg": mass_net_kg,
                "bbox_g": mass_bbox_kg * 1000.0,
                "net_g": mass_net_kg * 1000.0,
            },
            "output": {
                "dxf_named": os.path.abspath(out_dxf_named),
                "json_named": os.path.abspath(out_json_named),
                "fcstd_named": os.path.abspath(out_fcstd_named),
                "dxf": os.path.abspath(out_dxf),
                "json": os.path.abspath(out_json),
                "fcstd": os.path.abspath(out_fcstd),
            },
            "warnings": warnings,
        }

        write_json(out_json_named, result)

        # Save debug doc (named)
        doc.saveAs(out_fcstd_named)

        # Compat copies for web UI
        try_copy(out_dxf_named, out_dxf)
        try_copy(out_json_named, out_json)
        try_copy(out_fcstd_named, out_fcstd)

        print("OK: wrote DXF + JSON", flush=True)

    except Exception as e:
        # Always write named + compat error JSON
        err = {
            "ok": False,
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "error_type": type(e).__name__,
            "error": str(e),
            "traceback": traceback.format_exc(),
            "input": {
                "step_file": os.path.abspath(step_path),
                "step_filename": os.path.basename(step_path),
                "material": material_key,
                "density_kg_m3": density,
                "thickness_mm": args.thickness_mm,
                "k_factor": K_FACTOR,
                "k_standard": K_STANDARD,
            },
        }

        try:
            write_json(out_json_named, err)
        except Exception:
            pass

        try:
            write_json(out_json, err)
        except Exception:
            pass

        try:
            doc.saveAs(out_fcstd_named)
            try_copy(out_fcstd_named, out_fcstd)
        except Exception:
            pass

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

    os._exit(0)


if __name__ == "__main__":
    # Catch also SystemExit/argparse exits so we still emit result.json
    try:
        main()
    except BaseException as e:
        cwd = os.getcwd()
        err = {
            "ok": False,
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "error_type": type(e).__name__,
            "error": str(e),
            "traceback": traceback.format_exc(),
        }
        try:
            with open(os.path.join(cwd, "result.json"), "w", encoding="utf-8") as f:
                json.dump(err, f, indent=2, ensure_ascii=False)
        except Exception:
            pass
        print("FATAL:", err["error_type"], err["error"], flush=True)
        os._exit(1)