app.py

import streamlit as st
import os, base64, shutil, random
from pathlib import Path

@st.cache_data
def load_aframe_and_extras():
    # ↔️ This is your original A-Frame + components + moveCamera/fireRaycast boilerplate,
    #     with just the keydown mapping at the bottom tweaked to:
    #       W → down
    #       S → reset
    #       X → up
    return """
    <script src="https://aframe.io/releases/1.2.0/aframe.min.js"></script>
    <script src="https://unpkg.com/[email protected]/dist/aframe-event-set-component.min.js"></script>
    <script>
    let score = 0;
    AFRAME.registerComponent('draggable', { /* … your original code … */ });
    AFRAME.registerComponent('bouncing', { /* … your original code … */ });
    AFRAME.registerComponent('moving-light', { /* … your original code … */ });

    function moveCamera(direction) { /* … your original code … */ }
    function fireRaycast() { /* … your original code … */ }

    // —— Modified keys here ——
    document.addEventListener('keydown', function(event) {
      switch(event.key.toLowerCase()) {
        case 'w': moveCamera('down');    break;
        case 's': moveCamera('reset');   break;
        case 'x': moveCamera('up');      break;
        case 'q': moveCamera('rotateLeft');  break;
        case 'e': moveCamera('rotateRight'); break;
        case 'z': moveCamera('reset');      break;
        case 'c': moveCamera('ground');     break;
        case ' ': fireRaycast();            break;
      }
    });
    </script>
    """

def encode_file(file_path):
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode()

def create_aframe_entity(stem, file_type, position):
    """Only 3D models spin in place, scale=1, no bouncing/draggable."""
    if file_type == 'obj':
        return (
            f'<a-entity obj-model="obj: #{stem}" '
            f'position="{position}" rotation="0 0 0" scale="1 1 1" '
            f'animation="property: rotation; to: 0 360 0; loop: true; dur: 20000; easing: linear">'
            '</a-entity>'
        )
    if file_type in ('glb','gltf'):
        return (
            f'<a-entity gltf-model="#{stem}" '
            f'position="{position}" rotation="0 0 0" scale="1 1 1" '
            f'animation="property: rotation; to: 0 360 0; loop: true; dur: 20000; easing: linear">'
            '</a-entity>'
        )
    return ""  # images/videos handled separately

@st.cache_data
def generate_tilemap(files, directory, gw, gh):
    """
    - Encodes *all* your files into <a-assets>.
    - Then for each cell:
        1. Chooses a random PNG/WEBP → flat a-plane (ground).
        2. Chooses a random OBJ/GLB → spinning 3D entity above.
        3. If any MP4, chooses one → looping video plane just above ground.
    """
    img_exts = ('webp','png')
    model_exts = ('obj','glb','gltf')
    video_exts = ('mp4',)

    img_files   = [f for f in files if f.split('.')[-1] in img_exts]
    model_files = [f for f in files if f.split('.')[-1] in model_exts]
    vid_files   = [f for f in files if f.split('.')[-1] in video_exts]

    # --- Build assets ---
    assets = "<a-assets>"
    for f in files:
        stem = Path(f).stem
        ext  = f.split('.')[-1]
        data = encode_file(os.path.join(directory,f))
        if ext in model_exts:
            assets += f'<a-asset-item id="{stem}" src="data:application/octet-stream;base64,{data}"></a-asset-item>'
        elif ext in img_exts:
            assets += f'<img id="{stem}" src="data:image/{ext};base64,{data}">'
        elif ext in video_exts:
            assets += (
              f'<video id="{stem}" src="data:video/mp4;base64,{data}" '
              'loop="true" autoplay="true" muted="true"></video>'
            )
    assets += "</a-assets>"

    # --- Spawn per-cell entities ---
    entities = ""
    tile = 1
    sx = -(gw*tile)/2
    sz = -(gh*tile)/2

    for i in range(gw):
      for j in range(gh):
        x = sx + i*tile
        z = sz + j*tile

        # 1) image-ground
        if img_files:
          img = random.choice(img_files)
          stem = Path(img).stem
          entities += (
            f'<a-plane src="#{stem}" '
            f'width="{tile}" height="{tile}" '
            f'rotation="-90 0 0" position="{x} 0.01 {z}"></a-plane>'
          )

        # 2) spinning 3D model
        if model_files:
          mdl = random.choice(model_files)
          ext = mdl.split('.')[-1]
          stem = Path(mdl).stem
          entities += create_aframe_entity(stem, ext, f"{x} 0.5 {z}")

        # 3) video-layer (optional)
        if vid_files:
          vid = random.choice(vid_files)
          stem = Path(vid).stem
          entities += (
            f'<a-video src="#{stem}" '
            f'width="{tile}" height="{tile}" '
            f'rotation="-90 0 0" position="{x} 0.2 {z}" autoplay loop muted></a-video>'
          )

    return assets, entities

def main():
    st.set_page_config(layout="wide")
    with st.sidebar:
        st.markdown("### 🤖 3D AI …")
        uploaded = st.file_uploader("Add files:", accept_multiple_files=True)
        st.markdown("### 🎮 Camera")
        cols = st.columns(3)
        for btn,dir in zip(("⬅️","🔄↺","🔝"),("left","rotateLeft","reset")):
            cols[0 if btn in ("⬅️","🔄↺","🔝") else 1 if ... else 2].button(
                btn, on_click=lambda d=dir: st.session_state.update({'camera_move':d})
            )
        # … your original camera buttons …
        st.markdown("### 🗺️ Grid Size")
        gw = st.slider("Width", 1, 8, 8)
        gh = st.slider("Height",1, 5, 5)
        st.markdown("### 📁 Dir")
        directory = st.text_input("Path:", ".", key="dir")

    if not os.path.isdir(directory):
        st.sidebar.error("Invalid directory")
        return

    types = ('obj','glb','gltf','webp','png','mp4')
    files = [f for f in os.listdir(directory) if f.split('.')[-1] in types]

    # Build the <a-scene> shell
    scene = f"""
    <a-scene embedded style="height:600px; width:100%;">
      <a-entity id="rig" position="0 {max(gw,gh)} 0" rotation="-90 0 0">
        <a-camera fov="60" look-controls wasd-controls="enabled:false"
                  cursor="rayOrigin:mouse" raycaster="objects:.raycastable">
        </a-camera>
      </a-entity>
      <a-sky color="#87CEEB"></a-sky>
      <a-entity moving-light="color:#FFD700; speed:0.07 0.05 0.06; bounds:4 3 4" position="2 2 -2"></a-entity>
      <a-entity moving-light="color:#FF6347; speed:0.06 0.08 0.05; bounds:4 3 4" position="-2 1 2"></a-entity>
      <a-entity moving-light="color:#00CED1; speed:0.05 0.06 0.07; bounds:4 3 4" position="0 3 0"></a-entity>
      <a-text id="score" value="Score: 0" position="-1.5 1 -2" scale="0.5 0.5 0.5" color="white"></a-text>
    """

    assets, entities = generate_tilemap(files, directory, gw, gh)
    scene += assets + entities + "</a-scene>"

    # camera_move injection
    cm = st.session_state.get('camera_move', None)
    if cm:
        scene += f"<script>moveCamera('{cm}');</script>"
        st.session_state.pop('camera_move')

    # ────────────────────────────────────────────────────────────────────────────
    # Make sure OBJ/GLB loader is present *after* everything else
    loader = '<script src="https://cdn.jsdelivr.net/npm/[email protected]/dist/aframe-extras.loaders.min.js"></script>'
    st.components.v1.html(
      load_aframe_and_extras() + loader + scene,
      height=620
    )

if __name__ == "__main__":
    main()