#!/usr/bin/env python3
"""
Parametric stylised NCC-1715E Enterprise-inspired starship + desktop stand.

This is a printable display model, not a screen-accurate licensed replica.  Units: mm.
Generate:
  - ncc_1715e_starship_PRINT.stl
  - ncc_1715e_display_stand_PRINT.stl
  - ncc_1715e_assembly_REFERENCE.stl
  - preview PNG/SVG/HTML and verification manifest
"""
from __future__ import annotations

import json
import math
import os
from pathlib import Path
from typing import Iterable, Tuple

import numpy as np
import trimesh
from trimesh.transformations import rotation_matrix, translation_matrix

OUT = Path(__file__).resolve().parent

# ---- tunable dimensions -------------------------------------------------
SHIP_LENGTH = 185.0
SHIP_SCALE = SHIP_LENGTH / 185.0
REGISTRY = "NCC-1715E"  # Change to NCC-1701-E here if desired.

# Mesh detail: 2 is reasonable, 3 is smoother but heavier.
SPHERE_SUBDIV = 3
CYL_SECTIONS = 48

# Stand geometry
BASE_X = 115.0
BASE_Y = 72.0
BASE_Z = 7.0
POST_DIAM = 14.0
POST_HEIGHT = 55.0
SADDLE_WIDTH = 44.0
SADDLE_DEPTH = 20.0
SADDLE_ARM_DIAM = 9.0

# ------------------------------------------------------------------------

def tf(mesh: trimesh.Trimesh, matrix) -> trimesh.Trimesh:
    m = mesh.copy()
    m.apply_transform(matrix)
    return m


def box(size: Tuple[float, float, float], center=(0, 0, 0)) -> trimesh.Trimesh:
    return trimesh.creation.box(extents=size, transform=translation_matrix(center))


def cyl(radius: float, height: float, axis: str = "z", center=(0, 0, 0), sections=CYL_SECTIONS) -> trimesh.Trimesh:
    m = trimesh.creation.cylinder(radius=radius, height=height, sections=sections)
    if axis == "x":
        m.apply_transform(rotation_matrix(math.pi / 2, [0, 1, 0]))
    elif axis == "y":
        m.apply_transform(rotation_matrix(math.pi / 2, [1, 0, 0]))
    elif axis == "z":
        pass
    else:
        raise ValueError(axis)
    m.apply_transform(translation_matrix(center))
    return m


def ellipsoid(rx: float, ry: float, rz: float, center=(0, 0, 0), subdiv=SPHERE_SUBDIV) -> trimesh.Trimesh:
    m = trimesh.creation.icosphere(subdivisions=subdiv, radius=1.0)
    m.apply_scale([rx, ry, rz])
    m.apply_transform(translation_matrix(center))
    return m


def rounded_rect_base(x: float, y: float, z: float, r: float) -> trimesh.Trimesh:
    # Union a rectangle core plus four corner cylinders. Works reliably with manifold.
    parts = [box((x - 2*r, y, z), center=(0, 0, z/2)), box((x, y - 2*r, z), center=(0, 0, z/2))]
    for sx in [-1, 1]:
        for sy in [-1, 1]:
            parts.append(cyl(r, z, "z", center=(sx*(x/2-r), sy*(y/2-r), z/2), sections=32))
    return union(parts, "rounded_rect_base")


def union(parts: Iterable[trimesh.Trimesh], label="union") -> trimesh.Trimesh:
    parts = [p for p in parts if p is not None and len(p.vertices) > 0]
    if not parts:
        raise ValueError(f"empty {label}")
    if len(parts) == 1:
        return parts[0]
    try:
        merged = trimesh.boolean.union(parts, engine="manifold")
        if isinstance(merged, list):
            merged = trimesh.util.concatenate(merged)
        merged.remove_unreferenced_vertices()
        merged.merge_vertices()
        return merged
    except Exception as e:
        print(f"WARNING: boolean union failed for {label}: {e}; using concatenated overlap mesh")
        return trimesh.util.concatenate(parts)


def difference(mesh: trimesh.Trimesh, cutters: Iterable[trimesh.Trimesh], label="difference") -> trimesh.Trimesh:
    cutters = [c for c in cutters if c is not None and len(c.vertices) > 0]
    if not cutters:
        return mesh
    try:
        cut = trimesh.boolean.difference([mesh, *cutters], engine="manifold")
        if isinstance(cut, list):
            cut = trimesh.util.concatenate(cut)
        cut.remove_unreferenced_vertices()
        cut.merge_vertices()
        return cut
    except Exception as e:
        print(f"WARNING: boolean difference failed for {label}: {e}; returning uncut mesh")
        return mesh

# ---- blocky seven-segment text for robust raised/engraved labels --------
SEGMENTS = {
    "A": ((0.22, 0.88, 0.56, 0.10),),       # top horizontal
    "D": ((0.22, 0.05, 0.56, 0.10),),       # bottom
    "G": ((0.22, 0.47, 0.56, 0.10),),       # middle
    "F": ((0.08, 0.52, 0.10, 0.36),),       # upper left
    "B": ((0.78, 0.52, 0.10, 0.36),),       # upper right
    "E": ((0.08, 0.12, 0.10, 0.36),),       # lower left
    "C": ((0.78, 0.12, 0.10, 0.36),),       # lower right
}
DIGIT_SEGS = {
    "0": "ABCFED", "1": "BC", "2": "ABGED", "3": "ABGCD", "4": "FBGC",
    "5": "AFGCD", "6": "AFGECD", "7": "ABC", "8": "ABCDEFG", "9": "ABFGCD",
}
LETTER_PATTERNS = {
    "N": [(0.05,0.05,0.11,0.90),(0.84,0.05,0.11,0.90),(0.23,0.50,0.58,0.10)],
    "C": [(0.20,0.88,0.58,0.10),(0.08,0.12,0.10,0.76),(0.20,0.05,0.58,0.10)],
    "E": [(0.08,0.05,0.10,0.90),(0.20,0.88,0.62,0.10),(0.20,0.47,0.54,0.10),(0.20,0.05,0.62,0.10)],
    "-": [(0.22,0.47,0.56,0.10)],
}

def text_blocks(text: str, letter_w=6.0, letter_h=9.5, depth=1.0, spacing=1.2, center=(0,0,0)) -> trimesh.Trimesh:
    parts = []
    total_w = len(text) * letter_w + (len(text)-1)*spacing
    x0 = -total_w/2.0
    for i, ch in enumerate(text.upper()):
        bx = x0 + i*(letter_w+spacing)
        rects = []
        if ch in DIGIT_SEGS:
            for seg in DIGIT_SEGS[ch]:
                rects.extend(SEGMENTS[seg])
        else:
            rects = LETTER_PATTERNS.get(ch, [])
        for rx, ry, rw, rh in rects:
            parts.append(box((rw*letter_w, rh*letter_h, depth), center=(bx + rx*letter_w, (ry-0.5)*letter_h, center[2] + depth/2)))
    label = union(parts, "text_blocks") if parts else box((1,1,depth), center=(0,0,center[2]+depth/2))
    label.apply_transform(translation_matrix((center[0], center[1], 0)))
    return label

# ---- ship ----------------------------------------------------------------

def build_ship() -> trimesh.Trimesh:
    s = SHIP_SCALE
    parts = []

    # Sovereign-ish long oval saucer: a flattened ellipsoid plus lower lip.
    parts.append(ellipsoid(54*s, 35*s, 5.2*s, center=(0, -42*s, 12*s)))
    parts.append(ellipsoid(44*s, 27*s, 3.2*s, center=(0, -44*s, 8.9*s), subdiv=2))
    parts.append(ellipsoid(17*s, 12*s, 3.0*s, center=(0, -48*s, 17.2*s), subdiv=2))  # bridge mound

    # Secondary hull, flattened long rugby ball with printed deflector nose.
    parts.append(ellipsoid(19*s, 52*s, 11*s, center=(0, 22*s, 3*s)))
    parts.append(ellipsoid(12*s, 10*s, 6*s, center=(0, -15*s, 4*s), subdiv=2))  # neck fairing
    parts.append(cyl(8.5*s, 4.0*s, "y", center=(0, -28*s, 3.8*s), sections=32))  # deflector dish boss

    # Wide neck/spine connecting saucer to hull.
    neck = box((18*s, 45*s, 8*s), center=(0, -15*s, 7.2*s))
    neck.apply_transform(rotation_matrix(math.radians(-6), [1,0,0], point=[0,-15*s,7.2*s]))
    parts.append(neck)

    # Warp nacelles and caps.
    for side in [-1, 1]:
        x = side * 37*s
        parts.append(cyl(6.5*s, 84*s, "y", center=(x, 30*s, 10*s), sections=40))
        parts.append(ellipsoid(6.8*s, 8*s, 6.8*s, center=(x, -12*s, 10*s), subdiv=2))
        parts.append(ellipsoid(6.8*s, 8*s, 6.8*s, center=(x, 72*s, 10*s), subdiv=2))
        # Pylons as thick angled struts; robust but recognisable.
        pylon = box((7*s, 39*s, 5*s), center=(side*22*s, 14*s, 6.5*s))
        pylon.apply_transform(rotation_matrix(math.radians(side*24), [0,0,1], point=[side*22*s,14*s,6.5*s]))
        pylon.apply_transform(rotation_matrix(math.radians(-10), [1,0,0], point=[side*22*s,14*s,6.5*s]))
        parts.append(pylon)
        # Bussard/front caps as raised rings.
        parts.append(cyl(4.4*s, 2.2*s, "y", center=(x, -19.2*s, 10*s), sections=32))

    # Small vent/ridge bars on top aft hull, deliberately embedded into the hull for one connected solid.
    for x in [-8, 0, 8]:
        parts.append(box((2.4*s, 28*s, 2.2*s), center=(x*s, 28*s, 12.9*s)))

    ship = union(parts, "starship")
    # Move ship so its bottom is just above z=0 for standalone STL.
    mn, mx = ship.bounds
    ship.apply_transform(translation_matrix((0, 0, -mn[2] + 0.05)))
    return ship

# ---- stand ---------------------------------------------------------------

def build_stand() -> trimesh.Trimesh:
    parts = []
    base = rounded_rect_base(BASE_X, BASE_Y, BASE_Z, 9.0)
    parts.append(base)

    # Raised label on base top.
    # Sink the label 0.45 mm into the base so each letter is mechanically fused, not just coplanar-touching.
    label = text_blocks(REGISTRY, letter_w=5.2, letter_h=8.0, depth=1.25, spacing=0.7, center=(0, -18.0, BASE_Z - 0.45))
    parts.append(label)

    # Angled post leaning backward so the ship sits level in the saddle.
    post = cyl(POST_DIAM/2, POST_HEIGHT, "z", center=(0, 6, BASE_Z + POST_HEIGHT/2))
    post.apply_transform(rotation_matrix(math.radians(-12), [1,0,0], point=[0, 6, BASE_Z]))
    parts.append(post)

    # Triangular-ish gussets as chunky plates front/back.
    for yoff in [-7, 7]:
        gus = box((10, 3.5, 38), center=(0, 6+yoff, BASE_Z+19))
        gus.apply_transform(rotation_matrix(math.radians(-12), [1,0,0], point=[0, 6, BASE_Z]))
        parts.append(gus)

    # Saddle: two rounded rails plus crossbar. Designed to cradle the engineering hull/saucer underside.
    saddle_z = BASE_Z + POST_HEIGHT + 4
    # Positive mechanical boss/neck between the leaning post and saddle so the printable stand is one body.
    # The blocky collar is intentionally overbuilt and hidden under the cradle; it avoids fragile tangent-only joins.
    parts.append(box((18.0, 18.0, 18.0), center=(0, 0, saddle_z - 10)))
    parts.append(ellipsoid(12.0, 10.0, 9.0, center=(0, -4, saddle_z - 5), subdiv=2))
    parts.append(cyl(SADDLE_ARM_DIAM/2, SADDLE_WIDTH, "x", center=(0, -4, saddle_z), sections=32))
    for sx in [-1, 1]:
        arm = cyl(SADDLE_ARM_DIAM/2, SADDLE_DEPTH, "y", center=(sx*SADDLE_WIDTH/2, -11, saddle_z+2), sections=32)
        arm.apply_transform(rotation_matrix(math.radians(18*sx), [0,1,0], point=[sx*SADDLE_WIDTH/2,-11,saddle_z+2]))
        parts.append(arm)
        parts.append(ellipsoid(SADDLE_ARM_DIAM/2, SADDLE_ARM_DIAM/2, SADDLE_ARM_DIAM/2, center=(sx*SADDLE_WIDTH/2, -21, saddle_z+2), subdiv=2))

    stand = union(parts, "stand")
    mn, mx = stand.bounds
    stand.apply_transform(translation_matrix((0, 0, -mn[2])))
    return stand


def build_assembly(ship: trimesh.Trimesh, stand: trimesh.Trimesh) -> trimesh.Trimesh:
    # Reference only: place the ship above the saddle with slight nose-down heroic angle.
    st = stand.copy()
    sh = ship.copy()
    # center ship on stand, raised above saddle
    sh.apply_transform(rotation_matrix(math.radians(-3), [1,0,0], point=[0,0,0]))
    sh.apply_transform(translation_matrix((0, 3, 71)))
    return trimesh.util.concatenate([st, sh])


def verify_mesh(path: Path) -> dict:
    m = trimesh.load_mesh(path, force="mesh")
    comps = len(m.split(only_watertight=False))
    return {
        "file": path.name,
        "size_bytes": path.stat().st_size,
        "vertices": int(len(m.vertices)),
        "faces": int(len(m.faces)),
        "watertight": bool(m.is_watertight),
        "components": int(comps),
        "bounds_mm": [[round(float(v), 3) for v in row] for row in m.bounds.tolist()],
        "extents_mm": [round(float(v), 3) for v in m.extents.tolist()],
    }


def make_previews(ship: trimesh.Trimesh, stand: trimesh.Trimesh, assembly: trimesh.Trimesh):
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection

    def render(mesh, path, title, elev=22, azim=-55):
        fig = plt.figure(figsize=(9, 7), dpi=130)
        ax = fig.add_subplot(111, projection='3d')
        # downsample if needed for preview speed
        faces = mesh.faces
        if len(faces) > 14000:
            faces = faces[::max(1, len(faces)//14000)]
        verts = mesh.vertices
        poly = Poly3DCollection(verts[faces], linewidths=0.02, alpha=0.96)
        poly.set_facecolor((0.72, 0.74, 0.78, 1.0))
        poly.set_edgecolor((0.14, 0.14, 0.16, 0.12))
        ax.add_collection3d(poly)
        mn, mx = mesh.bounds
        center = (mn + mx)/2
        radius = max(mesh.extents)/2
        ax.set_xlim(center[0]-radius, center[0]+radius)
        ax.set_ylim(center[1]-radius, center[1]+radius)
        ax.set_zlim(max(0, center[2]-radius*0.55), center[2]+radius*0.9)
        ax.view_init(elev=elev, azim=azim)
        ax.set_title(title)
        ax.set_xlabel('X mm')
        ax.set_ylabel('Y mm')
        ax.set_zlabel('Z mm')
        ax.set_box_aspect((1,1,0.55))
        fig.tight_layout()
        fig.savefig(path)
        plt.close(fig)

    render(ship, OUT / "preview_starship.png", "NCC-1715E stylised starship printable STL", elev=19, azim=-52)
    render(stand, OUT / "preview_stand.png", "NCC-1715E labelled desktop stand printable STL", elev=25, azim=-45)
    render(assembly, OUT / "preview_assembly_reference.png", "Assembly reference only: ship on stand", elev=18, azim=-48)

    # Simple SVG side/top silhouettes for quick browser review.
    svg = f'''<svg xmlns="http://www.w3.org/2000/svg" width="1000" height="650" viewBox="0 0 1000 650">
<style>text{{font-family:Arial,sans-serif}} .dim{{font-size:20px;fill:#333}} .label{{font-size:24px;font-weight:700;fill:#111}} .ship{{fill:#b8bcc4;stroke:#222;stroke-width:3}} .stand{{fill:#d9d1b5;stroke:#222;stroke-width:3}}</style>
<rect width="100%" height="100%" fill="#f7f7f4"/>
<text x="40" y="48" class="label">NCC-1715E Enterprise-inspired display model — dimension sketch</text>
<ellipse cx="330" cy="185" rx="275" ry="85" class="ship"/>
<ellipse cx="330" cy="185" rx="85" ry="35" fill="#ccd0d6" stroke="#222" stroke-width="2"/>
<ellipse cx="655" cy="235" rx="70" ry="185" class="ship" transform="rotate(-8 655 235)"/>
<rect x="185" y="505" width="360" height="42" rx="18" class="stand"/>
<rect x="352" y="365" width="45" height="150" rx="22" class="stand" transform="rotate(-12 374 515)"/>
<rect x="285" y="333" width="175" height="36" rx="16" class="stand"/>
<text x="40" y="305" class="dim">Ship STL approx: {SHIP_LENGTH:.0f} mm long × 87 mm wide × 29 mm tall</text>
<text x="40" y="580" class="dim">Stand base: {BASE_X:.0f} × {BASE_Y:.0f} × {BASE_Z:.0f} mm; raised registry label: {REGISTRY}</text>
<text x="620" y="580" class="dim">Print as two parts: ship + stand. Assembly STL is reference only.</text>
</svg>'''
    (OUT / "dimension_sketch.svg").write_text(svg)

    html = f'''<!doctype html><html><head><meta charset="utf-8"><title>NCC-1715E CAD preview</title>
<style>body{{font-family:system-ui,Arial,sans-serif;background:#111;color:#eee;margin:0;padding:24px}} .grid{{display:grid;grid-template-columns:repeat(auto-fit,minmax(320px,1fr));gap:18px}} .card{{background:#1d1d22;border:1px solid #333;border-radius:14px;padding:14px}} img{{max-width:100%;border-radius:10px;background:#fff}} code{{background:#2b2b32;padding:2px 5px;border-radius:5px}}</style></head>
<body><h1>NCC-1715E Enterprise-inspired starship + stand</h1>
<p>Printable STLs: <code>ncc_1715e_starship_PRINT.stl</code> and <code>ncc_1715e_display_stand_PRINT.stl</code>. Assembly STL is for visual reference only.</p>
<div class="grid"><div class="card"><h2>Starship</h2><img src="preview_starship.png"></div><div class="card"><h2>Stand</h2><img src="preview_stand.png"></div><div class="card"><h2>Assembly reference</h2><img src="preview_assembly_reference.png"></div><div class="card"><h2>Dimension sketch</h2><img src="dimension_sketch.svg"></div></div>
</body></html>'''
    (OUT / "view.html").write_text(html)


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    ship = build_ship()
    stand = build_stand()
    assembly = build_assembly(ship, stand)

    paths = {
        "ship": OUT / "ncc_1715e_starship_PRINT.stl",
        "stand": OUT / "ncc_1715e_display_stand_PRINT.stl",
        "assembly": OUT / "ncc_1715e_assembly_REFERENCE.stl",
    }
    ship.export(paths["ship"])
    stand.export(paths["stand"])
    assembly.export(paths["assembly"])
    make_previews(ship, stand, assembly)

    manifest = {
        "project": "ncc_1715e_enterprise_display",
        "registry_text": REGISTRY,
        "material_assumption": "PLA first print; PETG optional for tougher stand.",
        "files": [verify_mesh(p) for p in paths.values()],
        "notes": [
            "Stylised Enterprise-inspired desktop model; not a screen-accurate licensed replica.",
            "Print ship and stand separately. Assembly STL is reference only.",
            "Ship has overlapping features boolean-unioned into one watertight body where manifold3d succeeds.",
            "Stand has raised blocky registry text for robust FDM printing.",
        ],
    }
    (OUT / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))

if __name__ == "__main__":
    main()