import './style.css';

import * as THREE from 'three';
import { GoogleGenAI, Modality, type LiveServerMessage } from '@google/genai';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import { GLTFLoader } from 'three/examples/jsm/loaders/GLTFLoader.js';
import { VRM, VRMExpressionPresetName, VRMLoaderPlugin, VRMUtils } from '@pixiv/three-vrm';

function requireElement<T>(value: T | null, message: string): T {
  if (!value) throw new Error(message);
  return value;
}

const app = requireElement(document.querySelector<HTMLDivElement>('#app'), 'App root not found');

app.innerHTML = `
  <aside class="panel">
    <span class="badge">TypeScript + VRM</span>
    <h1>Avatar lip sync MVP</h1>
    <p class="copy">Load a VRM avatar, upload an audio file, and drive mouth movement from the audio amplitude. This is amplitude-based lip sync, so it animates speaking energy rather than phoneme-accurate visemes.</p>
    <p class="copy" style="margin-top: 8px;">Gemini voice is set to <code>Kore</code>, which should read as a feminine voice.</p>

    <div class="stack">
      <section class="card">
        <h2>1. Avatar</h2>
        <p class="copy" style="margin: 0 0 12px; font-size: 0.95rem;">A sample avatar is loaded by default. Upload a different VRM to replace it.</p>
        <input id="vrmFile" type="file" accept=".vrm,model/gltf-binary" />
      </section>

      <section class="card">
        <h2>2. Audio</h2>
        <p class="copy" style="margin: 0 0 12px; font-size: 0.95rem;">A public-domain speech sample is bundled by default. Use the button below or upload your own file.</p>
        <div style="margin-bottom: 12px;">
          <button id="useSampleAudioButton">Use sample speech</button>
        </div>
        <input id="audioFile" type="file" accept="audio/*" />
        <div style="margin-top: 12px; display: flex; gap: 10px; flex-wrap: wrap;">
          <button id="playButton" disabled>Play</button>
          <button id="stopButton" disabled>Stop</button>
        </div>
        <audio id="audio" preload="auto"></audio>
      </section>

      <section class="card">
        <h2>3. Gemini reply</h2>
        <p class="copy" style="margin: 0 0 12px; font-size: 0.95rem;">Send a text prompt to Gemini Live, receive native audio, and pipe it into the avatar audio player.</p>
        <textarea id="promptInput" rows="4" style="width: 100%; resize: vertical; border-radius: 10px; border: 1px solid rgba(255,255,255,0.12); background: rgba(255,255,255,0.04); color: inherit; padding: 10px;" placeholder="Say something as the avatar..."></textarea>
        <div style="margin-top: 12px; display: flex; gap: 10px; flex-wrap: wrap;">
          <button id="sendPromptButton">Send to Gemini</button>
        </div>
        <p id="geminiStatus" class="footer" style="margin-bottom: 0;">Gemini idle</p>
      </section>

      <section class="card stats">
        <h2>4. Live values</h2>
        <div class="stat-row"><span>Avatar</span><strong id="avatarStatus">Not loaded</strong></div>
        <div class="stat-row"><span>Audio</span><strong id="audioStatus">Not loaded</strong></div>
        <div class="stat-row"><span>Mouth open</span><strong id="mouthValue">0.00</strong></div>
        <div class="meter"><div id="mouthMeter"></div></div>
      </section>
    </div>

    <p class="footer">If your VRM supports expression presets, this app drives <code>aa</code>. Otherwise it falls back to blendshape names like <code>A</code>, <code>aa</code>, or <code>mouthOpen</code>.</p>
  </aside>
  <main class="viewport"></main>
`;

const viewport = requireElement(app.querySelector<HTMLDivElement>('.viewport'), 'Missing viewport');
const vrmInput = requireElement(app.querySelector<HTMLInputElement>('#vrmFile'), 'Missing VRM input');
const audioInput = requireElement(app.querySelector<HTMLInputElement>('#audioFile'), 'Missing audio input');
const useSampleAudioButton = requireElement(app.querySelector<HTMLButtonElement>('#useSampleAudioButton'), 'Missing sample audio button');
const promptInput = requireElement(app.querySelector<HTMLTextAreaElement>('#promptInput'), 'Missing prompt input');
const sendPromptButton = requireElement(app.querySelector<HTMLButtonElement>('#sendPromptButton'), 'Missing Gemini send button');
const geminiStatus = requireElement(app.querySelector<HTMLElement>('#geminiStatus'), 'Missing Gemini status');
const playButton = requireElement(app.querySelector<HTMLButtonElement>('#playButton'), 'Missing play button');
const stopButton = requireElement(app.querySelector<HTMLButtonElement>('#stopButton'), 'Missing stop button');
const audioEl = requireElement(app.querySelector<HTMLAudioElement>('#audio'), 'Missing audio element');
const avatarStatus = requireElement(app.querySelector<HTMLElement>('#avatarStatus'), 'Missing avatar status');
const audioStatus = requireElement(app.querySelector<HTMLElement>('#audioStatus'), 'Missing audio status');
const mouthValue = requireElement(app.querySelector<HTMLElement>('#mouthValue'), 'Missing mouth value');
const mouthMeter = requireElement(app.querySelector<HTMLElement>('#mouthMeter'), 'Missing mouth meter');

const scene = new THREE.Scene();
scene.background = new THREE.Color(0x0f172a);

const camera = new THREE.PerspectiveCamera(35, viewport.clientWidth / viewport.clientHeight, 0.1, 100);
camera.position.set(0, 1.4, 2.2);

const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
renderer.setSize(viewport.clientWidth, viewport.clientHeight);
renderer.outputColorSpace = THREE.SRGBColorSpace;
viewport.appendChild(renderer.domElement);

const controls = new OrbitControls(camera, renderer.domElement);
controls.target.set(0, 1.2, 0);
controls.enableDamping = true;
controls.minDistance = 1;
controls.maxDistance = 5;
controls.update();

scene.add(new THREE.HemisphereLight(0xffffff, 0x1e293b, 1.4));
const dirLight = new THREE.DirectionalLight(0xffffff, 1.4);
dirLight.position.set(1, 1, 1);
scene.add(dirLight);

const floor = new THREE.Mesh(
  new THREE.CircleGeometry(3, 64),
  new THREE.MeshStandardMaterial({ color: 0x111827, roughness: 0.9, metalness: 0.05 })
);
floor.rotation.x = -Math.PI / 2;
floor.position.y = -0.01;
scene.add(floor);

const clock = new THREE.Clock();
const loader = new GLTFLoader();
loader.register((parser) => new VRMLoaderPlugin(parser));

let currentVrm: VRM | null = null;
let audioContext: AudioContext | null = null;
let analyser: AnalyserNode | null = null;
let sourceNode: MediaElementAudioSourceNode | null = null;
let dataArray: Uint8Array<ArrayBuffer> | null = null;
let vrmObjectUrl: string | null = null;
let audioObjectUrl: string | null = null;
let mouthOpenAmount = 0;
let audioEnergy = 0;
let geminiResponseParts: Uint8Array[] = [];

type PoseBones = {
  neck?: THREE.Object3D | null;
  spine?: THREE.Object3D | null;
  chest?: THREE.Object3D | null;
  upperChest?: THREE.Object3D | null;
  leftUpperArm?: THREE.Object3D | null;
  rightUpperArm?: THREE.Object3D | null;
  leftLowerArm?: THREE.Object3D | null;
  rightLowerArm?: THREE.Object3D | null;
  leftHand?: THREE.Object3D | null;
  rightHand?: THREE.Object3D | null;
};

let poseBones: PoseBones = {};

const geminiApiKey = import.meta.env.VITE_GEMINI_API_KEY as string | undefined;
const geminiModel = 'gemini-3.1-flash-live-preview';

function disposeCurrentVrm() {
  if (!currentVrm) return;
  scene.remove(currentVrm.scene);
  VRMUtils.deepDispose(currentVrm.scene);
  currentVrm = null;
  poseBones = {};
}

function cachePoseBones(vrm: VRM) {
  const humanoid = vrm.humanoid;
  poseBones = {
    neck: humanoid?.getNormalizedBoneNode('neck'),
    spine: humanoid?.getNormalizedBoneNode('spine'),
    chest: humanoid?.getNormalizedBoneNode('chest'),
    upperChest: humanoid?.getNormalizedBoneNode('upperChest'),
    leftUpperArm: humanoid?.getNormalizedBoneNode('leftUpperArm'),
    rightUpperArm: humanoid?.getNormalizedBoneNode('rightUpperArm'),
    leftLowerArm: humanoid?.getNormalizedBoneNode('leftLowerArm'),
    rightLowerArm: humanoid?.getNormalizedBoneNode('rightLowerArm'),
    leftHand: humanoid?.getNormalizedBoneNode('leftHand'),
    rightHand: humanoid?.getNormalizedBoneNode('rightHand'),
  };
}

function pcm16ToWav(pcmData: Uint8Array, sampleRate = 24000, channels = 1): Blob {
  const bytesPerSample = 2;
  const blockAlign = channels * bytesPerSample;
  const buffer = new ArrayBuffer(44 + pcmData.byteLength);
  const view = new DataView(buffer);

  const writeString = (offset: number, value: string) => {
    for (let i = 0; i < value.length; i += 1) view.setUint8(offset + i, value.charCodeAt(i));
  };

  writeString(0, 'RIFF');
  view.setUint32(4, 36 + pcmData.byteLength, true);
  writeString(8, 'WAVE');
  writeString(12, 'fmt ');
  view.setUint32(16, 16, true);
  view.setUint16(20, 1, true);
  view.setUint16(22, channels, true);
  view.setUint32(24, sampleRate, true);
  view.setUint32(28, sampleRate * blockAlign, true);
  view.setUint16(32, blockAlign, true);
  view.setUint16(34, 16, true);
  writeString(36, 'data');
  view.setUint32(40, pcmData.byteLength, true);
  new Uint8Array(buffer, 44).set(pcmData);
  return new Blob([buffer], { type: 'audio/wav' });
}

function concatUint8Arrays(chunks: Uint8Array[]): Uint8Array {
  const total = chunks.reduce((sum, chunk) => sum + chunk.byteLength, 0);
  const result = new Uint8Array(total);
  let offset = 0;
  for (const chunk of chunks) {
    result.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return result;
}

function base64ToUint8Array(base64: string): Uint8Array {
  const binary = atob(base64);
  const bytes = new Uint8Array(binary.length);
  for (let i = 0; i < binary.length; i += 1) bytes[i] = binary.charCodeAt(i);
  return bytes;
}

async function playGeneratedSpeech(pcmChunks: Uint8Array[]) {
  if (audioObjectUrl) URL.revokeObjectURL(audioObjectUrl);
  const wavBlob = pcm16ToWav(concatUint8Arrays(pcmChunks));
  audioObjectUrl = URL.createObjectURL(wavBlob);
  audioEl.src = audioObjectUrl;
  audioEl.load();
  audioStatus.textContent = 'Gemini response';
  playButton.disabled = false;
  stopButton.disabled = false;
  ensureAudioGraph();
  if (audioContext?.state === 'suspended') await audioContext.resume();
  await audioEl.play();
}

async function sendPromptToGemini(prompt: string) {
  if (!geminiApiKey) throw new Error('Missing VITE_GEMINI_API_KEY');

  geminiStatus.textContent = 'Connecting to Gemini…';
  sendPromptButton.disabled = true;
  geminiResponseParts = [];

  const ai = new GoogleGenAI({ apiKey: geminiApiKey });
  const session = await ai.live.connect({
    model: geminiModel,
    config: {
      responseModalities: [Modality.AUDIO],
      speechConfig: {
        voiceConfig: {
          prebuiltVoiceConfig: {
            voiceName: 'Kore',
          },
        },
      },
      systemInstruction: 'You are a concise speaking avatar. Reply conversationally in under 20 seconds.',
    },
    callbacks: {
      onopen: () => {
        geminiStatus.textContent = 'Gemini connected';
      },
      onmessage: async (message: LiveServerMessage) => {
        const serverContent = message.serverContent;
        const parts = serverContent?.modelTurn?.parts ?? [];
        for (const part of parts) {
          if (part.inlineData?.data) {
            geminiResponseParts.push(base64ToUint8Array(part.inlineData.data));
          }
        }

        if (serverContent?.turnComplete) {
          geminiStatus.textContent = 'Playing Gemini audio';
          if (geminiResponseParts.length > 0) await playGeneratedSpeech(geminiResponseParts);
          await session.close();
          sendPromptButton.disabled = false;
          geminiStatus.textContent = 'Gemini idle';
        }
      },
      onerror: (error: ErrorEvent) => {
        console.error(error);
        geminiStatus.textContent = 'Gemini error';
        sendPromptButton.disabled = false;
      },
      onclose: () => {
        sendPromptButton.disabled = false;
      },
    },
  });

  session.sendClientContent({
    turns: [{ role: 'user', parts: [{ text: prompt }] }],
    turnComplete: true,
  });
}

function setMouth(open: number) {
  const clamped = THREE.MathUtils.clamp(open, 0, 1);
  mouthOpenAmount = clamped;
  mouthValue.textContent = clamped.toFixed(2);
  mouthMeter.style.width = `${(clamped * 100).toFixed(1)}%`;

  if (!currentVrm) return;

  const manager = currentVrm.expressionManager;
  if (!manager) return;

  manager.setValue(VRMExpressionPresetName.Aa, clamped);
  for (const name of ['A', 'a', 'aa', 'Ah', 'mouthOpen']) {
    try {
      manager.setValue(name as VRMExpressionPresetName, clamped);
    } catch {
      // ignore unsupported expression names
    }
  }
}

function ensureAudioGraph() {
  if (audioContext && analyser && sourceNode && dataArray) return;

  audioContext = new AudioContext();
  analyser = audioContext.createAnalyser();
  analyser.fftSize = 2048;
  analyser.smoothingTimeConstant = 0.82;
  dataArray = new Uint8Array(new ArrayBuffer(analyser.frequencyBinCount));
  sourceNode = audioContext.createMediaElementSource(audioEl);
  sourceNode.connect(analyser);
  analyser.connect(audioContext.destination);
}

function computeMouthOpen() {
  if (!analyser || !dataArray || audioEl.paused) {
    audioEnergy = THREE.MathUtils.lerp(audioEnergy, 0, 0.18);
    setMouth(0);
    return;
  }

  analyser.getByteTimeDomainData(dataArray);
  let sum = 0;
  for (let i = 0; i < dataArray.length; i += 1) {
    const centered = (dataArray[i] - 128) / 128;
    sum += centered * centered;
  }

  const rms = Math.sqrt(sum / dataArray.length);
  const boosted = THREE.MathUtils.clamp((rms - 0.012) * 10.0, 0, 1);
  const exaggerated = THREE.MathUtils.clamp(Math.pow(boosted, 0.7) * 1.25, 0, 1);
  audioEnergy = THREE.MathUtils.lerp(audioEnergy, exaggerated, 0.35);
  const expressiveMouth = THREE.MathUtils.clamp(exaggerated * 1.35 + 0.02, 0, 1);
  setMouth(expressiveMouth);
}

function animatePresence(elapsed: number) {
  if (!currentVrm) return;

  const idleSway = Math.sin(elapsed * 1.4) * 0.025;
  const idleBob = Math.sin(elapsed * 2.1) * 0.012;
  const talkSway = Math.sin(elapsed * 9.0) * audioEnergy * 0.035;
  const talkNod = Math.sin(elapsed * 7.2) * audioEnergy * 0.03;
  const talkLift = audioEnergy * 0.045;

  currentVrm.scene.rotation.y = Math.PI + idleSway + talkSway;
  currentVrm.scene.rotation.x = talkNod * 0.35;
  currentVrm.scene.position.y = idleBob + talkLift;

  const {
    neck,
    spine,
    chest,
    upperChest,
    leftUpperArm,
    rightUpperArm,
    leftLowerArm,
    rightLowerArm,
    leftHand,
    rightHand,
  } = poseBones;

  if (neck) {
    neck.rotation.y = talkSway * 0.9;
    neck.rotation.x = talkNod * 0.8;
    neck.rotation.z = Math.sin(elapsed * 5.4) * audioEnergy * 0.04;
  }

  if (spine) {
    spine.rotation.y = idleSway * 0.45 + talkSway * 0.4;
    spine.rotation.x = talkNod * 0.2;
  }

  if (chest) {
    chest.rotation.x = mouthOpenAmount * 0.05 + audioEnergy * 0.04;
    chest.rotation.z = Math.sin(elapsed * 3.8) * 0.015 + talkSway * 0.3;
  }

  if (upperChest) {
    upperChest.rotation.y = talkSway * 0.35;
    upperChest.rotation.x = audioEnergy * 0.03;
  }

  if (leftUpperArm) {
    leftUpperArm.rotation.z = 1.12 + Math.sin(elapsed * 1.7) * 0.012 + audioEnergy * 0.02;
    leftUpperArm.rotation.x = 0.03 + Math.sin(elapsed * 2.3) * 0.01;
    leftUpperArm.rotation.y = -0.04 - talkSway * 0.1;
  }

  if (rightUpperArm) {
    rightUpperArm.rotation.z = -1.12 - Math.sin(elapsed * 1.9) * 0.012 - audioEnergy * 0.02;
    rightUpperArm.rotation.x = 0.03 + Math.sin(elapsed * 2.5) * 0.01;
    rightUpperArm.rotation.y = 0.04 - talkSway * 0.1;
  }

  if (leftLowerArm) {
    leftLowerArm.rotation.z = 0;
    leftLowerArm.rotation.x = -0.14 + Math.sin(elapsed * 2.8) * 0.008;
  }

  if (rightLowerArm) {
    rightLowerArm.rotation.z = 0;
    rightLowerArm.rotation.x = -0.14 + Math.sin(elapsed * 3.0) * 0.008;
  }

  if (leftHand) {
    leftHand.rotation.y = -0.08 + Math.sin(elapsed * 3.5) * 0.02;
    leftHand.rotation.z = -0.03;
  }

  if (rightHand) {
    rightHand.rotation.y = 0.08 - Math.sin(elapsed * 3.3) * 0.02;
    rightHand.rotation.z = 0.03;
  }
}

async function loadVrm(file: File) {
  disposeCurrentVrm();
  if (vrmObjectUrl) URL.revokeObjectURL(vrmObjectUrl);

  vrmObjectUrl = URL.createObjectURL(file);
  const gltf = await loader.loadAsync(vrmObjectUrl);
  const vrm = gltf.userData.vrm as VRM | undefined;
  if (!vrm) throw new Error('Loaded file is not a VRM avatar');

  VRMUtils.rotateVRM0(vrm);
  vrm.scene.rotation.y = Math.PI;
  scene.add(vrm.scene);
  currentVrm = vrm;
  cachePoseBones(vrm);
  avatarStatus.textContent = file.name;
}

async function loadDefaultVrm() {
  avatarStatus.textContent = 'Loading sample…';
  disposeCurrentVrm();
  if (vrmObjectUrl) {
    URL.revokeObjectURL(vrmObjectUrl);
    vrmObjectUrl = null;
  }

  const gltf = await loader.loadAsync('/default-avatar.vrm');
  const vrm = gltf.userData.vrm as VRM | undefined;
  if (!vrm) throw new Error('Default avatar is not a valid VRM');

  VRMUtils.rotateVRM0(vrm);
  vrm.scene.rotation.y = Math.PI;
  scene.add(vrm.scene);
  currentVrm = vrm;
  cachePoseBones(vrm);
  avatarStatus.textContent = 'default-avatar.vrm';
}

async function loadAudio(file: File) {
  if (audioObjectUrl) URL.revokeObjectURL(audioObjectUrl);

  audioObjectUrl = URL.createObjectURL(file);
  audioEl.src = audioObjectUrl;
  audioEl.load();
  audioStatus.textContent = file.name;
  playButton.disabled = false;
  stopButton.disabled = false;
  setMouth(0);
}

async function loadSampleAudio() {
  if (audioObjectUrl) {
    URL.revokeObjectURL(audioObjectUrl);
    audioObjectUrl = null;
  }

  audioEl.src = '/sample-speech.ogg';
  audioEl.load();
  audioStatus.textContent = 'sample-speech.ogg';
  playButton.disabled = false;
  stopButton.disabled = false;
  setMouth(0);
}

vrmInput.addEventListener('change', async () => {
  const file = vrmInput.files?.[0];
  if (!file) return;

  avatarStatus.textContent = 'Loading…';
  try {
    await loadVrm(file);
  } catch (error) {
    console.error(error);
    avatarStatus.textContent = 'Load failed';
  }
});

audioInput.addEventListener('change', async () => {
  const file = audioInput.files?.[0];
  if (!file) return;

  try {
    await loadAudio(file);
  } catch (error) {
    console.error(error);
    audioStatus.textContent = 'Load failed';
  }
});

useSampleAudioButton.addEventListener('click', async () => {
  try {
    await loadSampleAudio();
  } catch (error) {
    console.error(error);
    audioStatus.textContent = 'Sample load failed';
  }
});

sendPromptButton.addEventListener('click', async () => {
  const prompt = promptInput.value.trim();
  if (!prompt) return;

  try {
    await sendPromptToGemini(prompt);
  } catch (error) {
    console.error(error);
    geminiStatus.textContent = error instanceof Error ? error.message : 'Gemini request failed';
    sendPromptButton.disabled = false;
  }
});

playButton.addEventListener('click', async () => {
  try {
    ensureAudioGraph();
    if (audioContext?.state === 'suspended') await audioContext.resume();
    await audioEl.play();
  } catch (error) {
    console.error(error);
  }
});

stopButton.addEventListener('click', () => {
  audioEl.pause();
  audioEl.currentTime = 0;
  setMouth(0);
});

audioEl.addEventListener('ended', () => setMouth(0));

window.addEventListener('resize', () => {
  const width = viewport.clientWidth;
  const height = viewport.clientHeight;
  camera.aspect = width / height;
  camera.updateProjectionMatrix();
  renderer.setSize(width, height);
});

function tick() {
  const delta = clock.getDelta();
  const elapsed = clock.elapsedTime;
  currentVrm?.update(delta);
  computeMouthOpen();
  animatePresence(elapsed);
  controls.update();
  renderer.render(scene, camera);
  requestAnimationFrame(tick);
}

loadDefaultVrm().catch((error) => {
  console.error(error);
  avatarStatus.textContent = 'Sample load failed';
});

loadSampleAudio().catch((error) => {
  console.error(error);
  audioStatus.textContent = 'Sample load failed';
});

tick();