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

# 🚀 1. Load Aframe and Scripts: Let's get the party started with Aframe and friends! 🎉
@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 rotVelocity = {x: 0, y: 0, z: 0};
    let maxSpeed = 0.5;
    let accel = 0.02;
    let maxRotSpeed = 2;
    let rotAccel = 0.1;

    // 🖱️ 1.1 Draggable Component: Click and drag like a pro! 🏋️
    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});
        }
      }
    });

    // ⛹️ 1.2 Bouncing and Lights: Keep the scene lively! ✨
    AFRAME.registerComponent('bouncing', {/* … */});
    AFRAME.registerComponent('moving-light', {/* … */});

    // 📸 1.3 Move Camera: Snap to your favorite quadrant views! 🖼️
    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};
          break;
        case 'right':
          pos = {x: quadrantOffset, y: lowHeight, z: center.z};
          break;
        case 'front':
          pos = {x: center.x, y: lowHeight, z: -quadrantOffset};
          break;
        case 'back':
          pos = {x: center.x, y: lowHeight, z: quadrantOffset};
          break;
        case 'angle1':
          pos = {x: -quadrantOffset, y: lowHeight, z: -quadrantOffset};
          break;
        case 'angle2':
          pos = {x: quadrantOffset, y: lowHeight, z: -quadrantOffset};
          break;
        case 'angle3':
          pos = {x: -quadrantOffset, y: lowHeight, z: quadrantOffset};
          break;
        case 'angle4':
          pos = {x: quadrantOffset, y: lowHeight, z: quadrantOffset};
          break;
      }
      rig.setAttribute('position', pos);
      rig.setAttribute('rotation', rot);
    }

    // ✈️ 1.4 Fly Camera: Soar through the scene with style! 🛩️
    function flyCamera(action) {
      var rig = document.querySelector('#rig');
      var pos = rig.getAttribute('position');

      switch(action) {
        case 'moveLeft':
          var leftDir = new THREE.Vector3(-1, 0, 0);
          leftDir.applyQuaternion(rig.object3D.quaternion);
          pos.x += leftDir.x * 0.1;
          pos.y += leftDir.y * 0.1;
          pos.z += leftDir.z * 0.1;
          break;
        case 'moveRight':
          var rightDir = new THREE.Vector3(1, 0, 0);
          rightDir.applyQuaternion(rig.object3D.quaternion);
          pos.x += rightDir.x * 0.1;
          pos.y += rightDir.y * 0.1;
          pos.z += rightDir.z * 0.1;
          break;
        case 'rotateCCW':
          rotVelocity.y = Math.min(rotVelocity.y + rotAccel, maxRotSpeed);
          break;
        case 'rotateCW':
          rotVelocity.y = Math.max(rotVelocity.y - rotAccel, -maxRotSpeed);
          break;
        case 'pitchUp':
          rotVelocity.x = Math.max(rotVelocity.x - rotAccel, -maxRotSpeed);
          break;
        case 'pitchDown':
          rotVelocity.x = Math.min(rotVelocity.x + rotAccel, maxRotSpeed);
          break;
        case 'forward':
          var cam = document.querySelector('[camera]');
          var dir = new THREE.Vector3();
          cam.object3D.getWorldDirection(dir);
          pos.x += dir.x * 0.1;
          pos.y += dir.y * 0.1;
          pos.z += dir.z * 0.1;
          break;
        case 'backward':
          var cam = document.querySelector('[camera]');
          var dir = new THREE.Vector3();
          cam.object3D.getWorldDirection(dir);
          pos.x -= dir.x * 0.1;
          pos.y -= dir.y * 0.1;
          pos.z -= dir.z * 0.1;
          break;
        case 'stop':
          velocity = {x: 0, y: 0, z: 0};
          rotVelocity = {x: 0, y: 0, z: 0};
          moveCamera('center');
          break;
      }
      rig.setAttribute('position', pos);
    }

    // 🔄 1.5 Update Camera: Keep the camera flying smoothly! 🕊️
    function updateCameraPosition() {
      var rig = document.querySelector('#rig');
      var pos = rig.getAttribute('position');
      var rot = rig.getAttribute('rotation');

      pos.x += velocity.x;
      pos.y = Math.max(1, pos.y + velocity.y);
      pos.z += velocity.z;

      rot.x += rotVelocity.x;
      rot.y += rotVelocity.y;
      rot.z += rotVelocity.z;

      rig.setAttribute('position', pos);
      rig.setAttribute('rotation', rot);
      requestAnimationFrame(updateCameraPosition);
    }

    requestAnimationFrame(updateCameraPosition);

    // 🎯 1.6 Fire Raycast: Zap objects with a laser beam! ⚡
    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);
        }
      }
    }

    // ⌨️ 1.7 Keyboard Listener: Take control with your keys! 🎮
    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 'q': flyCamera('rotateCCW'); break;
        case 'e': flyCamera('rotateCW');  break;
        case 'z': flyCamera('pitchUp');   break;
        case 'c': flyCamera('pitchDown'); break;
        case 'a': flyCamera('moveLeft');  break;
        case 'd': flyCamera('moveRight'); break;
        case 'w': flyCamera('forward');   break;
        case 'x': flyCamera('backward');  break;
        case 's': flyCamera('stop');      break;
        case ' ': fireRaycast();          break;
      }
    });
    </script>
    """

# 🔐 2. Encode File: Sneakily encode files into base64! 🕵️
@st.cache_data
def encode_file(path):
    with open(path,'rb') as f: return base64.b64encode(f.read()).decode()

# 🏰 3. Create Aframe Entity: Build magical 3D models! 🪄
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 ''

# 🗺️ 4. Generate Tilemap: Craft the world one tile at a time! 🌍
@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]

    # 📦 4.1 Load Assets: Pack up all the goodies! 🎁
    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}" preload="auto" '
                       f'src="data:video/mp4;base64,{data}" '
                       'loop="true" autoplay="true" muted="true" playsinline webkit-playsinline></video>')
    assets += "</a-assets>"

    # 🏞️ 4.2 Build the Scene: Lay out the tiles and models! 🏡
    entities = ""
    model_counts = {f: 0 for f in models}  # Track spawned model counts 📊
    if vids:
        v = vids[0]; s = Path(v).stem
        # 📹 Adjust video dimensions to cover the tilemap while preserving aspect ratio 🎥
        # Assume a common video aspect ratio (16:9) and scale to fit 8x8 tilemap
        videoWidth = 8  # Match tilemap width
        videoHeight = 8  # Match tilemap height (square to fit tilemap)
        entities += (f'<a-video src="#{s}" width="{videoWidth}" height="{videoHeight}" '
                     f'rotation="-90 0 0" position="0 0.05 0" '
                     'material="shader:flat; src:#{s}" loop="true" autoplay="true"></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}")
          model_counts[m] += 1  # Increment count for this model 📈
    return assets, entities, model_counts

# 🎨 5. Main Function: Paint the canvas of your 3D world! 🖌️
def main():
    st.set_page_config(layout="wide")
    with st.sidebar:
        # 🧭 5.1 Camera Views UI: Pick your perfect angle! 📐
        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'}))

        # 📂 5.2 File Uploader: Drop your media treasures here! 💎
        st.markdown("### ➕ Add Media Files")
        ups = st.file_uploader(
            "🖼️ png/jpeg, 🏰 obj/glb/gltf, 📹 mp4:",
            accept_multiple_files=True
        )
        st.markdown("### 📋 Uploaded Model Files")
        directory = st.text_input("", ".", key="dir", placeholder="Enter directory path (e.g., '.')")
        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):
                    # Add emoji based on file type and show spawn count 🖌️
                    ext = f.split('.')[-1]
                    emoji = "🏰" if ext == 'obj' else "🏯"  # 🏰 for obj, 🏯 for glb/gltf
                    count = st.session_state.get('model_counts', {}).get(f, 0)
                    st.markdown(f"{i}. {emoji} {f} (Spawned: {count})")
            else:
                st.markdown("No model files found. 🕸️")

    # 🚨 5.3 Directory Check: Make sure the path exists! 🛤️
    if not os.path.isdir(directory):
        st.sidebar.error("Invalid directory 🚫")
        return

    # 📥 5.4 Handle Uploads: Bring in the new files! 📦
    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)
          emoji = "🖼️" if ext in ['webp', 'png', 'jpeg', 'jpg'] else "🏰" if ext in ['obj', 'glb', 'gltf'] else "📹"
          st.sidebar.success(f"Uploaded {emoji} {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

    # 🌐 5.5 Build the Scene: Create your 3D masterpiece! 🎬
    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>
    """

    # 🏗️ 5.6 Add Assets and Entities: Fill the world with wonders! 🌟
    assets, ents, model_counts = generate_tilemap(files, directory, 8, 8)
    st.session_state['model_counts'] = model_counts  # Store counts for display 📊
    scene += assets + ents + "</a-scene>"

    # 🎥 5.7 Apply Camera Actions: Lights, camera, action! 🎬
    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
    )

# 🏁 6. Run the Show: Let’s launch this 3D adventure! 🎭
if __name__ == "__main__":
    main()