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 src="https://unpkg.com/[email protected]/dist/aframe-look-at-component.min.js"></script>
    <script>
    let score = 0;
    let selectedEntity = null;

    AFRAME.registerComponent('draggable', {
      init: function () {
        this.el.addEventListener('mousedown', this.onMouseDown.bind(this));
        this.el.addEventListener('touchstart', this.onTouchStart.bind(this));
        this.el.addEventListener('mouseup', this.onMouseUp.bind(this));
        this.el.addEventListener('touchend', this.onTouchEnd.bind(this));
        this.el.addEventListener('mousemove', this.onMouseMove.bind(this));
        this.el.addEventListener('touchmove', this.onTouchMove.bind(this));
      },
      onMouseDown: function (evt) {
        selectedEntity = this.el;
      },
      onTouchStart: function (evt) {
        evt.preventDefault();
        selectedEntity = this.el;
      },
      onMouseUp: function () {
        selectedEntity = null;
      },
      onTouchEnd: function () {
        selectedEntity = null;
      },
      onMouseMove: function (evt) {
        if (selectedEntity) {
          let cam = document.querySelector('[camera]');
          let dir = new THREE.Vector3();
          cam.object3D.getWorldDirection(dir);
          let rc = new THREE.Raycaster(cam.object3D.position, dir);
          let plane = new THREE.Plane(new THREE.Vector3(0, 1, 0), 0);
          let pos = new THREE.Vector3();
          rc.ray.intersectPlane(plane, pos);
          selectedEntity.setAttribute('position', {x: pos.x, y: 0, z: pos.z});
        }
      },
      onTouchMove: function (evt) {
        evt.preventDefault();
        if (selectedEntity) {
          let cam = document.querySelector('[camera]');
          let dir = new THREE.Vector3();
          cam.object3D.getWorldDirection(dir);
          let rc = new THREE.Raycaster(cam.object3D.position, dir);
          let plane = new THREE.Plane(new THREE.Vector3(0, 1, 0), 0);
          let pos = new THREE.Vector3();
          rc.ray.intersectPlane(plane, pos);
          selectedEntity.setAttribute('position', {x: pos.x, y: 0, z: pos.z});
        }
      }
    });

    AFRAME.registerComponent('bouncing', {/* … */});
    AFRAME.registerComponent('moving-light', {/* … */});

    function moveCamera(view) {
      var rig = document.querySelector('#rig');
      var gw = 8, gh = 8;
      var center = {x: 0, y: 0, z: 0};
      var height = Math.max(gw, gh) * 1.5;
      var distance = Math.max(gw, gh) * 2; // Further increased distance to ensure the entire tilemap is in view
      var pos;

      switch(view) {
        case 'top':
          pos = {x: center.x, y: height, z: center.z};
          rig.setAttribute('position', pos);
          rig.setAttribute('rotation', {x: -90, y: 0, z: 0});
          break;
        case 'left':
          pos = {x: -distance, y: height, z: center.z};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'right':
          pos = {x: distance, y: height, z: center.z};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'front':
          pos = {x: center.x, y: height, z: -distance};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'back':
          pos = {x: center.x, y: height, z: distance};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'angle1':
          pos = {x: -distance, y: height, z: -distance};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'angle2':
          pos = {x: distance, y: height, z: -distance};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'angle3':
          pos = {x: -distance, y: height, z: distance};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
        case 'angle4':
          pos = {x: distance, y: height, z: distance};
          rig.setAttribute('position', pos);
          rig.setAttribute('look-at', center);
          break;
      }
    }

    function flyCamera(action) {
      var rig = document.querySelector('#rig');
      var pos = rig.getAttribute('position');
      var rot = rig.getAttribute('rotation');
      var speed = 0.1, rSpeed = 2;

      switch(action) {
        case 'forward':
          var dir = new THREE.Vector3(0, 0, -1);
          dir.applyQuaternion(rig.object3D.quaternion);
          pos.x += dir.x * speed;
          pos.y += dir.y * speed;
          pos.z += dir.z * speed;
          break;
        case 'backward':
          var dir = new THREE.Vector3(0, 0, 1);
          dir.applyQuaternion(rig.object3D.quaternion);
          pos.x += dir.x * speed;
          pos.y += dir.y * speed;
          pos.z += dir.z * speed;
          break;
        case 'panLeft':
          rot.y += rSpeed;
          break;
        case 'panRight':
          rot.y -= rSpeed;
          break;
        case 'rotateLeft':
          rot.y += rSpeed;
          break;
        case 'rotateRight':
          rot.y -= rSpeed;
          break;
        case 'reset':
          moveCamera('top');
          return;
      }
      rig.setAttribute('position', pos);
      rig.setAttribute('rotation', rot);
    }

    function fireRaycast() {
      var cam = document.querySelector('[camera]');
      var dir = new THREE.Vector3(); cam.object3D.getWorldDirection(dir);
      var rc = new THREE.Raycaster(cam.object3D.position, dir);
      var hits = rc.intersectObjects(
        Array.from(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);
        }
      }
    }

    document.addEventListener('keydown', e => {
      switch(e.key.toLowerCase()){
        case '1': moveCamera('top');    break;
        case '2': moveCamera('left');   break;
        case '3': moveCamera('right');  break;
        case '4': moveCamera('front');  break;
        case '5': moveCamera('back');   break;
        case '6': moveCamera('angle1'); break;
        case '7': moveCamera('angle2'); break;
        case '8': moveCamera('angle3'); break;
        case '9': moveCamera('angle4'); break;
        case 'w': flyCamera('forward');  break;
        case 'x': flyCamera('backward'); break;
        case 'a': flyCamera('panLeft');  break;
        case 'd': flyCamera('panRight'); break;
        case 'q': flyCamera('rotateLeft'); break;
        case 'e': flyCamera('rotateRight'); break;
        case 's': flyCamera('reset');    break;
        case ' ': fireRaycast();         break;
      }
    });
    </script>
    """

@st.cache_data
def encode_file(path):
    with open(path,'rb') as f: return base64.b64encode(f.read()).decode()

def create_aframe_entity(stem, ext, pos):
    ry = random.uniform(0,360)
    if ext == 'obj':
        return (f'<a-entity obj-model="obj:#{stem}" '
                f'position="{pos}" rotation="0 {ry} 0" scale="1 1 1" '
                'class="raycastable" draggable></a-entity>')
    if ext in ('glb','gltf'):
        return (f'<a-entity gltf-model="#{stem}" '
                f'position="{pos}" rotation="0 {ry} 0" scale="1 1 1" '
                'class="raycastable" draggable></a-entity>')
    return ''

@st.cache_data
def generate_tilemap(files, dirpath, gw=8, gh=8):
    img_exts   = ['webp','png','jpeg','jpg']
    model_exts = ['obj','glb','gltf']
    vid_exts   = ['mp4']

    imgs   = [f for f in files if f.split('.')[-1] in img_exts]
    models = [f for f in files if f.split('.')[-1] in model_exts]
    vids   = [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(dirpath,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 autoplay muted></video>')
    assets += "</a-assets>"

    entities = ""
    if vids:
        # Use the first video as a large ground plane
        v = vids[0]; s = Path(v).stem
        entities += (f'<a-video src="#{s}" width="{gw}" height="{gh}" '
                     f'rotation="-90 0 0" position="0 0.05 0" '
                     'loop autoplay></a-video>')
    else:
        # Generate individual tiles if no video
        sx = -gw/2; sz = -gh/2
        for i in range(gw):
          for j in range(gh):
            x,y,z = sx+i,0,sz+j
            if imgs:
              img = random.choice(imgs); s = Path(img).stem
              entities += (f'<a-plane src="#{s}" width="1" height="1" '
                           f'rotation="-90 0 0" position="{x} 0 {z}"></a-plane>')
    # Always generate models
    sx = -gw/2; sz = -gh/2
    for i in range(gw):
      for j in range(gh):
        x,y,z = sx+i,0,sz+j
        if models:
          m = random.choice(models); ext = m.split('.')[-1]; s = Path(m).stem
          entities += create_aframe_entity(s, ext, f"{x} 0 {z}")
    return assets, entities

def main():
    st.set_page_config(layout="wide")
    with st.sidebar:
        st.markdown("### 🧭 Camera Views")
        st.markdown("**Select View** 📷")
        cols = st.columns(3)
        cols[0].button("🔝 Top", on_click=lambda: st.session_state.update({'camera_view': 'top'}))
        cols[1].button("⬅️ Left", on_click=lambda: st.session_state.update({'camera_view': 'left'}))
        cols[2].button("➡️ Right", on_click=lambda: st.session_state.update({'camera_view': 'right'}))
        cols = st.columns(3)
        cols[0].button("⬆️ Front", on_click=lambda: st.session_state.update({'camera_view': 'front'}))
        cols[1].button("⬇️ Back", on_click=lambda: st.session_state.update({'camera_view': 'back'}))
        cols[2].button("↖️ Angle 1", on_click=lambda: st.session_state.update({'camera_view': 'angle1'}))
        cols = st.columns(3)
        cols[0].button("↗️ Angle 2", on_click=lambda: st.session_state.update({'camera_view': 'angle2'}))
        cols[1].button("↙️ Angle 3", on_click=lambda: st.session_state.update({'camera_view': 'angle3'}))
        cols[2].button("↘️ Angle 4", on_click=lambda: st.session_state.update({'camera_view': 'angle4'}))

        st.markdown("### ✈️ Flight Controls")
        st.markdown("**Fly Camera** 🎮")
        cols = st.columns(3)
        cols[0].button("Q Rotate Left", on_click=lambda: st.session_state.update({'fly_action': 'rotateLeft'}))
        cols[1].button("W Forward", on_click=lambda: st.session_state.update({'fly_action': 'forward'}))
        cols[2].button("E Rotate Right", on_click=lambda: st.session_state.update({'fly_action': 'rotateRight'}))
        cols = st.columns(3)
        cols[0].button("A Pan Left", on_click=lambda: st.session_state.update({'fly_action': 'panLeft'}))
        cols[1].button("S Reset", on_click=lambda: st.session_state.update({'fly_action': 'reset'}))
        cols[2].button("D Pan Right", on_click=lambda: st.session_state.update({'fly_action': 'panRight'}))
        cols = st.columns(1)
        cols[0].button("X Backward", on_click=lambda: st.session_state.update({'fly_action': 'backward'}))

        st.markdown("### ➕ Add Media Files")
        ups = st.file_uploader("Add files (png, jpeg, obj, glb, mp4, etc.):", accept_multiple_files=True)
        st.markdown("### 📋 Uploaded Model Files")
        directory = st.text_input("Path:", ".", key="dir")
        if os.path.isdir(directory):
            files = [f for f in os.listdir(directory) if f.split('.')[-1] in ['obj', 'glb', 'gltf']]
            if files:
                for i, f in enumerate(files, 1):
                    st.markdown(f"{i}. {f}")
            else:
                st.markdown("No model files found.")

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

    types = ['obj','glb','gltf','webp','png','jpeg','jpg','mp4']
    if ups:
      for up in ups:
        ext=Path(up.name).suffix.lower()[1:]
        if ext in 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 {up.name}")

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

    spot_h = max(8,8)*1.5

    scene = f"""
    <a-scene embedded style="height:100vh; width:100vw; position:fixed; top:0; left:0;">
      <a-entity id="rig" position="0 12 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-entity light="type:spot; color:#FFF; intensity:1; angle:45; penumbra:0.2"
                position="0 {spot_h} 0" rotation="-90 0 0"></a-entity>
      <a-text id="score" value="Score:0" position="-1.5 2 -3" scale="0.5 0.5 0.5" color="white"></a-text>
    """

    assets, ents = generate_tilemap(files, directory, 8, 8)
    scene += assets + ents + "</a-scene>"

    view_cmd = st.session_state.get('camera_view', 'top')
    if view_cmd:
      scene += f"<script>moveCamera('{view_cmd}');</script>"
      st.session_state.pop('camera_view', None)

    fly_cmd = st.session_state.get('fly_action')
    if fly_cmd:
      scene += f"<script>flyCamera('{fly_cmd}');</script>"
      st.session_state.pop('fly_action', None)

    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=1000
    )

if __name__ == "__main__":
    main()