app.py

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

@st.cache_data
def load_aframe_and_extras():
    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', {
      init: function () {
        this.el.setAttribute('class', 'raycastable');
        this.el.setAttribute('cursor-listener', '');
        this.dragHandler = this.dragMove.bind(this);
        this.el.sceneEl.addEventListener('mousemove', this.dragHandler);
        this.el.addEventListener('mousedown', this.onDragStart.bind(this));
        this.el.addEventListener('mouseup', this.onDragEnd.bind(this));
        this.camera = document.querySelector('[camera]');
      },
      remove: function () {
        this.el.removeAttribute('cursor-listener');
        this.el.sceneEl.removeEventListener('mousemove', this.dragHandler);
      },
      onDragStart: function (evt) {
        this.isDragging = true;
        this.el.emit('dragstart');
      },
      onDragEnd: function (evt) {
        this.isDragging = false;
        this.el.emit('dragend');
      },
      dragMove: function (evt) {
        if (!this.isDragging) return;
        var camera = this.camera;
        var vector = new THREE.Vector3(
          evt.clientX / window.innerWidth * 2 - 1,
          -(evt.clientY / window.innerHeight) * 2 + 1,
          0.5
        );
        vector.unproject(camera);
        var dir = vector.sub(camera.position).normalize();
        var distance = -camera.position.y / dir.y;
        var pos = camera.position.clone().add(dir.multiplyScalar(distance));
        this.el.setAttribute('position', pos);
      }
    });

    AFRAME.registerComponent('bouncing', {
      schema: {
        speed: {type: 'vec3', default: {x: 0.1, y: 0.1, z: 0.1}},
        dist:  {type: 'vec3', default: {x: 0.5, y: 0.5, z: 0.5}}
      },
      init: function () {
        this.originalPos = this.el.getAttribute('position');
        this.dir = {x:1,y:1,z:1};
      },
      tick: function (time, timeDelta) {
        var p = this.el.getAttribute('position');
        var s = this.data.speed, d = this.data.dist;
        ['x','y','z'].forEach(a => {
          p[a] += s[a]*this.dir[a]*(timeDelta/1000);
          if (Math.abs(p[a] - this.originalPos[a]) > d[a]) {
            this.dir[a] *= -1;
          }
        });
        this.el.setAttribute('position', p);
      },
      boost: function() {
        var s = this.data.speed;
        ['x','y','z'].forEach(a => s[a]*=1.5);
        this.data.speed = s;
        this.dir = { x:Math.random()>0.5?1:-1,
                     y:Math.random()>0.5?1:-1,
                     z:Math.random()>0.5?1:-1 };
      }
    });

    AFRAME.registerComponent('moving-light', {
      schema: {
        color: {type:'color', default:'#FFF'},
        speed: {type:'vec3', default:{x:0.1,y:0.1,z:0.1}},
        bounds:{type:'vec3', default:{x:5,y:5,z:5}}
      },
      init: function(){
        this.dir = {x:1,y:1,z:1};
        this.light = document.createElement('a-light');
        this.light.setAttribute('type','point');
        this.light.setAttribute('color', this.data.color);
        this.light.setAttribute('intensity','0.75');
        this.el.appendChild(this.light);
      },
      tick: function(time, dt){
        var p = this.el.getAttribute('position'),
            s = this.data.speed,
            b = this.data.bounds;
        ['x','y','z'].forEach(a=>{
          p[a] += s[a]*this.dir[a]*(dt/1000);
          if (Math.abs(p[a])>b[a]) this.dir[a]*=-1;
        });
        this.el.setAttribute('position', p);
      }
    });

    function moveCamera(direction) {
      var rig = document.querySelector('#rig');
      var pos = rig.getAttribute('position');
      var rot = rig.getAttribute('rotation');
      var speed = 0.5, rSpeed = 5;
      switch(direction) {
        case 'up': pos.y += speed; break;
        case 'down': pos.y -= speed; break;
        case 'forward': pos.z -= speed; break;
        case 'left': pos.x -= speed; break;
        case 'right': pos.x += speed; break;
        case 'rotateLeft': rot.y += rSpeed; break;
        case 'rotateRight': rot.y -= rSpeed; break;
        case 'reset': pos = {x:0,y:10,z:0}; rot = {x:-90,y:0,z:0}; break;
        case 'ground': pos = {x:0,y:1.6,z:0}; rot = {x:0,y:0,z:0}; break;
      }
      rig.setAttribute('position', pos);
      rig.setAttribute('rotation', rot);
    }

    function fireRaycast() {
      var camera = document.querySelector('[camera]'),
          dir = new THREE.Vector3();
      camera.object3D.getWorldDirection(dir);
      var rc = new THREE.Raycaster();
      rc.set(camera.object3D.position, dir);
      var hits = rc.intersectObjects(
        document.querySelectorAll('.raycastable').map(e=>e.object3D), true
      );
      if (hits.length>0) {
        var el = hits[0].object.el;
        if (el.components.bouncing) {
          el.components.bouncing.boost();
          score += 10;
          document.getElementById('score').setAttribute('value','Score: '+score);
        }
      }
    }

    // —— Key remap: W→down, S→reset, X→up —— 
    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>
    """

@st.cache_data
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):
    """1×1×1 scale, spin around Y, draggable but no bouncing."""
    anim = 'animation="property: rotation; to: 0 360 0; loop: true; dur: 20000; easing: linear"'
    if file_type == 'obj':
        return (
            f'<a-entity obj-model="obj: #{stem}" '
            f'position="{position}" rotation="0 0 0" scale="1 1 1" '
            f'class="raycastable" draggable {anim}></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'class="raycastable" draggable {anim}></a-entity>'
        )
    return ""

@st.cache_data
def generate_tilemap(files, directory, gw, gh):
    img_exts   = ['webp','png']
    model_exts = ['obj','glb','gltf']
    vid_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 vid_exts]

    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}" '
              f'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 vid_exts:
            assets += (
              f'<video id="{stem}" '
              f'src="data:video/mp4;base64,{data}" '
              'loop="true" autoplay="true" muted="true"></video>'
            )
    assets += "</a-assets>"

    entities = ""
    sx = -gw/2
    sz = -gh/2

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

        # 1) ground image
        if img_files:
          img = img_files[(i*gh + j) % len(img_files)]
          stem=Path(img).stem
          entities += (
            f'<a-plane src="#{stem}" width="1" height="1" '
            f'rotation="-90 0 0" position="{x} 0.01 {z}"></a-plane>'
          )

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

        # 3) video layer
        if vid_files:
          vid = vid_files[(i*gh + j) % len(vid_files)]
          stem = Path(vid).stem
          entities += (
            f'<a-video src="#{stem}" width="1" height="1" '
            f'rotation="-90 0 0" position="{x} 0.2 {z}" '
            'class="raycastable" draggable></a-video>'
          )

    return assets, entities

def main():
    st.set_page_config(layout="wide")
    with st.sidebar:
        st.markdown("### 🤖 3D AI Using Claude 3.5 Sonnet for AI Pair Programming")
        st.markdown(
          "[Open 3D Animation Toolkit]"
          "(https://huggingface.co/spaces/awacke1/3d_animation_toolkit)",
          unsafe_allow_html=True
        )
        st.markdown("### ⬆️ Upload")
        uploaded_files = st.file_uploader("Add files:", accept_multiple_files=True, key="file_uploader")

        st.markdown("### 🎮 Camera Controls")
        col1, col2, col3 = st.columns(3)
        with col1:
            st.button("⬅️", on_click=lambda: st.session_state.update({'camera_move':'left'}))
            st.button("🔄↺", on_click=lambda: st.session_state.update({'camera_move':'rotateLeft'}))
            st.button("🔝", on_click=lambda: st.session_state.update({'camera_move':'reset'}))
        with col2:
            st.button("⬆️", on_click=lambda: st.session_state.update({'camera_move':'up'}))
            st.button("👀", on_click=lambda: st.session_state.update({'camera_move':'ground'}))
            st.button("🔫", on_click=lambda: st.session_state.update({'camera_move':'fire'}))
        with col3:
            st.button("➡️", on_click=lambda: st.session_state.update({'camera_move':'right'}))
            st.button("⬇️", on_click=lambda: st.session_state.update({'camera_move':'down'}))
            st.button("⏩", on_click=lambda: st.session_state.update({'camera_move':'forward'}))

        st.markdown("### 🗺️ Grid Size")
        grid_width  = st.slider("Grid Width", 1, 8, 8)
        grid_height = st.slider("Grid Height",1, 5, 5)

        st.markdown("### 📁 Directory")
        directory = st.text_input("Enter path:", ".", key="directory_input")

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

    file_types = ['obj','glb','gltf','webp','png','mp4']
    if uploaded_files:
        for up in uploaded_files:
            ext = Path(up.name).suffix.lower()[1:]
            if ext in file_types:
                with open(os.path.join(directory, up.name),"wb") as f:
                    shutil.copyfileobj(up, f)
                st.sidebar.success(f"Uploaded: {up.name}")
            else:
                st.sidebar.warning(f"Skipped unsupported: {up.name}")

    files = [f for f in os.listdir(directory) if f.split('.')[-1] in file_types]

    # Build A-Frame scene shell
    aframe_scene = f"""
    <a-scene embedded style="height:600px; width:100%;">
      <a-entity id="rig" position="0 {max(grid_width,grid_height)} 0" rotation="-90 0 0">
        <a-camera fov="60" look-controls 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, grid_width, grid_height)
    aframe_scene += assets + entities + "</a-scene>"

    # Apply camera_move if any
    cam = st.session_state.get('camera_move')
    if cam:
        if cam == 'fire':
            aframe_scene += "<script>fireRaycast();</script>"
        else:
            aframe_scene += f"<script>moveCamera('{cam}');</script>"
        st.session_state.pop('camera_move')

    # Loader for OBJ/glTF
    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 + aframe_scene,
      height=630
    )

if __name__ == "__main__":
    main()