import os

import mujoco
import numpy as np

STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
MESH_FILE = os.path.join(STATIC_DIR, "so101_meshes.bin")
TEXTURED = {"cork": [0.80, 0.66, 0.48, 1.0]}
SKIP_PREFIX = ("wall",)


def export_meshes(model):
    os.makedirs(STATIC_DIR, exist_ok=True)
    manifest = []
    chunks = []
    offset = 0
    for i in range(model.nmesh):
        va, vn = model.mesh_vertadr[i], model.mesh_vertnum[i]
        fa, fn = model.mesh_faceadr[i], model.mesh_facenum[i]
        verts = model.mesh_vert[va : va + vn].astype(np.float32).tobytes()
        faces = model.mesh_face[fa : fa + fn].astype(np.uint16).tobytes()
        pad = (-len(faces)) % 4
        manifest.append(
            {"v_off": offset, "v_n": int(vn), "f_off": offset + len(verts), "f_n": int(fn)}
        )
        chunks += [verts, faces, b"\0" * pad]
        offset += len(verts) + len(faces) + pad
    with open(MESH_FILE, "wb") as f:
        f.write(b"".join(chunks))
    return manifest


def _rgba(model, g):
    mat = model.geom_matid[g]
    if mat >= 0:
        name = model.material(mat).name
        if name in TEXTURED:
            return TEXTURED[name]
        return [round(float(x), 3) for x in model.mat_rgba[mat]]
    return [round(float(x), 3) for x in model.geom_rgba[g]]


def scene_description(model):
    geoms = []
    bodies = set()
    for g in range(model.ngeom):
        group = model.geom_group[g]
        gtype = model.geom_type[g]
        name = model.geom(g).name
        if name.startswith(SKIP_PREFIX):
            continue
        if gtype == mujoco.mjtGeom.mjGEOM_MESH and group == 2:
            kind = "mesh"
        elif gtype == mujoco.mjtGeom.mjGEOM_BOX and group == 0:
            kind = "box"
        else:
            continue
        body = int(model.geom_bodyid[g])
        if body:
            bodies.add(body)
        geoms.append(
            {
                "kind": kind,
                "body": body,
                "mesh": int(model.geom_dataid[g]) if kind == "mesh" else -1,
                "size": [round(float(x), 5) for x in model.geom_size[g]],
                "pos": [round(float(x), 5) for x in model.geom_pos[g]],
                "quat": [round(float(x), 6) for x in model.geom_quat[g]],
                "rgba": _rgba(model, g),
            }
        )
    return geoms, sorted(bodies)


def body_poses(data, bodies):
    out = []
    for b in bodies:
        out += [round(float(x), 5) for x in data.xpos[b]]
        out += [round(float(x), 5) for x in data.xquat[b]]
    return out


def reset_message(run, model, data, manifest, mesh_url, label=""):
    geoms, bodies = scene_description(model)
    return {
        "kind": "reset",
        "run": run,
        "geoms": geoms,
        "bodies": bodies,
        "meshes": manifest,
        "mesh_url": mesh_url,
        "poses": [body_poses(data, bodies)],
        "label": label,
    }


def frames_message(run, start, poses, labels, done=False):
    return {"kind": "frames", "run": run, "start": start, "poses": poses, "labels": labels, "done": done}


HTML = """
<div class="f3a-wrap">
  <div class="f3a-view"></div>
  <div class="f3a-hud"><span class="f3a-label">loading 3D view</span></div>
  <div class="f3a-bar">
    <button class="f3a-replay" type="button">Replay</button>
    <span class="f3a-hint">drag to orbit, scroll to zoom, right-drag to pan</span>
  </div>
</div>
"""

CSS = """
.f3a-wrap { position: relative; width: 100%; height: 520px; border-radius: 8px; overflow: hidden; background: #d9dadd; }
.f3a-view { position: absolute; inset: 0; }
.f3a-view canvas { display: block; width: 100%; height: 100%; }
.f3a-hud { position: absolute; left: 10px; top: 8px; right: 10px; pointer-events: none; font: 13px/1.4 ui-monospace, monospace; color: #111; }
.f3a-label { background: rgba(255,255,255,0.8); padding: 2px 6px; border-radius: 4px; }
.f3a-bar { position: absolute; left: 10px; bottom: 8px; display: flex; gap: 10px; align-items: center; font: 12px ui-sans-serif, sans-serif; color: #333; }
.f3a-replay { padding: 3px 10px; border-radius: 4px; border: 1px solid #888; background: #fff; cursor: pointer; }
"""

JS = """
const view = element.querySelector('.f3a-view');
const labelEl = element.querySelector('.f3a-label');
const st = { ready: false, run: null, frames: [], labels: [], shown: -1, playT0: 0, done: false, msgs: 0, bodyGroups: {}, meshCache: {}, meshUrl: null };
element.f3aState = st;
let THREE = null;
let renderer, scene, camera, controls, root, dynamicRoot;

function applyPose(flat) {
  if (!flat) return;
  const ids = st.bodies || [];
  for (let i = 0; i < ids.length; i++) {
    const g = st.bodyGroups[ids[i]];
    if (!g) continue;
    const o = i * 7;
    g.position.set(flat[o], flat[o + 1], flat[o + 2]);
    g.quaternion.set(flat[o + 4], flat[o + 5], flat[o + 6], flat[o + 3]);
  }
}

async function loadMeshes(url, manifest) {
  if (st.meshUrl === url) return;
  const buf = await (await fetch(url)).arrayBuffer();
  st.meshCache = {};
  manifest.forEach((m, i) => {
    const pos = new Float32Array(buf, m.v_off, m.v_n * 3);
    const idx = new Uint16Array(buf, m.f_off, m.f_n * 3);
    const geo = new THREE.BufferGeometry();
    geo.setAttribute('position', new THREE.BufferAttribute(pos, 3));
    geo.setIndex(new THREE.BufferAttribute(idx, 1));
    geo.computeVertexNormals();
    st.meshCache[i] = geo;
  });
  st.meshUrl = url;
}

async function buildScene(msg) {
  await loadMeshes(msg.mesh_url, msg.meshes);
  if (dynamicRoot) { root.remove(dynamicRoot); }
  dynamicRoot = new THREE.Group();
  root.add(dynamicRoot);
  st.bodyGroups = {};
  st.bodies = msg.bodies;
  for (const b of msg.bodies) {
    const g = new THREE.Group();
    st.bodyGroups[b] = g;
    dynamicRoot.add(g);
  }
  for (const gd of msg.geoms) {
    let geo;
    if (gd.kind === 'mesh') geo = st.meshCache[gd.mesh];
    else geo = new THREE.BoxGeometry(gd.size[0] * 2, gd.size[1] * 2, gd.size[2] * 2);
    if (!geo) continue;
    const c = gd.rgba;
    const mat = new THREE.MeshStandardMaterial({ color: new THREE.Color(c[0], c[1], c[2]), roughness: 0.75, metalness: 0.05 });
    const mesh = new THREE.Mesh(geo, mat);
    mesh.position.set(gd.pos[0], gd.pos[1], gd.pos[2]);
    mesh.quaternion.set(gd.quat[1], gd.quat[2], gd.quat[3], gd.quat[0]);
    mesh.castShadow = true;
    mesh.receiveShadow = true;
    (gd.body ? st.bodyGroups[gd.body] : dynamicRoot).add(mesh);
  }
  applyPose(msg.poses[0]);
}

function handle(msg) {
  if (!msg || !st.ready) return;
  if (msg.kind === 'reset') {
    st.run = msg.run; st.frames = []; st.labels = []; st.shown = -1; st.playT0 = performance.now(); st.done = false;
    st.building = true;
    buildScene(msg).then(() => { st.building = false; labelEl.textContent = msg.label || 'ready'; });
    return;
  }
  st.msgs += 1;
  if (msg.kind === 'frames' && msg.run === st.run) {
    msg.poses.forEach((p, i) => { st.frames[msg.start + i] = p; st.labels[msg.start + i] = msg.labels[i]; });
    if (msg.done) st.done = true;
  }
}

function available() {
  let n = st.shown;
  while (st.frames[n + 1]) n += 1;
  return n;
}

function tick() {
  requestAnimationFrame(tick);
  if (!renderer) return;
  const now = performance.now();
  if (!st.building) {
    const avail = available();
    const due = Math.floor((now - st.playT0) * 0.03);
    if (due > avail) st.playT0 = now - (avail + 1) / 0.03;
    const target = Math.min(due, avail);
    if (target > st.shown) {
      applyPose(st.frames[target]);
      labelEl.textContent = st.labels[target] || '';
      st.shown = target;
    }
  }
  controls.update();
  renderer.render(scene, camera);
}

function resize() {
  if (!renderer) return;
  const w = view.clientWidth || 640, h = view.clientHeight || 480;
  renderer.setSize(w, h, false);
  camera.aspect = w / h;
  camera.updateProjectionMatrix();
}

element.querySelector('.f3a-replay').addEventListener('click', () => { st.shown = -1; st.playT0 = performance.now(); });

(async () => {
  THREE = await import('https://esm.sh/three@0.160.0');
  const { OrbitControls } = await import('https://esm.sh/three@0.160.0/examples/jsm/controls/OrbitControls.js');
  renderer = new THREE.WebGLRenderer({ antialias: true });
  renderer.setPixelRatio(Math.min(window.devicePixelRatio || 1, 2));
  renderer.shadowMap.enabled = true;
  renderer.shadowMap.type = THREE.PCFSoftShadowMap;
  view.appendChild(renderer.domElement);
  scene = new THREE.Scene();
  scene.background = new THREE.Color(0xd9dadd);
  camera = new THREE.PerspectiveCamera(45, 1, 0.01, 20);
  camera.position.set(0.62, 0.45, 0.42);
  controls = new OrbitControls(camera, renderer.domElement);
  controls.target.set(0.14, 0.05, -0.03);
  controls.enableDamping = true;
  scene.add(new THREE.HemisphereLight(0xffffff, 0x8a8a8a, 1.6));
  const sun = new THREE.DirectionalLight(0xffffff, 1.6);
  sun.position.set(0.6, 1.4, 0.5);
  sun.castShadow = true;
  sun.shadow.mapSize.set(2048, 2048);
  const sc = sun.shadow.camera; sc.left = -0.6; sc.right = 0.6; sc.top = 0.6; sc.bottom = -0.6; sc.near = 0.1; sc.far = 4;
  sun.shadow.bias = -0.0005;
  scene.add(sun);
  root = new THREE.Group();
  root.rotation.x = -Math.PI / 2;
  scene.add(root);
  resize();
  new ResizeObserver(resize).observe(view);
  st.ready = true;
  handle(props.value);
  requestAnimationFrame(tick);
})().catch((e) => { labelEl.textContent = '3D view failed: ' + e; });

watch('value', () => handle(props.value));
"""