import base64
import os
import numpy as np
import streamlit.components.v1 as components
from pathlib import Path

def _make_component():
    space_id = os.environ.get("SPACE_ID", "")
    if space_id:
        slug = space_id.replace("/", "-").lower()
        url  = f"https://{slug}.hf.space/app/static/webgpu_component.html"
        return components.declare_component("webgpu_embedder", url=url)
    # Local dev: serve from frontend/ folder
    return components.declare_component(
        "webgpu_embedder",
        path=str(Path(__file__).parent / "frontend"),
    )

_component = _make_component()

def webgpu_embed(papers: list[dict], run: bool, key: str = None):
    result = _component(papers=papers, run=run, key=key, default=None)
    if result is None:
        return None
    if "error" in result:
        raise RuntimeError(result["error"])
    data = base64.b64decode(result["embeddings_b64"])
    arr  = np.frombuffer(data, dtype=np.float32).reshape(result["n_papers"], result["n_dims"])
    return arr.copy()