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/aframe-event-set-component@5.0.0/dist/aframe-event-set-component.min.js"></script>
    <script>
    let score = 0;
    let selectedEntity = null;
    let velocity = {x: 0, y: 0, z: 0};
    let maxSpeed = 0.5;
    let accel = 0.02;

    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; // y=12 for center view
      var lowHeight = 3; // y=3 for closer views
      var quadrantOffset = Math.max(gw, gh) / 4; // 2 units to center on quadrants
      var pos, rot;

      rot = {x: -90, y: 0, z: 0}; // Always top-down view

      switch(view) {
        case 'center':
          pos = {x: center.x, y: height, z: center.z};
          break;
        case 'left':
          pos = {x: -quadrantOffset, y: lowHeight, z: center.z}; // Center on left side
          break;
        case 'right':
          pos = {x: quadrantOffset, y: lowHeight, z: center.z}; // Center on right side
          break;
        case 'front':
          pos = {x: center.x, y: lowHeight, z: -quadrantOffset}; // Center on top side
          break;
        case 'back':
          pos = {x: center.x, y: lowHeight, z: quadrantOffset}; // Center on bottom side
          break;
        case 'angle1':
          pos = {x: -quadrantOffset, y: lowHeight, z: -quadrantOffset}; // Top-left quadrant
          break;
        case 'angle2':
          pos = {x: quadrantOffset, y: lowHeight, z: -quadrantOffset}; // Top-right quadrant
          break;
        case 'angle3':
          pos = {x: -quadrantOffset, y: lowHeight, z: quadrantOffset}; // Bottom-left quadrant
          break;
        case 'angle4':
          pos = {x: quadrantOffset, y: lowHeight, z: quadrantOffset}; // Bottom-right quadrant
          break;
      }
      rig.setAttribute('position', pos);
      rig.setAttribute('rotation', rot);
    }

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

      switch(action) {
        case 'forward':
          var cam = document.querySelector('[camera]');
          var dir = new THREE.Vector3();
          cam.object3D.getWorldDirection(dir); // Raycast direction (downward for top-down view)
          var pos = rig.getAttribute('position');
          pos.x += dir.x * 0.1;
          pos.y += dir.y * 0.1;
          pos.z += dir.z * 0.1;
          rig.setAttribute('position', pos);
          break;
        case 'backward':
          velocity.z = Math.min(velocity.z + accel, maxSpeed);
          break;
        case 'stop':
          velocity = {x: 0, y: 0, z: 0};
          moveCamera('center');
          break;
        case 'rollCCW':
          rot.z += rSpeed;
          break;
        case 'rollCW':
          rot.z -= rSpeed;
          break;
      }
      rig.setAttribute('rotation', rot);
    }

    function updateCameraPosition() {
      var rig = document.querySelector('#rig');
      var pos = rig.getAttribute('position');
      pos.x += velocity.x;
      pos.y += velocity.y;
      pos.z += velocity.z;
      rig.setAttribute('position', pos);
      requestAnimationFrame(updateCameraPosition);
    }

    requestAnimationFrame(updateCameraPosition);

    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('center'); 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 's': flyCamera('stop');     break;
        case 'q': flyCamera('rollCCW');  break;
        case 'e': flyCamera('rollCW');   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:
        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:
        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>')
    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 Quadrant** 📷")
        cols = st.columns(3)
        cols[0].button("↖️ Top-Left", on_click=lambda: st.session_state.update({'camera_view': 'angle1'}))
        cols[1].button("⬆️ Top", on_click=lambda: st.session_state.update({'camera_view': 'front'}))
        cols[2].button("↗️ Top-Right", on_click=lambda: st.session_state.update({'camera_view': 'angle2'}))
        cols = st.columns(3)
        cols[0].button("⬅️ Left", on_click=lambda: st.session_state.update({'camera_view': 'left'}))
        cols[1].button("🔳 Center", on_click=lambda: st.session_state.update({'camera_view': 'center'}))
        cols[2].button("➡️ Right", on_click=lambda: st.session_state.update({'camera_view': 'right'}))
        cols = st.columns(3)
        cols[0].button("↙️ Bottom-Left", on_click=lambda: st.session_state.update({'camera_view': 'angle3'}))
        cols[1].button("⬇️ Bottom", on_click=lambda: st.session_state.update({'camera_view': 'back'}))
        cols[2].button("↘️ Bottom-Right", on_click=lambda: st.session_state.update({'camera_view': 'angle4'}))

        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', 'center')
    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/aframe-extras@6.1.1/dist/aframe-extras.loaders.min.js"></script>'

    st.components.v1.html(
      load_aframe_and_extras() + loader + scene,
      height=1000
    )

if __name__ == "__main__":
    main()