from __future__ import annotations

import math
import re
from dataclasses import dataclass

import numpy as np

from .mesh import FaceVertex, Meshlet, ObjMesh
from .meshlets import build_compatible_edge_to_triangles
from .progress import Heartbeat, log
from .stripifier import DEFAULT_LKH_EXE, default_stripify_budget
from .exporter import _make_chains


@dataclass
class LegacyExportResult:
    header: str
    chains: int
    face_words: int
    meshes: int = 1
    components: int = 0
    triangles: int = 0
    source: str | None = None


def _identifier(name: str) -> str:
    out = re.sub(r"\W+", "_", name.strip())
    if not out:
        out = "mesh"
    if out[0].isdigit():
        out = "_" + out
    return out


def _fmt_float(v: float) -> str:
    if not math.isfinite(v):
        return "0.0f"
    text = f"{float(v):.9g}"
    if "e" not in text and "E" not in text and "." not in text:
        text += ".0"
    return text + "f"


def _fmt_vec(values) -> str:
    return "{ " + ", ".join(_fmt_float(float(x)) for x in values) + " }"


def export_legacy_mesh3d_header(
    mesh: ObjMesh,
    *,
    name: str,
    color_type: str = "tgx::RGB565",
    meshlets: list[Meshlet] | None = None,
    use_lkh: bool = True,
    lkh_exe=DEFAULT_LKH_EXE,
    lkh_time_limit_s: float | None = None,
    texture_symbols: dict[str, str] | None = None,
    extra_includes: list[str] | None = None,
    split_by_material: bool = True,
    header_filename: str | None = None,
    single_header: bool = True,
) -> LegacyExportResult:
    """Export a Mesh3D validation reference.

    When meshlets are provided, the face stream is generated with the same meshlet order and
    triangle strips as the meshlet formats. This makes comparisons meaningful: the intended
    difference is that meshlet formats may skip invisible meshlets before reading their payload.
    """
    symbol = _identifier(name)
    has_texcoords = mesh.has_texcoords()
    has_normals = mesh.has_normals()
    elem_words = 1 + (1 if has_texcoords else 0) + (1 if has_normals else 0)
    extra_includes = extra_includes or []
    texture_symbols = texture_symbols or {}
    groups = _material_groups(mesh) if split_by_material else [("", list(range(len(mesh.triangles))))]
    if not groups:
        groups = [("", [])]
    if meshlets is not None:
        groups = _meshlet_groups(meshlets)

    encoded_groups: list[dict] = []
    total_chains = 0
    total_face_words = 0
    total_components = 0
    for order, (material, tri_indices) in enumerate(groups):
        if meshlets is None:
            chains, components = _build_material_chains(
                mesh,
                tri_indices,
                use_lkh=use_lkh,
                lkh_exe=lkh_exe,
                lkh_time_limit_s=lkh_time_limit_s,
            )
        else:
            chains = _build_reference_chains(mesh, [m for m in meshlets if m.material == material], use_lkh=use_lkh, lkh_exe=lkh_exe, lkh_time_limit_s=lkh_time_limit_s)
            components = len(chains)
        face_words = 1 + sum(1 + (len(chain) + 2) * elem_words for chain in chains)
        group_bb_min, group_bb_max = _group_bounding_box(mesh, tri_indices)
        suffix = "" if order == 0 else f"_{order + 1}"
        encoded_groups.append(
            dict(
                material=material,
                tri_indices=tri_indices,
                chains=chains,
                face_words=face_words,
                bb_min=group_bb_min,
                bb_max=group_bb_max,
                face_symbol=f"{symbol}{suffix}_face",
                mesh_symbol=f"{symbol}{suffix}",
            )
        )
        total_chains += len(chains)
        total_face_words += face_words
        total_components += components

    vertices, texcoords, normals = _remap_encoded_groups_to_stream_order(mesh, encoded_groups, has_texcoords, has_normals)

    header: list[str] = []
    header.append("#pragma once")
    header.append("")
    header.append("#include <tgx.h>")
    header.append("")
    header.append("// Mesh3D generated by tgx_mesh with optimized triangle strips.")
    header.append(f"// Source model : {mesh.name}")
    header.append(f"// Symbol       : {symbol}")
    header.append(f"// Color type   : {color_type}")
    header.append(f"// Triangles    : {len(mesh.triangles)}")
    header.append(f"// Submeshes    : {len(encoded_groups)}")
    header.append(f"// Components   : {total_components}")
    header.append(f"// Chains       : {total_chains} ({100.0 * total_chains / max(1, len(mesh.triangles)):.2f}% of triangles)")
    header.append("")

    definitions: list[str] = []

    definitions.append(f"const tgx::fVec3 {symbol}_vert_array[{len(vertices)}] PROGMEM = {{")
    for v in vertices:
        definitions.append(f"    {_fmt_vec(v)},")
    definitions.append("};")
    definitions.append("")

    tex_name = "nullptr"
    if has_texcoords:
        tex_name = f"{symbol}_tex_array"
        definitions.append(f"const tgx::fVec2 {tex_name}[{len(texcoords)}] PROGMEM = {{")
        for vt in texcoords:
            definitions.append(f"    {_fmt_vec(vt)},")
        definitions.append("};")
        definitions.append("")

    norm_name = "nullptr"
    if has_normals:
        norm_name = f"{symbol}_norm_array"
        definitions.append(f"const tgx::fVec3 {norm_name}[{len(normals)}] PROGMEM = {{")
        for n in normals:
            definitions.append(f"    {_fmt_vec(n)},")
        definitions.append("};")
        definitions.append("")

    for group in encoded_groups:
        definitions.append(f"const uint16_t {group['face_symbol']}[{group['face_words']}] PROGMEM = {{")
        for chain in group["chains"]:
            definitions.append(f"{len(chain)},")
            first = chain[0][1]
            definitions.append(_face_corner(first[0], has_texcoords, has_normals))
            definitions.append(_face_corner(first[1], has_texcoords, has_normals))
            definitions.append(_face_corner(first[2], has_texcoords, has_normals))
            for direction, tri in chain[1:]:
                assert direction is not None
                definitions.append(_face_corner(tri[2], has_texcoords, has_normals, direction=direction))
        definitions.append("0")
        definitions.append("};")
        definitions.append("")

    # Define linked Mesh3D declarations from tail to head, so each next pointer
    # can refer to a declaration already visible to C++.
    for i in range(len(encoded_groups) - 1, -1, -1):
        group = encoded_groups[i]
        material = mesh.materials.get(group["material"])
        color = material.diffuse if material is not None else (0.75, 0.75, 0.75)
        ambiant = material.ambiant_strength if material is not None else 0.1
        diffuse = material.diffuse_strength if material is not None else 0.7
        specular = material.specular_strength if material is not None else 0.6
        exponent = material.specular_exponent if material is not None else 32
        texture = texture_symbols.get(group["material"], "")
        texture_ptr = f"&{texture}" if texture else "nullptr"
        next_ptr = f"&{encoded_groups[i + 1]['mesh_symbol']}" if i + 1 < len(encoded_groups) else "nullptr"
        gmin = group["bb_min"]
        gmax = group["bb_max"]
        object_prefix = "const" if single_header else "extern const"
        definitions.append(f"{object_prefix} tgx::Mesh3D<{color_type}> {group['mesh_symbol']} PROGMEM =")
        definitions.append("    {")
        definitions.append("    1,")
        definitions.append(f"    {len(vertices)},")
        definitions.append(f"    {len(texcoords) if has_texcoords else 0},")
        definitions.append(f"    {len(normals) if has_normals else 0},")
        definitions.append(f"    {len(group['tri_indices'])},")
        definitions.append(f"    {group['face_words']},")
        definitions.append(f"    {symbol}_vert_array,")
        definitions.append(f"    {tex_name},")
        definitions.append(f"    {norm_name},")
        definitions.append(f"    {group['face_symbol']},")
        definitions.append(f"    {texture_ptr},")
        definitions.append(f"    {{ {_fmt_float(color[0])}, {_fmt_float(color[1])}, {_fmt_float(color[2])} }},")
        definitions.append(f"    {_fmt_float(ambiant)},")
        definitions.append(f"    {_fmt_float(diffuse)},")
        definitions.append(f"    {_fmt_float(specular)},")
        definitions.append(f"    {int(exponent)},")
        definitions.append(f"    {next_ptr},")
        definitions.append("    {")
        definitions.append(f"    {_fmt_float(gmin[0])}, {_fmt_float(gmax[0])},")
        definitions.append(f"    {_fmt_float(gmin[1])}, {_fmt_float(gmax[1])},")
        definitions.append(f"    {_fmt_float(gmin[2])}, {_fmt_float(gmax[2])}")
        definitions.append("    },")
        definitions.append(f"    \"{group['mesh_symbol']}\"")
        definitions.append("    };")
        definitions.append("")

    if single_header:
        for include in extra_includes:
            header.append(f'#include "{include}"')
        header.append("")
        header.extend(definitions)
        source = None
    else:
        for group in encoded_groups:
            header.append(f"extern const tgx::Mesh3D<{color_type}> {group['mesh_symbol']};")
        header.append("")
        source_lines = [f'#include "{header_filename or (symbol + ".h")}"']
        for include in extra_includes:
            source_lines.append(f'#include "{include}"')
        source_lines.append("")
        source_lines.extend(definitions)
        source = "\n".join(source_lines)

    return LegacyExportResult(
        "\n".join(header),
        total_chains,
        total_face_words,
        len(encoded_groups),
        total_components,
        len(mesh.triangles),
        source,
    )


def _remap_encoded_groups_to_stream_order(mesh: ObjMesh, encoded_groups: list[dict], has_texcoords: bool, has_normals: bool):
    v_order: list[int] = []
    vt_order: list[int] = []
    vn_order: list[int] = []
    v_map: dict[int, int] = {}
    vt_map: dict[int, int] = {}
    vn_map: dict[int, int] = {}

    def remap_corner(corner: FaceVertex) -> FaceVertex:
        if corner.v not in v_map:
            v_map[corner.v] = len(v_order)
            v_order.append(corner.v)
        vt = corner.vt
        if has_texcoords and corner.vt >= 0:
            if corner.vt not in vt_map:
                vt_map[corner.vt] = len(vt_order)
                vt_order.append(corner.vt)
            vt = vt_map[corner.vt]
        vn = corner.vn
        if has_normals and corner.vn >= 0:
            if corner.vn not in vn_map:
                vn_map[corner.vn] = len(vn_order)
                vn_order.append(corner.vn)
            vn = vn_map[corner.vn]
        return FaceVertex(v_map[corner.v], vt, vn)

    for group in encoded_groups:
        remapped_chains = []
        for chain in group["chains"]:
            remapped_chain = []
            for direction, tri in chain:
                remapped_chain.append((direction, tuple(remap_corner(c) for c in tri)))
            remapped_chains.append(remapped_chain)
        group["chains"] = remapped_chains

    vertices = mesh.vertices[v_order] if v_order else mesh.vertices[:0]
    texcoords = mesh.texcoords[vt_order] if has_texcoords and vt_order else mesh.texcoords[:0]
    normals = mesh.normals[vn_order] if has_normals and vn_order else mesh.normals[:0]
    return vertices, texcoords, normals


def _build_reference_chains(mesh: ObjMesh, meshlets: list[Meshlet] | None, *, use_lkh: bool, lkh_exe, lkh_time_limit_s: float | None = None) -> list:
    if meshlets is None:
        return [[(None, tri.corners)] for tri in mesh.triangles]

    chains = []
    for meshlet in meshlets:
        chains.extend(_make_chains(mesh, meshlet.triangles, use_lkh=use_lkh, lkh_exe=lkh_exe, lkh_time_limit_s=lkh_time_limit_s))
    return chains


def _material_groups(mesh: ObjMesh) -> list[tuple[str, list[int]]]:
    groups: dict[str, list[int]] = {}
    order: list[str] = []
    for tri_index, tri in enumerate(mesh.triangles):
        if tri.material not in groups:
            groups[tri.material] = []
            order.append(tri.material)
        groups[tri.material].append(tri_index)
    return [(material, groups[material]) for material in order]


def _meshlet_groups(meshlets: list[Meshlet]) -> list[tuple[str, list[int]]]:
    groups: dict[str, list[int]] = {}
    order: list[str] = []
    for meshlet in meshlets:
        if meshlet.material not in groups:
            groups[meshlet.material] = []
            order.append(meshlet.material)
        groups[meshlet.material].extend(meshlet.triangles)
    return [(material, groups[material]) for material in order]


def _connected_components(mesh: ObjMesh, tri_indices: list[int]) -> list[list[int]]:
    component_set = set(tri_indices)
    adjacency = [set() for _ in mesh.triangles]
    for tris in build_compatible_edge_to_triangles(mesh).values():
        local = [ti for ti in tris if ti in component_set]
        if len(local) < 2:
            continue
        for a in local:
            adjacency[a].update(b for b in local if b != a)
    remaining = set(tri_indices)
    components: list[list[int]] = []
    while remaining:
        start = remaining.pop()
        stack = [start]
        component = [start]
        while stack:
            current = stack.pop()
            for neighbor in adjacency[current]:
                if neighbor in remaining:
                    remaining.remove(neighbor)
                    stack.append(neighbor)
                    component.append(neighbor)
        component.sort()
        components.append(component)
    components.sort(key=lambda item: (-len(item), item[0] if item else 0))
    return components


def _build_material_chains(mesh: ObjMesh, tri_indices: list[int], *, use_lkh: bool, lkh_exe, lkh_time_limit_s: float | None) -> tuple[list, int]:
    chains = []
    components = _connected_components(mesh, tri_indices)
    total_triangles = sum(len(component) for component in components)
    if total_triangles >= 512 or len(components) >= 16:
        log(f"legacy strips: {len(components)} compatible components, {total_triangles} triangles")
    heartbeat = Heartbeat("legacy strips")
    processed = 0
    for index, component in enumerate(components, start=1):
        budget = lkh_time_limit_s if lkh_time_limit_s is not None else default_stripify_budget(len(component))
        if len(component) >= 16 or total_triangles >= 512 or len(components) >= 16:
            log(f"legacy strips: component {index}/{len(components)}, {len(component)} triangles, budget {budget:.2f}s")
        before = len(chains)
        chains.extend(_make_chains(mesh, component, use_lkh=use_lkh, lkh_exe=lkh_exe, lkh_time_limit_s=lkh_time_limit_s))
        processed += len(component)
        if total_triangles >= 512 or len(components) >= 16:
            heartbeat.ping(f"{processed}/{total_triangles} triangles, {len(chains)} chains (+{len(chains) - before})")
    if total_triangles >= 512 or len(components) >= 16:
        log(f"legacy strips done: {len(chains)} chains")
    return chains, len(components)


def _group_bounding_box(mesh: ObjMesh, tri_indices: list[int]) -> tuple[np.ndarray, np.ndarray]:
    if not tri_indices:
        return mesh.bounding_box()
    vertices = sorted({corner.v for tri_index in tri_indices for corner in mesh.triangles[tri_index].corners})
    pts = mesh.vertices[vertices]
    return np.min(pts, axis=0), np.max(pts, axis=0)


def _face_corner(corner: FaceVertex, has_texcoords: bool, has_normals: bool, *, direction: int = 0) -> str:
    v = corner.v | (32768 if direction else 0)
    parts = [str(v)]
    if has_texcoords:
        parts.append(str(corner.vt))
    if has_normals:
        parts.append(str(corner.vn))
    return ", ".join(parts) + ","
