Ai_Assistant/client/_archive/main-VR.js

import * as THREE from "three";
import { VRButton } from "three/addons/webxr/VRButton.js";
import { XRControllerModelFactory } from "three/addons/webxr/XRControllerModelFactory.js";
import { OrbitControls } from "three/addons/controls/OrbitControls.js";
import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";
import { VRMUtils } from "@pixiv/three-vrm";
import { createVRMAnimationClip } from "@pixiv/three-vrm-animation";

import { VRM_PATH, WS_URL, HTTP_URL } from "./config.js";
import { loadVRM } from "./vrmLoader.js";
import { AudioManager } from "./audioManager.js";
import { AnimationManager } from "./animationManager.js";
import { loadMixamoAnimation } from "./loadMixamoAnimation.js";
import { initWebcam, takePictureAndUpload } from "./webcam.js";

// --- Globals ---
let renderer, scene, camera, dolly;
let controls; // Desktop debug controls
let vrm, currentMixer, currentAction;
let animationMgr, audioMgr;
let gltfLoader; // Loader with plugins registered

// Audio Recording State
let mediaRecorder = null;
let audioChunks = [];
let isRecording = false;
let currentStream = null; // Track the current media stream

// VR Controllers
let controller0, controller1;
let controllerGrip0, controllerGrip1;

// Recording indicator (visual feedback in VR)
let recordingIndicator = null;
let isInVRSession = false;

const clock = new THREE.Clock();
const START_POSITION = new THREE.Vector3(0, 0, 1.5); // Offset user slightly back

// --- Helper: Trim Animation Clip (from app.js) ---
function trimAnimationClip(clip, startTime, endTime) {
    startTime = Math.max(0, startTime || 0);
    const fullDuration = clip.duration;
    endTime =
        typeof endTime === "number" && endTime >= 0
            ? Math.min(fullDuration, endTime)
            : fullDuration;

    if (endTime <= startTime) {
        console.warn("trimAnimationClip: invalid range", {
            startTime,
            endTime,
            duration: fullDuration,
        });
        return null;
    }

    const newTracks = [];

    for (const track of clip.tracks) {
        const times = track.times;
        const values = track.values;
        const stride = track.getValueSize();
        const keptTimes = [];
        const keptValues = [];

        for (let i = 0; i < times.length; i++) {
            const t = times[i];
            if (t >= startTime && t <= endTime) {
                keptTimes.push(t - startTime);
                const baseIndex = i * stride;
                for (let s = 0; s < stride; s++) {
                    keptValues.push(values[baseIndex + s]);
                }
            }
        }

        if (keptTimes.length > 0) {
            const TimesArray = Float32Array;
            const ValuesArray = Float32Array;
            const newTimes = new TimesArray(keptTimes);
            const newValues = new ValuesArray(keptValues);
            let NewTrack;
            try {
                NewTrack = new track.constructor(
                    track.name,
                    newTimes,
                    newValues,
                    track.getInterpolation
                        ? track.getInterpolation()
                        : undefined,
                );
            } catch (err) {
                NewTrack = new track.constructor(
                    track.name,
                    newTimes,
                    newValues,
                );
            }
            newTracks.push(NewTrack);
        }
    }

    const newDuration = endTime - startTime;
    const newName = `${clip.name || "vrma_clip"}_trimmed_${startTime.toFixed(3)}-${endTime.toFixed(3)}`;
    return new THREE.AnimationClip(newName, newDuration, newTracks);
}

// --- Create Recording Indicator (visible in VR) ---
function createRecordingIndicator() {
    const geometry = new THREE.SphereGeometry(0.02, 16, 16);
    const material = new THREE.MeshBasicMaterial({
        color: 0xff0000,
        transparent: true,
        opacity: 0,
    });
    const indicator = new THREE.Mesh(geometry, material);
    indicator.name = "recordingIndicator";
    return indicator;
}

// --- Update Recording Indicator ---
function updateRecordingIndicator(visible) {
    if (recordingIndicator) {
        recordingIndicator.material.opacity = visible ? 1 : 0;
        // Pulse effect when recording
        if (visible) {
            const pulse = Math.sin(Date.now() * 0.01) * 0.3 + 0.7;
            recordingIndicator.material.opacity = pulse;
        }
    }
}

// --- Initialization ---
async function init() {
    // 1. Renderer
    renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
    renderer.setSize(window.innerWidth, window.innerHeight);
    renderer.setPixelRatio(window.devicePixelRatio);
    renderer.xr.enabled = true;
    renderer.xr.setReferenceSpaceType("local-floor"); // Positions 0,0,0 at floor level
    renderer.shadowMap.enabled = true;
    renderer.shadowMap.type = THREE.PCFSoftShadowMap;
    document.body.appendChild(renderer.domElement);

    // 2. VR Button
    document.body.appendChild(VRButton.createButton(renderer));

    // 3. Scene
    scene = new THREE.Scene();
    scene.background = new THREE.Color(0x505050); // Dark grey background
    // scene.background = null; // Uncomment for passthrough if using AR

    // 4. Camera & Dolly
    // WebXR updates the camera position/rotation relative to the reference space.
    // To move the "player", we move a parent group (dolly).
    dolly = new THREE.Group();
    dolly.position.copy(START_POSITION);
    scene.add(dolly);

    camera = new THREE.PerspectiveCamera(
        30,
        window.innerWidth / window.innerHeight,
        0.1,
        20,
    );
    dolly.add(camera);

    // Desktop Controls (OrbitControls)
    // Note: OrbitControls usually controls the camera directly.
    // When in VR, the camera is controlled by the headset.
    controls = new OrbitControls(camera, renderer.domElement);
    controls.target.set(0, 1.0, 0);
    controls.update();

    // 5. Lighting
    const dirLight = new THREE.DirectionalLight(0xffffff, 1);
    dirLight.position.set(3, 15, -5);
    dirLight.castShadow = true;
    scene.add(dirLight);
    scene.add(new THREE.AmbientLight(0xffffff, 1.0));

    // 6. Controllers (Visualization)
    const controllerModelFactory = new XRControllerModelFactory();

    // Grip 0 (Left Controller)
    controllerGrip0 = renderer.xr.getControllerGrip(0);
    controllerGrip0.add(
        controllerModelFactory.createControllerModel(controllerGrip0),
    );
    dolly.add(controllerGrip0);

    // Grip 1 (Right Controller)
    controllerGrip1 = renderer.xr.getControllerGrip(1);
    controllerGrip1.add(
        controllerModelFactory.createControllerModel(controllerGrip1),
    );
    dolly.add(controllerGrip1);

    // Pointers (optional, simple lines)
    const geometry = new THREE.BufferGeometry().setFromPoints([
        new THREE.Vector3(0, 0, 0),
        new THREE.Vector3(0, 0, -1),
    ]);
    const line = new THREE.Line(geometry);
    line.name = "line";
    line.scale.z = 5;

    controller0 = renderer.xr.getController(0);
    controller0.add(line.clone());
    dolly.add(controller0);

    controller1 = renderer.xr.getController(1);
    controller1.add(line.clone());
    dolly.add(controller1);

    // Create and add recording indicator to right controller
    recordingIndicator = createRecordingIndicator();
    recordingIndicator.position.set(0, 0.05, -0.05); // Position above controller
    controller1.add(recordingIndicator);

    // 7. Load VRM
    try {
        const vrmData = await loadVRM(VRM_PATH, scene);
        vrm = vrmData.vrm;
        gltfLoader = vrmData.loader; // Save for VRMA loading

        console.log("VRM loaded successfully");

        // Dispatch event for UI
        window.dispatchEvent(new CustomEvent("vrm-loaded"));

        // Shadows
        vrm.scene.traverse((obj) => {
            if (obj.isMesh) obj.castShadow = true;
        });

        currentMixer = new THREE.AnimationMixer(vrm.scene);

        // 8. Managers
        audioMgr = new AudioManager(vrm);
        animationMgr = new AnimationManager(
            vrm,
            audioMgr,
            renderer,
            scene,
            camera,
            controls,
        );

        // Unlock audio on VR session start/interaction
        // VRButton click acts as interaction.
        // We can also listen to 'select' on controllers.
        controller0.addEventListener("select", unlockAudio);
        controller1.addEventListener("select", unlockAudio);
        window.addEventListener("click", unlockAudio, { once: true });
        window.addEventListener("touchend", unlockAudio, { once: true });

        // Push-to-Talk using GRIP BUTTON (squeeze) on RIGHT controller only
        // This avoids conflict with trigger (select) used for pointing/interaction
        // Meta Quest: Grip button = squeeze events
        controller1.addEventListener("squeezestart", startRecording);
        controller1.addEventListener("squeezeend", stopRecording);

        // Also allow left controller grip as alternative
        controller0.addEventListener("squeezestart", startRecording);
        controller0.addEventListener("squeezeend", stopRecording);
    } catch (err) {
        console.error("Failed to load VRM:", err);
    }

    // 9. Helper Grids
    const gridHelper = new THREE.GridHelper(10, 10);
    scene.add(gridHelper);

    // 10. Webcam Init (Optional/Background)
    initWebcam()
        .then(() => console.log("Webcam ready"))
        .catch((e) => console.warn("Webcam not available:", e));

    // 11. WebSocket
    connectWebSocket();

    // 12. VR Session Event Handlers
    renderer.xr.addEventListener("sessionstart", onVRSessionStart);
    renderer.xr.addEventListener("sessionend", onVRSessionEnd);

    // 13. Animation Loop
    renderer.setAnimationLoop(animate);

    // Resize
    window.addEventListener("resize", onWindowResize);
}

// --- VR Session Handlers ---
function onVRSessionStart() {
    console.log("🥽 VR Session Started");
    isInVRSession = true;

    // Add class to body for CSS
    document.body.classList.add("vr-active");

    // Disable OrbitControls in VR (headset controls the camera)
    if (controls) {
        controls.enabled = false;
    }

    // Unlock audio when entering VR
    unlockAudio();
}

function onVRSessionEnd() {
    console.log("👋 VR Session Ended");
    isInVRSession = false;

    // Remove class from body
    document.body.classList.remove("vr-active");

    // Re-enable OrbitControls for desktop preview
    if (controls) {
        controls.enabled = true;
    }

    // Stop any ongoing recording
    if (isRecording) {
        stopRecording();
    }
}

// --- Audio Unlock (matches app.js behavior) ---
function unlockAudio() {
    if (
        audioMgr &&
        audioMgr.audioContext &&
        audioMgr.audioContext.state === "suspended"
    ) {
        audioMgr.audioContext
            .resume()
            .then(() => {
                console.log("🔓 AudioContext unlocked");

                // Also unlock <audio> element if needed (matches app.js)
                if (audioMgr.audioElement) {
                    audioMgr.audioElement
                        .play()
                        .then(() => {
                            audioMgr.audioElement.pause(); // immediately pause so it's silent
                            audioMgr.audioElement.currentTime = 0;
                        })
                        .catch((err) => {
                            console.warn("Audio element unlock failed:", err);
                        });
                }
            })
            .catch((err) => {
                console.warn("AudioContext resume failed:", err);
            });
    }
}

function onWindowResize() {
    camera.aspect = window.innerWidth / window.innerHeight;
    camera.updateProjectionMatrix();
    renderer.setSize(window.innerWidth, window.innerHeight);
}

function animate() {
    const delta = clock.getDelta();

    // Update mixers and managers
    if (currentMixer) currentMixer.update(delta);
    if (animationMgr) animationMgr.update(delta);
    if (vrm) vrm.update(delta);

    // Update lip sync if the method exists
    if (audioMgr && typeof audioMgr.updateLipSync === "function") {
        audioMgr.updateLipSync(delta);
    }

    // Update recording indicator
    updateRecordingIndicator(isRecording);

    // Update controls only when not in VR
    if (!isInVRSession && controls) {
        controls.update();
    }

    // Render
    renderer.render(scene, camera);
}

// --- WebSocket Handling ---
function connectWebSocket() {
    const ws = new WebSocket(WS_URL);

    ws.onopen = () => {
        console.log("✅ WebSocket connected");
        window.dispatchEvent(new CustomEvent("ws-connected"));
    };

    ws.onerror = (err) => {
        console.error("❌ WS error", err);
    };

    ws.onclose = () => {
        console.log("🔌 WebSocket disconnected, attempting reconnect in 3s...");
        window.dispatchEvent(new CustomEvent("ws-disconnected"));
        setTimeout(connectWebSocket, 3000);
    };

    ws.onmessage = async ({ data }) => {
        let msg;
        try {
            msg = JSON.parse(data);
            console.log("📩 Received:", msg.type, msg);
        } catch {
            return;
        }

        // Handle different message types
        switch (msg.type) {
            case "start_animation":
                await handleStartAnimation(msg);
                break;
            case "start_vrma":
                await handleStartVRMA(msg);
                break;
            case "start_mixamo":
                await handleStartMixamo(msg);
                break;
            case "take_picture":
                console.log("📸 Take picture request");
                await takePictureAndUpload();
                break;
            case "transcription_result":
                // Handle transcription feedback from server
                console.log("📝 Transcription result:", msg.text);
                break;
            case "transcription_error":
                console.error("❌ Transcription error:", msg.error);
                break;
            default:
                console.log("Unknown message type:", msg.type);
        }
    };
}

async function handleStartAnimation(msg) {
    const {
        audio_path,
        audio_text,
        audio_duraction,
        expression = "neutral",
    } = msg;

    if (!audioMgr) {
        console.warn("AudioManager not initialized");
        return;
    }

    audioMgr.setExpression(expression);

    try {
        await audioMgr.setupAudio(audio_path);
        await audioMgr.analyzeAudio();

        // Play audio - explicitly play for VR/mobile compatibility
        try {
            await audioMgr.audioElement.play();
            console.log("🔊 Audio playing");
        } catch (err) {
            console.warn("Audio play failed:", err);
            // Try unlocking and playing again
            unlockAudio();
            setTimeout(async () => {
                try {
                    await audioMgr.audioElement.play();
                } catch (e) {
                    console.error("Audio play retry failed:", e);
                }
            }, 100);
        }

        if (animationMgr) {
            animationMgr.play();
        }
    } catch (e) {
        console.error("Failed to play animation:", e);
    }
}

async function handleStartVRMA(msg) {
    if (!gltfLoader || !vrm) {
        console.warn("VRMA handler: loader or VRM not ready");
        return;
    }

    const {
        animation_url,
        play_once = false,
        crop_start = 0.0,
        crop_end = 0.0,
    } = msg;

    try {
        console.log("🎬 Loading VRMA animation:", animation_url);

        const gltfVrma = await gltfLoader.loadAsync(animation_url);
        const vrmAnimation = gltfVrma.userData.vrmAnimations[0];
        let clip = createVRMAnimationClip(vrmAnimation, vrm);

        // Handle cropping
        const startTime = Math.max(0, parseFloat(crop_start) || 0);
        const endTime = Math.max(
            0,
            clip.duration - (parseFloat(crop_end) || 0),
        );

        if (startTime > 0 || (parseFloat(crop_end) || 0) > 0) {
            const trimmed = trimAnimationClip(clip, startTime, endTime);
            if (trimmed) {
                clip = trimmed;
            } else {
                console.warn("VRMA trim returned null — using original clip");
            }
        }

        // Notify animation manager
        if (animationMgr) animationMgr.setVRMAPlaying(true);

        // Stop previous action gracefully
        if (currentAction && currentAction.isRunning()) {
            currentAction.crossFadeTo(currentAction, 0.2, false);
        }

        // Ensure mixer exists
        if (!currentMixer) {
            currentMixer = new THREE.AnimationMixer(vrm.scene);
        }

        // Create and configure new action
        const newAction = currentMixer.clipAction(clip);

        if (play_once) {
            newAction.setLoop(THREE.LoopOnce, 0);
            newAction.clampWhenFinished = true;
            newAction.enabled = true;
        } else {
            newAction.setLoop(THREE.LoopRepeat, Infinity);
            newAction.clampWhenFinished = false;
        }

        newAction.reset().play();

        // Crossfade from previous action
        if (currentAction && currentAction !== newAction) {
            currentAction.crossFadeTo(newAction, 0.5, false);
        }

        // Cleanup previous listener
        if (currentMixer._vrmaFinishedListener) {
            currentMixer.removeEventListener(
                "finished",
                currentMixer._vrmaFinishedListener,
            );
            currentMixer._vrmaFinishedListener = null;
        }

        // Setup finish listener
        const onFinished = (e) => {
            if (e.action === newAction) {
                if (animationMgr) animationMgr.setVRMAPlaying(false);
                if (play_once) newAction.stop();
                if (currentMixer._vrmaFinishedListener) {
                    currentMixer.removeEventListener(
                        "finished",
                        currentMixer._vrmaFinishedListener,
                    );
                    currentMixer._vrmaFinishedListener = null;
                }
            }
        };

        currentMixer._vrmaFinishedListener = onFinished;
        currentMixer.addEventListener("finished", onFinished);
        currentAction = newAction;

        console.log("✅ VRMA animation started");
    } catch (err) {
        console.error("Failed to load VRMA:", err);
    }
}

async function handleStartMixamo(msg) {
    if (!vrm) {
        console.warn("Mixamo handler: VRM not ready");
        return;
    }

    const { animation_url } = msg;

    try {
        console.log("🎬 Loading Mixamo Animation:", animation_url);
        const clip = await loadMixamoAnimation(animation_url, vrm);

        if (!currentMixer) {
            currentMixer = new THREE.AnimationMixer(vrm.scene);
        }

        const newAction = currentMixer.clipAction(clip);
        newAction.reset().play();

        if (currentAction && currentAction !== newAction) {
            currentAction.crossFadeTo(newAction, 0.5, false);
        }
        currentAction = newAction;

        console.log("✅ Mixamo animation started");
    } catch (err) {
        console.error("Failed to load Mixamo:", err);
    }
}

// --- Audio Recording Logic (Push-to-Talk) ---

async function startRecording() {
    if (isRecording) return;

    console.log("🎤 Starting recording (grip button pressed)...");

    try {
        // Request microphone access with optimized settings for voice
        const stream = await navigator.mediaDevices.getUserMedia({
            audio: {
                echoCancellation: true,
                noiseSuppression: true,
                autoGainControl: true,
                sampleRate: 16000, // Good for speech recognition
            },
        });

        currentStream = stream;

        // Use webm for better compatibility, or try mp4/opus if available
        let mimeType = "audio/webm";
        if (MediaRecorder.isTypeSupported("audio/webm;codecs=opus")) {
            mimeType = "audio/webm;codecs=opus";
        } else if (MediaRecorder.isTypeSupported("audio/mp4")) {
            mimeType = "audio/mp4";
        }

        mediaRecorder = new MediaRecorder(stream, { mimeType });
        audioChunks = [];

        mediaRecorder.ondataavailable = (event) => {
            if (event.data.size > 0) {
                audioChunks.push(event.data);
            }
        };

        mediaRecorder.onstop = async () => {
            console.log("🔄 Processing recording...");

            const audioBlob = new Blob(audioChunks, { type: mimeType });

            // Only upload if we have meaningful audio (> ~1kb typically means some content)
            if (audioBlob.size > 1000) {
                await uploadAudio(audioBlob);
            } else {
                console.log("⚠️ Recording too short, skipping upload");
            }

            // Stop all tracks to release microphone
            if (currentStream) {
                currentStream.getTracks().forEach((track) => track.stop());
                currentStream = null;
            }
        };

        mediaRecorder.start(100); // Collect data every 100ms for responsiveness
        isRecording = true;

        // Dispatch event for UI
        window.dispatchEvent(new CustomEvent("recording-start"));

        console.log("✅ Recording started");
    } catch (err) {
        console.error("❌ Error accessing microphone:", err);
        isRecording = false;
    }
}

function stopRecording() {
    if (!isRecording || !mediaRecorder) return;

    console.log("⏹️ Stopping recording (grip button released)...");

    try {
        if (mediaRecorder.state !== "inactive") {
            mediaRecorder.stop();
        }
    } catch (err) {
        console.warn("Error stopping MediaRecorder:", err);
    }

    isRecording = false;

    // Dispatch event for UI
    window.dispatchEvent(new CustomEvent("recording-stop"));
}

async function uploadAudio(blob) {
    const formData = new FormData();

    // Determine file extension based on MIME type
    let extension = ".webm";
    if (blob.type.includes("mp4")) {
        extension = ".mp4";
    } else if (blob.type.includes("ogg")) {
        extension = ".ogg";
    }

    formData.append("file", blob, `recording${extension}`);

    console.log("📤 Uploading audio...", blob.size, "bytes");

    try {
        const res = await fetch(`${HTTP_URL}/upload-audio/`, {
            method: "POST",
            body: formData,
        });

        if (!res.ok) {
            throw new Error(`HTTP error! status: ${res.status}`);
        }

        const result = await res.json();
        console.log("✅ Audio upload response:", result);
    } catch (err) {
        console.error("❌ Audio upload failed:", err);
    }
}

// Kickoff
init();