from __future__ import annotations

import io
import json
import math
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any

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


EPS = 1e-7


@dataclass
class ModelAnalysis:
    vertices: int
    faces: int
    watertight: bool
    volume_mm3: float | None
    extents_mm: list[float]
    centre_mm: list[float]
    principal_axis: list[float]
    suggested_scale_warning: str | None
    repair_notes: list[str]


def load_mesh(path_or_bytes: str | Path | bytes) -> trimesh.Trimesh:
    source: Any
    if isinstance(path_or_bytes, bytes):
        source = io.BytesIO(path_or_bytes)
    else:
        source = str(path_or_bytes)
    loaded = trimesh.load(source, file_type="stl", force="mesh", process=False)
    if isinstance(loaded, trimesh.Scene):
        if not loaded.geometry:
            raise ValueError("The STL did not contain any mesh geometry.")
        loaded = trimesh.util.concatenate(tuple(loaded.geometry.values()))
    if not isinstance(loaded, trimesh.Trimesh) or len(loaded.faces) == 0:
        raise ValueError("The STL could not be read as a triangle mesh.")
    return loaded


def repair_mesh(mesh: trimesh.Trimesh) -> tuple[trimesh.Trimesh, list[str]]:
    mesh = mesh.copy()
    notes: list[str] = []
    before_v, before_f = len(mesh.vertices), len(mesh.faces)
    mesh.remove_unreferenced_vertices()
    mesh.update_faces(mesh.unique_faces())
    mesh.update_faces(mesh.nondegenerate_faces())
    mesh.remove_unreferenced_vertices()
    if len(mesh.vertices) != before_v or len(mesh.faces) != before_f:
        notes.append("Removed duplicate, degenerate, or unreferenced geometry.")
    try:
        trimesh.repair.fix_normals(mesh, multibody=True)
        notes.append("Recalculated face winding and normals.")
    except Exception:
        notes.append("Normal repair was partially completed.")
    if not mesh.is_watertight:
        try:
            filled = trimesh.repair.fill_holes(mesh)
            if filled:
                notes.append("Filled simple boundary holes.")
        except Exception:
            pass
    if mesh.volume < 0:
        mesh.invert()
        notes.append("Corrected inverted mesh orientation.")
    return mesh, notes


def principal_axes(mesh: trimesh.Trimesh) -> tuple[np.ndarray, np.ndarray]:
    points = mesh.vertices - mesh.vertices.mean(axis=0)
    cov = np.cov(points.T)
    values, vectors = np.linalg.eigh(cov)
    order = np.argsort(values)[::-1]
    return values[order], vectors[:, order]


def analyse_mesh(mesh: trimesh.Trimesh, repair_notes: list[str] | None = None) -> ModelAnalysis:
    values, vectors = principal_axes(mesh)
    extents = np.asarray(mesh.extents, dtype=float)
    warning = None
    if float(extents.max()) < 5:
        warning = "The model is very small. It may have been authored in metres or inches rather than millimetres."
    elif float(extents.max()) > 1000:
        warning = "The model is unusually large. Check that the STL units are millimetres."
    volume = float(abs(mesh.volume)) if mesh.is_volume else None
    return ModelAnalysis(
        vertices=int(len(mesh.vertices)),
        faces=int(len(mesh.faces)),
        watertight=bool(mesh.is_watertight),
        volume_mm3=volume,
        extents_mm=[round(float(v), 3) for v in extents],
        centre_mm=[round(float(v), 3) for v in mesh.centroid],
        principal_axis=[round(float(v), 6) for v in vectors[:, 0]],
        suggested_scale_warning=warning,
        repair_notes=repair_notes or [],
    )


def _unit(vector: np.ndarray, fallback: np.ndarray | None = None) -> np.ndarray:
    vector = np.asarray(vector, dtype=float)
    length = np.linalg.norm(vector)
    if length < EPS:
        if fallback is None:
            raise ValueError("Cannot normalise a zero-length direction.")
        return _unit(fallback)
    return vector / length


def choose_cut_normal(mesh: trimesh.Trimesh, point: np.ndarray, surface_normal: np.ndarray | None = None) -> np.ndarray:
    _, axes = principal_axes(mesh)
    centre_direction = point - mesh.centroid
    candidates = [centre_direction, axes[:, 0]]
    if surface_normal is not None:
        # A surface normal often points out of the model rather than along a limb,
        # but it is useful when the click is close to the model centre.
        candidates.append(np.cross(surface_normal, np.array([0.0, 0.0, 1.0])))
    for candidate in candidates:
        if np.linalg.norm(candidate) > EPS:
            normal = _unit(candidate)
            # Keep the cut mostly vertical for flat-print fidget geometry.
            flat = np.array([normal[0], normal[1], 0.0])
            if np.linalg.norm(flat) > 0.18:
                return _unit(flat)
            return normal
    return np.array([1.0, 0.0, 0.0])


def estimate_local_scale(mesh: trimesh.Trimesh, point: np.ndarray) -> dict[str, float]:
    distances = np.linalg.norm(mesh.vertices - point, axis=1)
    nearest = np.partition(distances, min(200, len(distances) - 1))[: min(200, len(distances))]
    local = float(np.percentile(nearest, 80)) if len(nearest) else float(mesh.extents.min() * 0.2)
    overall_min = float(max(mesh.extents.min(), 1.0))
    radius = float(np.clip(local * 0.48, 2.8, max(3.0, overall_min * 0.28)))
    wall = float(np.clip(radius * 0.30, 1.2, 3.2))
    return {"radius": radius, "wall": wall, "local_span": local * 2.0}


def rank_joint_candidates(mesh: trimesh.Trimesh, point: np.ndarray, movement: str, clearance: float) -> list[dict[str, Any]]:
    scale = estimate_local_scale(mesh, point)
    dims = np.sort(np.asarray(mesh.extents, dtype=float))
    slenderness = float(dims[-1] / max(dims[-2], EPS))
    flatness = float(dims[0] / max(dims[1], EPS))
    near_edge = float(np.linalg.norm(point - mesh.centroid) / max(np.linalg.norm(mesh.extents) * 0.5, EPS))

    hinge_score = 78.0 + min(slenderness * 4.0, 9.0) + (7.0 if movement in {"bend", "wiggle"} else 0.0)
    ball_score = 70.0 + (10.0 if movement in {"ball", "rotate"} else 0.0) + max(0.0, (1.1 - near_edge) * 8.0)
    flexi_score = 74.0 + min(slenderness * 5.0, 11.0) + (5.0 if flatness < 0.55 else 0.0)

    candidates = [
        {
            "type": "hinge",
            "name": "Pinned Hinge",
            "score": round(min(97.0, hinge_score), 1),
            "movement": "Single-axis bend",
            "print_mode": "Assembly joint",
            "supports": "Usually none after auto-layout",
            "radius_mm": round(scale["radius"], 2),
            "clearance_mm": clearance,
            "reason": "Strong, predictable and easiest to validate for FDM printing.",
        },
        {
            "type": "ball_socket",
            "name": "Snap Ball Joint",
            "score": round(min(96.0, ball_score), 1),
            "movement": "Multi-directional rotation",
            "print_mode": "Two-piece snap assembly",
            "supports": "Depends on socket opening",
            "radius_mm": round(scale["radius"], 2),
            "clearance_mm": clearance,
            "reason": "Best where the surrounding body is thick enough for a socket wall.",
        },
        {
            "type": "flexi",
            "name": "Flexi Hinge",
            "score": round(min(98.0, flexi_score), 1),
            "movement": "Side-to-side fidget motion",
            "print_mode": "Pinned articulated segment",
            "supports": "None in recommended orientation",
            "radius_mm": round(scale["radius"] * 0.9, 2),
            "clearance_mm": clearance,
            "reason": "Designed for tails, necks, tentacles and long creature bodies.",
        },
    ]
    return sorted(candidates, key=lambda item: item["score"], reverse=True)



def suggest_joint_locations(
    mesh: trimesh.Trimesh,
    movement: str = "auto",
    clearance: float = 0.35,
    max_suggestions: int = 5,
) -> list[dict[str, Any]]:
    """Generate conservative candidate cut locations along the dominant body axis.

    This is intentionally deterministic. It does not claim that every candidate is
    printable; it gives SHiRE ranked locations to present for Ray's approval before
    any geometry is changed.
    """
    if max_suggestions < 1:
        return []
    _, axes = principal_axes(mesh)
    axis = _unit(axes[:, 0])
    projections = np.dot(mesh.vertices - mesh.centroid, axis)
    lo, hi = float(np.min(projections)), float(np.max(projections))
    span = hi - lo
    if span < EPS:
        return []

    fractions = np.linspace(0.18, 0.82, max(max_suggestions, 3))
    slab_half = max(span * 0.035, 0.35)
    labels = ["End region", "Lower-middle", "Centre", "Upper-middle", "Opposite end"]
    suggestions: list[dict[str, Any]] = []

    for index, fraction in enumerate(fractions[:max_suggestions]):
        target = lo + span * float(fraction)
        mask = np.abs(projections - target) <= slab_half
        local_vertices = mesh.vertices[mask]
        if len(local_vertices) < 8:
            nearest_idx = np.argsort(np.abs(projections - target))[: min(200, len(mesh.vertices))]
            local_vertices = mesh.vertices[nearest_idx]
        if len(local_vertices) == 0:
            continue

        local_centre = local_vertices.mean(axis=0)
        # Force the point onto the intended principal-axis station while preserving
        # the local cross-section centre in the two perpendicular directions.
        current_proj = float(np.dot(local_centre - mesh.centroid, axis))
        point = local_centre + axis * (target - current_proj)
        normal = choose_cut_normal(mesh, point)
        scale = estimate_local_scale(mesh, point)
        candidates = rank_joint_candidates(mesh, point, movement, clearance)
        top = candidates[0]

        end_penalty = abs(float(fraction) - 0.5) * 10.0
        tiny_penalty = 14.0 if scale["radius"] <= 3.0 else 0.0
        confidence = float(np.clip(top["score"] - end_penalty - tiny_penalty, 0.0, 99.0))
        suggestions.append(
            {
                "suggestion_id": f"S{index + 1:02d}",
                "label": labels[index] if index < len(labels) else f"Body station {index + 1}",
                "point": {"x": float(point[0]), "y": float(point[1]), "z": float(point[2])},
                "cut_normal": {"x": float(normal[0]), "y": float(normal[1]), "z": float(normal[2])},
                "axis_fraction": round(float(fraction), 4),
                "local_radius_mm": round(float(scale["radius"]), 3),
                "recommended_joint": top,
                "alternatives": candidates[1:],
                "confidence": round(confidence, 1),
                "warning": (
                    "Very small local cross-section; SHiRE should thicken or move the joint."
                    if scale["radius"] <= 3.0
                    else None
                ),
            }
        )

    suggestions.sort(key=lambda item: item["confidence"], reverse=True)
    return suggestions

def _slice_with_gap(mesh: trimesh.Trimesh, origin: np.ndarray, normal: np.ndarray, gap: float) -> tuple[trimesh.Trimesh, trimesh.Trimesh]:
    n = _unit(normal)
    try:
        positive = trimesh.intersections.slice_mesh_plane(
            mesh, plane_normal=n, plane_origin=origin + n * (gap * 0.5), cap=True
        )
        negative = trimesh.intersections.slice_mesh_plane(
            mesh, plane_normal=-n, plane_origin=origin - n * (gap * 0.5), cap=True
        )
    except Exception:
        positive = trimesh.intersections.slice_mesh_plane(
            mesh, plane_normal=n, plane_origin=origin + n * (gap * 0.5), cap=False
        )
        negative = trimesh.intersections.slice_mesh_plane(
            mesh, plane_normal=-n, plane_origin=origin - n * (gap * 0.5), cap=False
        )
    if len(positive.faces) == 0 or len(negative.faces) == 0:
        raise ValueError("The selected point did not split the model into two usable sections.")
    return negative, positive


def _align_z(mesh: trimesh.Trimesh, direction: np.ndarray) -> trimesh.Trimesh:
    direction = _unit(direction)
    z = np.array([0.0, 0.0, 1.0])
    dot = float(np.clip(np.dot(z, direction), -1.0, 1.0))
    if abs(dot - 1.0) < 1e-8:
        return mesh
    if abs(dot + 1.0) < 1e-8:
        axis = np.array([1.0, 0.0, 0.0])
        angle = math.pi
    else:
        axis = _unit(np.cross(z, direction))
        angle = math.acos(dot)
    mesh.apply_transform(rotation_matrix(angle, axis))
    return mesh


def cylinder_between(radius: float, start: np.ndarray, end: np.ndarray, sections: int = 48) -> trimesh.Trimesh:
    direction = end - start
    height = float(np.linalg.norm(direction))
    cylinder = trimesh.creation.cylinder(radius=radius, height=height, sections=sections)
    _align_z(cylinder, direction)
    cylinder.apply_translation((start + end) * 0.5)
    return cylinder


def safe_union(meshes: list[trimesh.Trimesh]) -> trimesh.Trimesh:
    try:
        return trimesh.boolean.union(meshes, engine="manifold", check_volume=False)
    except Exception:
        return trimesh.util.concatenate(meshes)


def safe_difference(base: trimesh.Trimesh, cutters: list[trimesh.Trimesh]) -> trimesh.Trimesh:
    try:
        result = trimesh.boolean.difference([base, *cutters], engine="manifold", check_volume=False)
        if isinstance(result, list):
            result = trimesh.util.concatenate(result)
        return result
    except Exception as exc:
        raise ValueError(f"Boolean subtraction failed: {exc}") from exc


def forge_hinge(mesh: trimesh.Trimesh, point: np.ndarray, normal: np.ndarray, radius: float, clearance: float, flexi: bool = False) -> dict[str, trimesh.Trimesh]:
    gap = max(clearance * 2.0, 0.65)
    a, b = _slice_with_gap(mesh, point, normal, gap)
    build_axis = np.array([0.0, 0.0, 1.0])
    thickness = float(max(mesh.extents[2], radius * 1.7))
    pin_r = max(1.25, radius * (0.30 if flexi else 0.34))
    lug_r = pin_r + clearance + max(1.15, radius * 0.20)
    total_h = min(max(thickness, lug_r * 3.2), max(thickness * 1.15, 8.0))
    centre_h = total_h * 0.38
    side_h = (total_h - centre_h - clearance * 4.0) * 0.5
    if side_h < 1.0:
        side_h = max(1.0, total_h * 0.22)
        centre_h = total_h - side_h * 2.0 - clearance * 4.0

    # Offset the hinge slightly toward the cut so both lugs intersect their body section.
    centre = point.copy()
    centre[2] = float(np.clip(point[2], mesh.bounds[0, 2] + total_h * 0.5, mesh.bounds[1, 2] - total_h * 0.5)) if thickness > total_h else float(mesh.centroid[2])

    central_lug = trimesh.creation.cylinder(radius=lug_r, height=centre_h, sections=64)
    central_lug.apply_translation(centre)

    z_low = centre[2] - centre_h * 0.5 - clearance * 2.0 - side_h * 0.5
    z_high = centre[2] + centre_h * 0.5 + clearance * 2.0 + side_h * 0.5
    low_lug = trimesh.creation.cylinder(radius=lug_r, height=side_h, sections=64)
    high_lug = low_lug.copy()
    low_lug.apply_translation([centre[0], centre[1], z_low])
    high_lug.apply_translation([centre[0], centre[1], z_high])

    # Necks bridge each lug into its side of the original model.
    neck_len = radius * 1.2
    neck_r = lug_r * 0.62
    neck_a = cylinder_between(neck_r, centre - normal * neck_len, centre)
    neck_b = cylinder_between(neck_r, centre + normal * neck_len, centre)

    a_plus = safe_union([a, central_lug, neck_a])
    b_plus = safe_union([b, low_lug, high_lug, neck_b])

    hole = trimesh.creation.cylinder(radius=pin_r + clearance, height=total_h + 2.0, sections=64)
    hole.apply_translation(centre)
    a_final = safe_difference(a_plus, [hole])
    b_final = safe_difference(b_plus, [hole])

    pin_h = total_h + 0.7
    pin = trimesh.creation.cylinder(radius=pin_r, height=pin_h, sections=64)
    cap_r = pin_r + 0.7
    cap_h = 0.6
    cap1 = trimesh.creation.cylinder(radius=cap_r, height=cap_h, sections=64)
    cap2 = cap1.copy()
    cap1.apply_translation([0, 0, -pin_h * 0.5 - cap_h * 0.5])
    cap2.apply_translation([0, 0, pin_h * 0.5 + cap_h * 0.5])
    pin = safe_union([pin, cap1, cap2])
    pin.apply_translation(centre)

    assembled = trimesh.util.concatenate([a_final, b_final, pin])
    return {"part_a": a_final, "part_b": b_final, "pin": pin, "assembled": assembled}


def forge_ball_socket(mesh: trimesh.Trimesh, point: np.ndarray, normal: np.ndarray, radius: float, clearance: float) -> dict[str, trimesh.Trimesh]:
    gap = max(clearance * 1.7, 0.55)
    a, b = _slice_with_gap(mesh, point, normal, gap)
    ball_r = radius
    wall = max(1.2, radius * 0.27)
    socket_outer = ball_r + clearance + wall
    centre = point + normal * (ball_r * 0.20)

    ball = trimesh.creation.icosphere(subdivisions=3, radius=ball_r)
    ball.apply_translation(centre)
    neck_start = centre - normal * (ball_r * 1.65)
    neck_end = centre - normal * (ball_r * 0.55)
    neck = cylinder_between(max(1.25, ball_r * 0.34), neck_start, neck_end)
    a_final = safe_union([a, ball, neck])

    outer = trimesh.creation.icosphere(subdivisions=3, radius=socket_outer)
    inner = trimesh.creation.icosphere(subdivisions=3, radius=ball_r + clearance)
    outer.apply_translation(centre)
    inner.apply_translation(centre)
    shell = safe_difference(outer, [inner])

    # Create a snap opening facing the ball stem.
    mouth_r = max(ball_r * 0.72, ball_r - wall * 0.55)
    mouth_start = centre - normal * (socket_outer * 1.5)
    mouth_end = centre - normal * (ball_r * 0.05)
    mouth = cylinder_between(mouth_r, mouth_start, mouth_end, sections=64)
    shell = safe_difference(shell, [mouth])
    socket_neck = cylinder_between(max(1.4, socket_outer * 0.48), centre + normal * (socket_outer * 0.35), centre + normal * (socket_outer * 1.55))
    b_final = safe_union([b, shell, socket_neck])
    assembled = trimesh.util.concatenate([a_final, b_final])
    return {"part_a": a_final, "part_b": b_final, "assembled": assembled}


def layout_parts(parts: dict[str, trimesh.Trimesh], gap: float) -> trimesh.Trimesh:
    printable = [(name, mesh.copy()) for name, mesh in parts.items() if name != "assembled"]
    cursor = 0.0
    laid_out: list[trimesh.Trimesh] = []
    for _, mesh in printable:
        mesh.apply_translation(-mesh.bounds[0])
        mesh.apply_translation([cursor, 0.0, 0.0])
        cursor = float(mesh.bounds[1, 0] + gap)
        laid_out.append(mesh)
    return trimesh.util.concatenate(laid_out)


def export_forge_package(
    output_dir: Path,
    source_name: str,
    mesh: trimesh.Trimesh,
    point: np.ndarray,
    normal: np.ndarray,
    joint_type: str,
    clearance: float,
    layout_gap: float,
) -> dict[str, Any]:
    scale = estimate_local_scale(mesh, point)
    radius = scale["radius"] * (0.88 if joint_type == "flexi" else 1.0)
    if joint_type in {"hinge", "flexi"}:
        parts = forge_hinge(mesh, point, normal, radius, clearance, flexi=joint_type == "flexi")
    elif joint_type == "ball_socket":
        parts = forge_ball_socket(mesh, point, normal, radius, clearance)
    else:
        raise ValueError(f"Unsupported joint type: {joint_type}")

    package_id = uuid.uuid4().hex[:12]
    package = output_dir / package_id
    package.mkdir(parents=True, exist_ok=False)
    files: dict[str, str] = {}
    for name, part in parts.items():
        filename = f"{source_name}_{name}.stl"
        part.export(package / filename)
        files[name] = filename
    layout = layout_parts(parts, layout_gap)
    layout_name = f"{source_name}_PRINT_LAYOUT.stl"
    layout.export(package / layout_name)
    files["print_layout"] = layout_name

    report = {
        "engine": "FlexiForge MVP 0.1",
        "joint_type": joint_type,
        "clearance_mm": clearance,
        "joint_radius_mm": round(radius, 3),
        "selection_point_mm": [round(float(v), 4) for v in point],
        "cut_normal": [round(float(v), 6) for v in normal],
        "files": files,
        "important": [
            "This MVP creates real boolean joint geometry and separate printable parts.",
            "Run a small clearance test before committing to a long print.",
            "The hinge pin is a separate printable component; the ball joint requires snap assembly.",
            "Inspect the generated parts in a slicer before printing.",
        ],
    }
    (package / "FlexiForge_Report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    return {"package_id": package_id, "files": files, "report": report, "directory": package}
