"""Čtení a zápis STL (binární i ASCII) bez externích závislostí.

Vše je v NumPy, mesh je reprezentován jako:
    vertices : (V, 3) float64
    faces    : (F, 3) int32   -- indexy do vertices
"""
from __future__ import annotations

import io
import os
import struct

import numpy as np


# --------------------------------------------------------------------------- #
#  Čtení
# --------------------------------------------------------------------------- #
def _is_binary_stl(data: bytes) -> bool:
    """Rozpozná binární STL. ASCII STL začíná 'solid' a obsahuje 'facet normal'.

    Pozor: některé binární STL mají v hlavičce také 'solid', proto se rozhoduje
    podle konzistence délky souboru s deklarovaným počtem trojúhelníků.
    """
    if len(data) < 84:
        return False
    n_tri = struct.unpack("<I", data[80:84])[0]
    expected = 84 + n_tri * 50
    if expected == len(data):
        return True
    # tolerantně: některé exportéry přidávají na konec smetí
    if n_tri > 0 and expected <= len(data) <= expected + 4096:
        head = data[:512].lstrip().lower()
        if not head.startswith(b"solid"):
            return True
        # 'solid' v hlavičce, ale sedí délka -> binární
        if b"facet normal" not in data[:2048].lower():
            return True
    head = data[:512].lstrip().lower()
    return not head.startswith(b"solid")


def _read_binary_stl(data: bytes):
    n_tri = struct.unpack("<I", data[80:84])[0]
    body = data[84:84 + n_tri * 50]
    if len(body) < n_tri * 50:
        n_tri = len(body) // 50
        body = body[:n_tri * 50]
    rec = np.frombuffer(body, dtype=np.uint8).reshape(n_tri, 50)
    floats = rec[:, :48].copy().view("<f4").reshape(n_tri, 12)
    normals = floats[:, 0:3].astype(np.float64)
    tri = floats[:, 3:12].astype(np.float64).reshape(n_tri, 3, 3)
    return tri, normals


def _read_ascii_stl(text: str):
    verts = []
    normals = []
    cur = []
    for raw in text.splitlines():
        line = raw.strip()
        if not line:
            continue
        low = line.lower()
        if low.startswith("vertex"):
            p = line.split()
            cur.append((float(p[1]), float(p[2]), float(p[3])))
        elif low.startswith("facet normal"):
            p = line.split()
            try:
                normals.append((float(p[2]), float(p[3]), float(p[4])))
            except (ValueError, IndexError):
                normals.append((0.0, 0.0, 0.0))
        elif low.startswith("endfacet"):
            if len(cur) == 3:
                verts.append(cur)
            cur = []
    if not verts:
        raise ValueError("ASCII STL neobsahuje žádné trojúhelníky.")
    tri = np.asarray(verts, dtype=np.float64)
    nrm = np.asarray(normals, dtype=np.float64)
    if nrm.shape[0] != tri.shape[0]:
        nrm = np.zeros((tri.shape[0], 3))
    return tri, nrm


def read_stl_triangles(path_or_bytes):
    """Vrátí (triangles (F,3,3), file_normals (F,3))."""
    if isinstance(path_or_bytes, (bytes, bytearray)):
        data = bytes(path_or_bytes)
    else:
        with open(path_or_bytes, "rb") as fh:
            data = fh.read()
    if len(data) < 15:
        raise ValueError("Soubor je příliš malý na to, aby byl STL.")
    if _is_binary_stl(data):
        return _read_binary_stl(data)
    return _read_ascii_stl(data.decode("utf-8", errors="replace"))


def weld_vertices(tri: np.ndarray, tol: float = 1e-7):
    """Sešije duplicitní vrcholy (STL je „polévka trojúhelníků“).

    tol je absolutní tolerance ve stejných jednotkách jako model (mm).
    """
    pts = tri.reshape(-1, 3)
    if tol <= 0:
        uniq, inv = np.unique(pts, axis=0, return_inverse=True)
    else:
        keys = np.round(pts / tol).astype(np.int64)
        uniq_keys, idx, inv = np.unique(keys, axis=0, return_index=True,
                                        return_inverse=True)
        uniq = pts[idx]
    faces = inv.reshape(-1, 3).astype(np.int64)
    # zahodit degenerované trojúhelníky
    ok = (faces[:, 0] != faces[:, 1]) & (faces[:, 1] != faces[:, 2]) & \
         (faces[:, 0] != faces[:, 2])
    return uniq, faces[ok]


def read_stl(path_or_bytes, weld_tol: float = 1e-6):
    """Načte STL a vrátí (vertices, faces)."""
    tri, _ = read_stl_triangles(path_or_bytes)
    return weld_vertices(tri, weld_tol)


# --------------------------------------------------------------------------- #
#  OBJ / PLY (ASCII)
# --------------------------------------------------------------------------- #
def read_obj(path_or_bytes):
    """Wavefront OBJ — vrcholy a stěny (n-úhelníky se rozdělí vějířem)."""
    if isinstance(path_or_bytes, (bytes, bytearray)):
        text = bytes(path_or_bytes).decode("utf-8", errors="replace")
    else:
        with open(path_or_bytes, "r", encoding="utf-8", errors="replace") as fh:
            text = fh.read()

    verts = []
    faces = []
    for raw in text.splitlines():
        if not raw or raw[0] == "#":
            continue
        if raw.startswith("v "):
            p = raw.split()
            verts.append((float(p[1]), float(p[2]), float(p[3])))
        elif raw.startswith("f "):
            idx = []
            for tok in raw.split()[1:]:
                a = tok.split("/")[0]
                if not a:
                    continue
                i = int(a)
                idx.append(i - 1 if i > 0 else len(verts) + i)
            for k in range(1, len(idx) - 1):        # vějířová triangulace
                faces.append((idx[0], idx[k], idx[k + 1]))
    if not verts or not faces:
        raise ValueError("OBJ neobsahuje použitelné vrcholy nebo stěny.")
    V = np.asarray(verts, dtype=np.float64)
    F = np.asarray(faces, dtype=np.int64)
    ok = (F >= 0).all(axis=1) & (F < V.shape[0]).all(axis=1)
    F = F[ok]
    ok = (F[:, 0] != F[:, 1]) & (F[:, 1] != F[:, 2]) & (F[:, 0] != F[:, 2])
    return V, F[ok]


def read_ply(path_or_bytes):
    """PLY — podporován ASCII i binary_little_endian s float32 x,y,z."""
    if isinstance(path_or_bytes, (bytes, bytearray)):
        data = bytes(path_or_bytes)
    else:
        with open(path_or_bytes, "rb") as fh:
            data = fh.read()
    end = data.find(b"end_header")
    if end < 0:
        raise ValueError("PLY bez hlavičky.")
    header = data[:end].decode("ascii", errors="replace").splitlines()
    body_start = data.find(b"\n", end) + 1

    fmt = "ascii"
    n_v = n_f = 0
    props = []
    cur = None
    for ln in header:
        p = ln.split()
        if not p:
            continue
        if p[0] == "format":
            fmt = p[1]
        elif p[0] == "element":
            cur = p[1]
            if cur == "vertex":
                n_v = int(p[2])
            elif cur == "face":
                n_f = int(p[2])
        elif p[0] == "property" and cur == "vertex" and p[1] != "list":
            props.append((p[1], p[2]))

    if fmt == "ascii":
        toks = data[body_start:].decode("ascii", errors="replace").split()
        k = 0
        np_ = len(props)
        V = np.empty((n_v, 3))
        names = [n for _, n in props]
        ix, iy, iz = names.index("x"), names.index("y"), names.index("z")
        for i in range(n_v):
            row = toks[k:k + np_]
            k += np_
            V[i] = (float(row[ix]), float(row[iy]), float(row[iz]))
        faces = []
        for _ in range(n_f):
            c = int(toks[k]); k += 1
            idx = [int(x) for x in toks[k:k + c]]
            k += c
            for m in range(1, c - 1):
                faces.append((idx[0], idx[m], idx[m + 1]))
        return V, np.asarray(faces, dtype=np.int64)

    raise ValueError("Binární PLY zatím není podporováno — ulož prosím jako "
                     "STL, OBJ nebo ASCII PLY.")


def read_mesh(path_or_bytes, name: str | None = None):
    """Načte STL / OBJ / ASCII PLY podle přípony nebo obsahu."""
    ext = ""
    if name:
        ext = os.path.splitext(name)[1].lower()
    elif isinstance(path_or_bytes, (str, os.PathLike)):
        ext = os.path.splitext(str(path_or_bytes))[1].lower()

    if ext == ".obj":
        return read_obj(path_or_bytes)
    if ext == ".ply":
        return read_ply(path_or_bytes)
    if ext in (".stl", ""):
        try:
            return read_stl(path_or_bytes)
        except Exception:
            if ext == ".stl":
                raise
    for fn in (read_obj, read_ply, read_stl):
        try:
            return fn(path_or_bytes)
        except Exception:
            continue
    raise ValueError("Formát souboru se nepodařilo rozpoznat "
                     "(podporováno STL, OBJ, ASCII PLY).")


# --------------------------------------------------------------------------- #
#  Zápis
# --------------------------------------------------------------------------- #
def face_normals(vertices: np.ndarray, faces: np.ndarray,
                 normalize: bool = True) -> np.ndarray:
    v0 = vertices[faces[:, 0]]
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]
    n = np.cross(v1 - v0, v2 - v0)
    if normalize:
        ln = np.linalg.norm(n, axis=1)
        ln[ln == 0] = 1.0
        n = n / ln[:, None]
    return n


def write_stl_binary(path_or_stream, vertices: np.ndarray, faces: np.ndarray,
                     header: str = "Babofkator") -> None:
    n = face_normals(vertices, faces)
    tri = vertices[faces]                       # (F,3,3)
    f_cnt = faces.shape[0]
    rec = np.zeros((f_cnt, 50), dtype=np.uint8)
    payload = np.empty((f_cnt, 12), dtype="<f4")
    payload[:, 0:3] = n
    payload[:, 3:12] = tri.reshape(f_cnt, 9)
    rec[:, :48] = payload.view(np.uint8).reshape(f_cnt, 48)

    head = header.encode("ascii", errors="replace")[:79]
    head = head + b" " * (80 - len(head))

    own = isinstance(path_or_stream, (str, bytes, os.PathLike))
    fh = open(path_or_stream, "wb") if own else path_or_stream
    try:
        fh.write(head)
        fh.write(struct.pack("<I", f_cnt))
        fh.write(rec.tobytes())
    finally:
        if own:
            fh.close()


def stl_bytes(vertices, faces, header="Babofkator") -> bytes:
    buf = io.BytesIO()
    write_stl_binary(buf, vertices, faces, header)
    return buf.getvalue()
