#!/usr/bin/env python3
"""Convert .npy face embeddings to pickle-free .f32emb before HF Space deploy.

Uses only stdlib so GitHub Actions runners do not need numpy pre-installed.
"""
from __future__ import annotations

import os
import struct
import sys

_FR = os.path.dirname(os.path.abspath(__file__))
if _FR not in sys.path:
    sys.path.insert(0, _FR)

MAGIC = b"CEPF32E1"
HEADER = struct.Struct("<8sII")
NPY_MAGIC = b"\x93NUMPY"


def _parse_npy_header(raw: bytes) -> tuple[str, int]:
    if not raw.startswith(NPY_MAGIC):
        raise ValueError("Not a .npy file")
    major, minor = raw[6], raw[7]
    if (major, minor) == (1, 0):
        hlen = struct.unpack("<H", raw[8:10])[0]
        header = raw[10 : 10 + hlen].decode("latin1").strip()
        offset = 10 + hlen
    elif (major, minor) >= (2, 0):
        hlen = struct.unpack("<I", raw[8:12])[0]
        header = raw[12 : 12 + hlen].decode("latin1").strip()
        offset = 12 + hlen
    else:
        raise ValueError(f"Unsupported .npy version {(major, minor)}")
    return header, offset


def _shape_from_header(header: str) -> tuple[int, ...]:
    # e.g. {'descr': '<f4', 'fortran_order': False, 'shape': (2, 512), }
    start = header.find("'shape':")
    if start < 0:
        raise ValueError("Missing shape in npy header")
    chunk = header[start:]
    open_paren = chunk.find("(")
    close_paren = chunk.find(")", open_paren)
    inner = chunk[open_paren + 1 : close_paren]
    parts = [p.strip() for p in inner.split(",") if p.strip()]
    return tuple(int(p) for p in parts)


def npy_to_f32emb(npy_path: str, f32_path: str) -> None:
    with open(npy_path, "rb") as fh:
        raw = fh.read()
    header, offset = _parse_npy_header(raw)
    shape = _shape_from_header(header)
    if len(shape) == 1:
        n_vec, dim = 1, shape[0]
    elif len(shape) == 2:
        n_vec, dim = shape
    else:
        raise ValueError(f"Unsupported embedding shape {shape}")
    data = raw[offset : offset + n_vec * dim * 4]
    if len(data) != n_vec * dim * 4:
        raise ValueError(f"Truncated npy payload in {npy_path}")
    os.makedirs(os.path.dirname(f32_path) or ".", exist_ok=True)
    with open(f32_path, "wb") as out:
        out.write(HEADER.pack(MAGIC, n_vec, dim))
        out.write(data)


def export_tree(root: str) -> int:
    if not os.path.isdir(root):
        return 0
    converted = 0
    for dirpath, _, files in os.walk(root):
        for fname in files:
            if not fname.endswith(".npy"):
                continue
            npy_path = os.path.join(dirpath, fname)
            f32_path = npy_path[:-4] + ".f32emb"
            try:
                npy_to_f32emb(npy_path, f32_path)
                converted += 1
                print(f"Exported {npy_path} -> {f32_path}")
            except Exception as exc:
                print(f"WARN: could not export {npy_path}: {exc}", file=sys.stderr)
    return converted


def main() -> int:
    roots = [
        os.path.join(_FR, "faces_db"),
        os.path.join(_FR, "temp_faces_db"),
    ]
    total = sum(export_tree(r) for r in roots)
    print(f"Converted {total} embedding file(s) to .f32emb")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())