import hashlib
import json
import os
import struct
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from OCP.BRepMesh import BRepMesh_IncrementalMesh
from OCP.Message import Message_ProgressRange
from OCP.RWGltf import RWGltf_CafWriter
from OCP.STEPCAFControl import STEPCAFControl_Reader
from OCP.TColStd import TColStd_IndexedDataMapOfStringString
from OCP.TCollection import TCollection_AsciiString, TCollection_ExtendedString
from OCP.TDF import TDF_LabelSequence
from OCP.TDocStd import TDocStd_Document
from OCP.XCAFApp import XCAFApp_Application
from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ShapeTool


CONVERTER_NAME = "NodeDcBimConverter"
CONVERTER_VERSION = "0.3.0"
UPLOADS_DIR = Path(os.environ.get("NODEDC_BIM_CONVERTER_UPLOADS_DIR", "/beam/uploads")).resolve()
POLL_INTERVAL_SECONDS = float(os.environ.get("NODEDC_BIM_CONVERTER_INTERVAL_SECONDS", "10"))
PROCESS_ONCE = os.environ.get("NODEDC_BIM_CONVERTER_ONCE", "").lower() in {"1", "true", "yes"}
REPROCESS_READY_ON_VERSION_CHANGE = os.environ.get(
    "NODEDC_BIM_CONVERTER_REPROCESS_READY_ON_VERSION_CHANGE", ""
).lower() in {"1", "true", "yes"}
STEP_MESH_LINEAR_DEFLECTION = float(os.environ.get("NODEDC_BIM_CONVERTER_STEP_LINEAR_DEFLECTION", "0.1"))
STEP_MESH_ANGULAR_DEFLECTION = float(os.environ.get("NODEDC_BIM_CONVERTER_STEP_ANGULAR_DEFLECTION", "0.1"))
STEP_EXTENSIONS = {".step", ".stp"}


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def src_from_path(path: Path) -> str:
    return f"uploads/{path.resolve().relative_to(UPLOADS_DIR).as_posix()}"


def manifest_path(source_path: Path) -> Path:
    return source_path.with_name(f"{source_path.name}.beam.json")


def read_json(path: Path) -> dict[str, Any]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}


def write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(f"{path.name}.tmp")
    tmp_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    tmp_path.replace(path)


def hash_id(value: str) -> str:
    return hashlib.sha1(value.encode("utf-8")).hexdigest()[:16]


def json_safe(value: Any) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, dict):
        return {str(k): json_safe(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_safe(v) for v in value]
    return str(value)


def node_to_metadata(node: Any, parent_path: str = "") -> dict[str, Any]:
    name = node.name or "root"
    node_path = f"{parent_path}/{name}" if parent_path else name
    children = [node_to_metadata(child, node_path) for child in node.children]
    return {
        "id": hash_id(node_path),
        "name": name,
        "path": node_path,
        "hasShape": bool(node.obj),
        "metadata": json_safe(node.metadata or {}),
        "children": children,
    }


def count_nodes(node: dict[str, Any]) -> int:
    return 1 + sum(count_nodes(child) for child in node.get("children", []))


def metadata_dict(value: Any) -> dict[str, Any]:
    safe_value = json_safe(value or {})
    return safe_value if isinstance(safe_value, dict) else {"value": safe_value}


def get_shape_type(value: Any) -> str | None:
    try:
        shape_type = value.ShapeType()
        return str(shape_type) if shape_type else None
    except Exception:
        return None


def collect_solids(workplane: Any) -> list[Any]:
    solids: list[Any] = []
    seen: set[int] = set()

    for value in workplane.vals():
        candidates = []
        try:
            candidates = list(value.Solids())
        except Exception:
            candidates = []
        if not candidates and get_shape_type(value) == "Solid":
            candidates = [value]
        for solid in candidates:
            key = hash(solid)
            if key in seen:
                continue
            seen.add(key)
            solids.append(solid)

    return solids


def unique_assembly_name(parent: Any, requested_name: Any) -> str:
    base_name = str(requested_name or "component")
    if base_name not in parent.objects:
        return base_name

    index = 2
    while True:
        candidate = f"{base_name}__{index:03d}"
        if candidate not in parent.objects:
            return candidate
        index += 1


def import_step_with_unique_component_names(source_path: Path, cq_module: Any) -> Any:
    original_add = cq_module.Assembly.add

    def add_with_unique_names(self: Any, arg: Any, *args: Any, **kwargs: Any) -> Any:
        if isinstance(arg, cq_module.Assembly):
            positional_name = args[0] if args else None
            requested_name = kwargs.get("name", positional_name) or arg.name
            unique_name = unique_assembly_name(self, requested_name)

            if unique_name != requested_name:
                arg.metadata = {
                    **metadata_dict(arg.metadata),
                    "originalName": str(requested_name),
                }
                if args:
                    args = (unique_name, *args[1:])
                else:
                    kwargs = {**kwargs, "name": unique_name}

        return original_add(self, arg, *args, **kwargs)

    cq_module.Assembly.add = add_with_unique_names
    try:
        return cq_module.Assembly.importStep(str(source_path))
    finally:
        cq_module.Assembly.add = original_add


def import_step_assembly(source_path: Path) -> tuple[Any, str, str]:
    import cadquery as cq

    try:
        return cq.Assembly.importStep(str(source_path)), "assembly", getattr(cq, "__version__", "unknown")
    except ValueError as exc:
        error_message = str(exc)
        if "Unique name is required" in error_message and "already in the assembly" in error_message:
            return (
                import_step_with_unique_component_names(source_path, cq),
                "assembly-unique-names",
                getattr(cq, "__version__", "unknown"),
            )
        if "does not contain an assembly" not in error_message:
            raise

    step_model = cq.importers.importStep(str(source_path))
    solids = collect_solids(step_model)
    assy = cq.Assembly(name=f"{source_path.stem}_root")

    if len(solids) > 1:
        for index, solid in enumerate(solids, start=1):
            assy.add(solid, name=f"{source_path.stem}_solid_{index:03d}")
        return assy, "split-solids", getattr(cq, "__version__", "unknown")

    assy.add(step_model, name=source_path.stem)
    return assy, "single-shape", getattr(cq, "__version__", "unknown")


def read_glb_json(glb_path: Path) -> tuple[dict[str, Any], list[tuple[bytes, bytes]]]:
    data = glb_path.read_bytes()
    if len(data) < 20:
        raise ValueError("GLB is too small.")

    magic, version, declared_length = struct.unpack("<4sII", data[:12])
    if magic != b"glTF" or version != 2:
        raise ValueError("Only binary glTF v2 files are supported.")
    if declared_length != len(data):
        raise ValueError("GLB declared length does not match file size.")

    gltf_json: dict[str, Any] | None = None
    chunks: list[tuple[bytes, bytes]] = []
    offset = 12
    while offset < len(data):
        if offset + 8 > len(data):
            raise ValueError("Invalid GLB chunk header.")
        chunk_length, chunk_type = struct.unpack("<I4s", data[offset:offset + 8])
        offset += 8
        chunk_data = data[offset:offset + chunk_length]
        offset += chunk_length
        if chunk_type == b"JSON":
            gltf_json = json.loads(chunk_data.decode("utf-8"))
        else:
            chunks.append((chunk_type, chunk_data))

    if gltf_json is None:
        raise ValueError("GLB JSON chunk is missing.")
    return gltf_json, chunks


def write_glb_json(glb_path: Path, gltf_json: dict[str, Any], chunks: list[tuple[bytes, bytes]]) -> None:
    json_bytes = json.dumps(gltf_json, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    json_padding = (4 - (len(json_bytes) % 4)) % 4
    json_chunk = json_bytes + (b" " * json_padding)

    body = struct.pack("<I4s", len(json_chunk), b"JSON") + json_chunk
    for chunk_type, chunk_data in chunks:
        chunk_padding = (4 - (len(chunk_data) % 4)) % 4
        padded_chunk = chunk_data + (b"\x00" * chunk_padding)
        body += struct.pack("<I4s", len(padded_chunk), chunk_type) + padded_chunk

    glb_path.write_bytes(struct.pack("<4sII", b"glTF", 2, 12 + len(body)) + body)


def normalize_glb_node_names(glb_path: Path) -> tuple[dict[str, Any], int]:
    gltf_json, chunks = read_glb_json(glb_path)
    nodes = gltf_json.get("nodes", [])
    meshes = gltf_json.get("meshes", [])
    renamed = 0

    for node in nodes:
        mesh_index = node.get("mesh")
        if not isinstance(mesh_index, int) or mesh_index < 0 or mesh_index >= len(meshes):
            continue
        mesh_name = meshes[mesh_index].get("name")
        node_name = node.get("name")
        if mesh_name and (not node_name or str(node_name).startswith("NAUO")):
            node["name"] = mesh_name
            renamed += 1

    if renamed:
        write_glb_json(glb_path, gltf_json, chunks)
    return gltf_json, renamed


def gltf_node_name(gltf_json: dict[str, Any], node: dict[str, Any], index: int) -> str:
    name = node.get("name")
    if name:
        return str(name)
    meshes = gltf_json.get("meshes", [])
    mesh_index = node.get("mesh")
    if isinstance(mesh_index, int) and 0 <= mesh_index < len(meshes):
        mesh_name = meshes[mesh_index].get("name")
        if mesh_name:
            return str(mesh_name)
    return f"node_{index}"


def gltf_node_to_metadata(gltf_json: dict[str, Any], node_index: int, parent_path: str = "") -> dict[str, Any]:
    nodes = gltf_json.get("nodes", [])
    node = nodes[node_index]
    name = gltf_node_name(gltf_json, node, node_index)
    node_path = f"{parent_path}/{name}#{node_index}" if parent_path else f"{name}#{node_index}"
    children = [
        gltf_node_to_metadata(gltf_json, child_index, node_path)
        for child_index in node.get("children", [])
        if isinstance(child_index, int) and 0 <= child_index < len(nodes)
    ]
    return {
        "id": hash_id(node_path),
        "name": name,
        "path": node_path,
        "hasShape": isinstance(node.get("mesh"), int),
        "metadata": {},
        "children": children,
    }


def gltf_component_tree(gltf_json: dict[str, Any], fallback_name: str) -> dict[str, Any]:
    nodes = gltf_json.get("nodes", [])
    scenes = gltf_json.get("scenes", [])
    scene_index = gltf_json.get("scene", 0)
    scene = scenes[scene_index] if isinstance(scene_index, int) and 0 <= scene_index < len(scenes) else {}
    root_indexes = [idx for idx in scene.get("nodes", []) if isinstance(idx, int) and 0 <= idx < len(nodes)]

    if len(root_indexes) == 1:
        return gltf_node_to_metadata(gltf_json, root_indexes[0])

    children = [gltf_node_to_metadata(gltf_json, idx, fallback_name) for idx in root_indexes]
    return {
        "id": hash_id(fallback_name),
        "name": fallback_name,
        "path": fallback_name,
        "hasShape": False,
        "metadata": {},
        "children": children,
    }


def convert_step_xcaf_to_glb(source_path: Path, glb_path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("MDTV-XCAF"))
    app.NewDocument(TCollection_ExtendedString("MDTV-XCAF"), doc)

    print(f"[{CONVERTER_NAME}] STEPCAF read {source_path.name}", flush=True)
    reader = STEPCAFControl_Reader()
    reader.SetNameMode(True)
    reader.SetColorMode(True)
    reader.SetLayerMode(True)
    read_status = reader.ReadFile(str(source_path))
    if "RetDone" not in str(read_status):
        raise ValueError(f"STEPCAF read failed: {read_status}")
    if not reader.Transfer(doc):
        raise ValueError("STEPCAF transfer failed.")

    shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
    labels = TDF_LabelSequence()
    shape_tool.GetFreeShapes(labels)
    if labels.Length() == 0:
        raise ValueError("STEPCAF document has no free shapes.")

    print(
        f"[{CONVERTER_NAME}] meshing {source_path.name}: "
        f"{labels.Length()} free shape(s), deflection={STEP_MESH_LINEAR_DEFLECTION}",
        flush=True,
    )
    for index in range(1, labels.Length() + 1):
        shape = XCAFDoc_ShapeTool.GetShape_s(labels.Value(index))
        if not shape.IsNull():
            BRepMesh_IncrementalMesh(
                shape,
                STEP_MESH_LINEAR_DEFLECTION,
                False,
                STEP_MESH_ANGULAR_DEFLECTION,
                True,
            )

    glb_path.parent.mkdir(parents=True, exist_ok=True)
    print(f"[{CONVERTER_NAME}] GLB export {source_path.name}", flush=True)
    writer = RWGltf_CafWriter(TCollection_AsciiString(str(glb_path)), True)
    file_info = TColStd_IndexedDataMapOfStringString()
    if not writer.Perform(doc, file_info, Message_ProgressRange()):
        raise ValueError("GLB export failed.")
    if not glb_path.exists() or glb_path.stat().st_size == 0:
        raise ValueError("GLB export produced no output.")

    gltf_json, renamed_nodes = normalize_glb_node_names(glb_path)
    tree = gltf_component_tree(gltf_json, source_path.stem)
    stats = {
        "freeShapeCount": labels.Length(),
        "gltfNodeCount": len(gltf_json.get("nodes", [])),
        "gltfMeshCount": len(gltf_json.get("meshes", [])),
        "renamedNodeCount": renamed_nodes,
        "linearDeflection": STEP_MESH_LINEAR_DEFLECTION,
        "angularDeflection": STEP_MESH_ANGULAR_DEFLECTION,
    }
    return tree, stats


def convert_step_cadquery_to_glb(source_path: Path, glb_path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    assy, import_strategy, cadquery_version = import_step_assembly(source_path)
    tree = node_to_metadata(assy)
    glb_path.parent.mkdir(parents=True, exist_ok=True)
    assy.export(str(glb_path), tolerance=0.1, angularTolerance=0.1)
    stats = {
        "fallback": True,
        "importStrategy": import_strategy,
        "cadqueryVersion": cadquery_version,
    }
    return tree, stats


def convert_step_to_glb(source_path: Path, glb_path: Path, metadata_path: Path) -> dict[str, Any]:
    import_strategy = "occt-xcaf-gltf"
    try:
        tree, stats = convert_step_xcaf_to_glb(source_path, glb_path)
    except Exception as exc:
        print(f"[{CONVERTER_NAME}] OCCT/XCAF path failed for {source_path}: {exc}", file=sys.stderr, flush=True)
        tree, stats = convert_step_cadquery_to_glb(source_path, glb_path)
        import_strategy = stats.get("importStrategy") or "cadquery"

    metadata = {
        "source": src_from_path(source_path),
        "artifact": src_from_path(glb_path),
        "format": "step",
        "targetFormat": "glb",
        "importStrategy": import_strategy,
        "componentTree": tree,
        "componentCount": count_nodes(tree),
        "generatedAt": now_iso(),
        "converter": {
            "name": CONVERTER_NAME,
            "version": CONVERTER_VERSION,
            "engine": "occt-xcaf" if import_strategy == "occt-xcaf-gltf" else "cadquery",
            "cadqueryVersion": stats.get("cadqueryVersion"),
        },
        "stats": stats,
    }
    write_json_atomic(metadata_path, metadata)
    return metadata


def should_process(source_path: Path, manifest: dict[str, Any], glb_path: Path) -> bool:
    status = manifest.get("status")
    if status == "ready" and glb_path.exists():
        return REPROCESS_READY_ON_VERSION_CHANGE and manifest.get("converterVersion") != CONVERTER_VERSION
    if status == "processing":
        updated_at = manifest.get("updatedAt")
        if updated_at:
            try:
                updated = datetime.fromisoformat(str(updated_at).replace("Z", "+00:00"))
                if (datetime.now(timezone.utc) - updated).total_seconds() < 300:
                    return False
            except Exception:
                pass
    return source_path.exists()


def process_one(source_path: Path) -> bool:
    manifest_file = manifest_path(source_path)
    stem = source_path.stem
    glb_path = source_path.with_name(f"{stem}.glb")
    metadata_path = source_path.with_name(f"{stem}.metadata.json")
    manifest = read_json(manifest_file)

    if not should_process(source_path, manifest, glb_path):
        return False

    base_manifest = {
        "createdAt": manifest.get("createdAt") or now_iso(),
        "componentTreeRequired": True,
        "converterName": CONVERTER_NAME,
        "converterVersion": CONVERTER_VERSION,
        "sourceFormat": "step",
        "sourceSrc": src_from_path(source_path),
        "targetFormat": "glb",
    }

    write_json_atomic(
        manifest_file,
        {
            **base_manifest,
            "message": "Preparing GLB and component tree.",
            "status": "processing",
            "updatedAt": now_iso(),
        },
    )
    print(f"[{CONVERTER_NAME}] converting {source_path}", flush=True)

    try:
        metadata = convert_step_to_glb(source_path, glb_path, metadata_path)
        write_json_atomic(
            manifest_file,
            {
                **base_manifest,
                "artifactSrc": src_from_path(glb_path),
                "artifactType": "gltf",
                "componentCount": metadata.get("componentCount"),
                "metadataSrc": src_from_path(metadata_path),
                "message": "GLB and component tree are ready.",
                "status": "ready",
                "updatedAt": now_iso(),
            },
        )
        print(f"[{CONVERTER_NAME}] ready {glb_path}", flush=True)
        return True
    except Exception as exc:
        write_json_atomic(
            manifest_file,
            {
                **base_manifest,
                "error": str(exc),
                "message": "Failed to prepare GLB and component tree.",
                "status": "failed",
                "updatedAt": now_iso(),
            },
        )
        print(f"[{CONVERTER_NAME}] failed {source_path}: {exc}", file=sys.stderr, flush=True)
        return False


def scan_once() -> int:
    if not UPLOADS_DIR.exists():
        print(f"[{CONVERTER_NAME}] uploads dir does not exist: {UPLOADS_DIR}", flush=True)
        return 0

    processed = 0
    for source_path in sorted(UPLOADS_DIR.rglob("*")):
        if not source_path.is_file():
            continue
        if source_path.suffix.lower() not in STEP_EXTENSIONS:
            continue
        if process_one(source_path):
            processed += 1
    return processed


def main() -> None:
    print(f"[{CONVERTER_NAME}] watching {UPLOADS_DIR}", flush=True)
    while True:
        scan_once()
        if PROCESS_ONCE:
            return
        time.sleep(POLL_INTERVAL_SECONDS)


if __name__ == "__main__":
    main()