import * as ort from './vendor/ort.webgpu.bundle.min.mjs';
import {Tokenizer} from './vendor/tokenizers.mjs';
import {WhisperFeatureExtractor, ClapFeatureExtractor} from './vendor/transformers.web.min.js';

ort.env.wasm.wasmPaths = new URL('./vendor/', import.meta.url).href;
ort.env.wasm.numThreads = 1;
ort.env.wasm.proxy = false;

let model = null;
let loading = null;
let languageLoading = null;
let soundLoading = null;
let modelFileSizes = {};
let currentId = null;
const textCache = new Map();
const truncatedTexts = new Set();
const progress = message => postMessage({id: currentId, progress: message});

const BROAD = ['animal', 'nature', 'music', 'environment'];
const ANIMAL8 = ['dog', 'cat', 'bird', 'insect', 'chicken_and_rooster', 'crow', 'frog', 'cricket'];
const ALIASES = {
  animal: ['animal sound', 'animal', '動物音', '動物の音', '動物の鳴き声'],
  nature: ['nature sound', 'nature', '自然音', '自然の音'],
  music: ['music', 'musical', '音楽'],
  environment: ['environmental sound', 'environment sound', 'environment', '環境音', '環境の音'],
  dog: ['dog', '犬', 'いぬ', '犬の鳴き声'], cat: ['cat', '猫', 'ねこ', '猫の鳴き声'],
  bird: ['bird', '鳥', 'とり', '鳥の鳴き声'], insect: ['insect', '虫', '昆虫', '虫の音'],
  chicken_and_rooster: ['chicken', 'rooster', 'hen', '鶏', 'にわとり', 'ニワトリ', '雄鶏'],
  crow: ['crow', 'カラス', 'からす'], frog: ['frog', 'カエル', 'かえる'],
  cricket: ['cricket', 'コオロギ', 'こおろぎ'],
};
const OR_CUES = ['または', 'もしくは', 'あるいは', 'いずれか', 'or', 'either'];
const NEGATIVE_CUES = ['含まれない', '聞こえない', '存在しない', '該当しない', '満たさない', 'ではない', 'なし', 'ない', 'without', 'absent', 'not', 'no '];

async function asset(path) {
  const response = await fetch(new URL(path, import.meta.url));
  if (!response.ok) throw new Error(`モデルファイルの取得に失敗しました (${response.status})。Spaceへのアクセス権を確認してください。`);
  return response;
}

async function modelBytes(path, label) {
  const response = await asset(path);
  const filename = path.split('/').at(-1).split('?')[0];
  const total = Number(modelFileSizes[filename]?.bytes || response.headers.get('Content-Length') || 0);
  const report = received => {
    const amount = (received / 1e6).toFixed(1);
    progress(total
      ? `${label}を端末に読み込み中… ${Math.min(100, Math.round(received / total * 100))}%（${amount} / ${(total / 1e6).toFixed(1)} MB）`
      : `${label}を端末に読み込み中… ${amount} MB`);
  };
  report(0);
  if (!response.body) {
    const bytes = new Uint8Array(await response.arrayBuffer());
    report(bytes.length);
    return bytes;
  }
  const reader = response.body.getReader();
  const chunks = []; let received = 0; let last = performance.now();
  while (true) {
    const {done, value} = await reader.read(); if (done) break;
    chunks.push(value); received += value.length;
    if (performance.now() - last > 400) {
      report(received);
      last = performance.now();
    }
  }
  const result = new Uint8Array(received); let offset = 0;
  for (const chunk of chunks) {result.set(chunk, offset); offset += chunk.length;}
  report(received);
  return result;
}

async function createSession(path, label) {
  const bytes = await modelBytes(path, label);
  progress(`${label}を準備中…`);
  return ort.InferenceSession.create(bytes, {
    executionProviders: ['wasm'], graphOptimizationLevel: 'all',
  });
}

async function load() {
  if (model?.audio && model?.head && model?.sound) return status();
  if (loading) return loading;
  loading = (async () => {
    if (!model) {
      const manifest = await (await asset('./models/manifest.json')).json();
      modelFileSizes = manifest.files || {};
      progress('Edge AIモデルを端末に読み込みます。');
      const text = await createSession('./models/text_q8.onnx', '質問モデル');
      const routerCore = await createSession('./models/modality_router_core_fp32.onnx?v=20260921-router-ensemble-web-v1', 'ルーティングモデル 1/2');
      const routerNatural = await createSession('./models/modality_router_natural_fp32.onnx?v=20260921-router-ensemble-web-v1', 'ルーティングモデル 2/2');
      const [tokenizerJson, tokenizerConfig, qwenConfig, clapConfig] = await Promise.all([
        'tokenizer.json', 'tokenizer_config.json',
      ].map(async name => (await asset('./models/tokenizer/' + name)).json()).concat([
        asset('./models/qwen_preprocessor_config.json').then(response => response.json()),
        asset('./models/clap_preprocessor_config.json').then(response => response.json()),
      ]));
      model = {
        audio: null, text, head: null, routers: [routerCore, routerNatural], sound: null,
        tokenizer: new Tokenizer(tokenizerJson, tokenizerConfig),
        qwenExtractor: new WhisperFeatureExtractor(qwenConfig),
        clapExtractor: new ClapFeatureExtractor(clapConfig),
        manifest, backend: 'wasm', precision: 'int8',
      };
    }
    await loadLanguage();
    await loadSound();
    progress('モデルの準備ができました。データはこの端末内だけで処理されます。');
    return status();
  })();
  try {return await loading;} finally {loading = null;}
}

async function loadLanguage() {
  if (model.audio && model.head) return {audio: model.audio, head: model.head};
  if (languageLoading) return languageLoading;
  languageLoading = (async () => {
    progress('発話モデルを端末に読み込み中…');
    const audio = await createSession('./models/qwen_audio_q8.onnx', '発話モデル');
    const head = await createSession('./models/qwen_typed_head_q8.onnx?v=20260921-noulmatch-v4', '発話の判断モデル');
    model.audio = audio; model.head = head;
    return {audio, head};
  })();
  try {return await languageLoading;} finally {languageLoading = null;}
}

async function loadSound() {
  if (model.sound) return model.sound;
  if (soundLoading) return soundLoading;
  soundLoading = (async () => {
    progress('非言語モデルを端末に読み込み中…');
    const audio = await createSession('./models/clap_audio_q8.onnx', '非言語モデル');
    const heads = await createSession('./models/clap_heads_q8.onnx', '非言語の判断モデル');
    model.sound = {audio, heads};
    return model.sound;
  })();
  try {return await soundLoading;} finally {soundLoading = null;}
}

function status() {
  return {ready: true, backend: model.backend, precision: model.precision,
    model: model.manifest.model, on_device: true, language_ready: Boolean(model.audio), sound_ready: Boolean(model.sound)};
}

function optionsFor(q) {
  if (!q || typeof q !== 'object' || typeof q.instructions !== 'string' || !q.instructions.trim()) throw new Error('各質問にinstructions（質問文）を指定してください。');
  if (q.instructions.length > 2000) throw new Error('質問文は2000文字以内にしてください。');
  let pairs;
  if (q.type === 'noul') {
    const c = q.criteria ?? {false: 'いいえ。その条件を満たさない。', true: 'はい。その条件を満たす。'};
    if (!c || Array.isArray(c) || Object.keys(c).sort().join(',') !== 'false,true') throw new Error('Noulのcriteriaにはfalseとtrueを指定してください。');
    pairs = [['false', c.false], ['true', c.true]];
  } else if (q.type === 'choice') {
    if (!q.criteria || typeof q.criteria !== 'object' || Array.isArray(q.criteria)) throw new Error('Choiceのcriteriaは候補IDと説明のオブジェクトです。');
    pairs = Object.entries(q.criteria);
    if (pairs.length < 2 || pairs.length > 32) throw new Error('Choiceの候補数は2〜32個です。');
  } else if (q.type === 'score') {
    if (!Array.isArray(q.criteria) || q.criteria.length < 2 || q.criteria.length > 10) throw new Error('Scoreのcriteriaは2〜10段階の説明を順に並べた配列です。');
    pairs = q.criteria.map((value, index) => [String(index), value]);
  } else throw new Error('typeはnoul・choice・scoreのいずれかです。');
  if (pairs.some(([key, value]) => !key || typeof value !== 'string' || !value.trim() || value.length > 1000)) throw new Error('各候補には空でない説明を1000文字以内で指定してください。');
  return pairs;
}

function contains(text, alias) {
  const value = text.toLocaleLowerCase().replace(/\s+/g, ' ').trim();
  const target = alias.toLocaleLowerCase().replace(/\s+/g, ' ').trim();
  if (/^[\x00-\x7F]+$/.test(target) && /[a-z]/.test(target)) {
    const escaped = target.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
    return new RegExp(`(^|[^a-z])${escaped}([^a-z]|$)`).test(value);
  }
  return value.includes(target);
}

function labelsIn(text, allowed) {
  return allowed.filter(label => ALIASES[label].some(alias => contains(text, alias)));
}

function hasAny(text, cues) { return cues.some(cue => contains(text, cue)); }
function operatorFor(text) { return hasAny(text, OR_CUES) ? 'or' : 'and'; }
function isNegative(text) { return hasAny(text, NEGATIVE_CUES); }

function softmax(values) {
  const maximum = Math.max(...values);
  const exponentials = values.map(value => Math.exp(value - maximum));
  const total = exponentials.reduce((sum, value) => sum + value, 0);
  return exponentials.map(value => value / total);
}

async function learnedRoute(contextVector, questionVector, optionVectors, count) {
  const inputs = {
    context: new ort.Tensor('float32', contextVector, [1, 256]),
    question: new ort.Tensor('float32', questionVector, [1, 256]),
    options: new ort.Tensor('float32', optionVectors, [1, count, 256]),
  };
  const core = await model.routers[0].run(inputs);
  const natural = await model.routers[1].run(inputs);
  const mix = (left, right) => Array.from(left, (value, index) => 0.4 * value + 0.6 * right[index]);
  const questionLogits = mix(core.question_logits.data, natural.question_logits.data);
  const optionLogits = mix(core.option_logits.data, natural.option_logits.data);
  const weights = softmax(questionLogits);
  const optionWeights = Array.from({length: count}, (_, index) =>
    softmax(optionLogits.slice(index * 3, index * 3 + 3)));
  const entropy = -weights.reduce((sum, value) => sum + value * Math.log(Math.max(value, 1e-12)), 0) / Math.log(3);
  const both = weights[2] >= 0.10 || entropy >= 0.40;
  const mode = both ? 'joint' : (weights[0] >= weights[1] ? 'language' : 'sound');
  return {mode, weights, optionWeights, entropy};
}

function publicRoute(route) {
  const experts = route.mode === 'language' ? ['speech'] : route.mode === 'sound' ? ['nonverbal'] : ['speech', 'nonverbal'];
  const value = {mode: route.mode, experts, method: 'learned_sparse_router_ensemble_v1', entropy: Number(route.entropy.toFixed(4)),
    weights: Object.fromEntries(['speech', 'nonverbal', 'joint'].map((name, index) => [name, Number(route.weights[index].toFixed(4))]))};
  if (route.mode === 'joint') value.fusion = 'learned_option_weighted_logit_fusion_v1';
  return value;
}

function typedAnswer(question, pairs, probabilities) {
  const total = probabilities.reduce((sum, value) => sum + value, 0);
  const ps = probabilities.map(value => value / total);
  const entropy = -ps.reduce((sum, value) => sum + (value ? value * Math.log(value) : 0), 0) / Math.log(ps.length);
  const answer = {type: question.type, probabilities: Object.fromEntries(pairs.map(([key], index) => [key, ps[index]])),
    confidence: Math.max(0, Math.min(1, 1 - entropy)), confidence_kind: 'normalized_entropy_not_calibrated'};
  if (question.type === 'noul') answer.noul = ps[1];
  else if (question.type === 'choice') answer.choice = pairs[ps.indexOf(Math.max(...ps))][0];
  else {answer.score = ps.reduce((sum, value, index) => sum + index * value, 0); answer.legend = Object.fromEntries(pairs);}
  return answer;
}

async function textVector(value, warnings) {
  if (textCache.has(value)) {
    if (truncatedTexts.has(value)) warnings.add('長い文はモデルの上限128トークンに切り詰めました。');
    return textCache.get(value);
  }
  const tokenized = model.tokenizer.encode(value, {add_special_tokens: true});
  let tokens = Array.from(tokenized.ids);
  if (tokens.length > 128) {truncatedTexts.add(value); tokens = [...tokens.slice(0, 127), 2]; warnings.add('長い文はモデルの上限128トークンに切り詰めました。');}
  const ids = new BigInt64Array(128); ids.fill(3n);
  const mask = new BigInt64Array(128);
  tokens.forEach((token, index) => {ids[index] = BigInt(token); mask[index] = 1n;});
  const output = await model.text.run({input_ids: new ort.Tensor('int64', ids, [1, 128]), attention_mask: new ort.Tensor('int64', mask, [1, 128])});
  const vector = Float32Array.from(output.features.data); textCache.set(value, vector);
  if (textCache.size > 512) {const key = textCache.keys().next().value; textCache.delete(key); truncatedTexts.delete(key);}
  return vector;
}

function qwenWindows(samples) {
  const width = 80000, hop = 40000;
  const last = Math.max(0, samples.length - width); const starts = [];
  for (let index = 0; index <= last; index += hop) starts.push(index);
  if (!starts.length || starts.at(-1) !== last) starts.push(last);
  return starts.map(start => samples.slice(start, Math.min(start + width, samples.length)));
}

async function qwenFeatures(samples) {
  const language = await loadLanguage();
  const windows = qwenWindows(samples); const features = new Float32Array(windows.length * 2048);
  for (let index = 0; index < windows.length; index += 1) {
    progress(`端末内で発話を分析中 ${index + 1}/${windows.length}`);
    const extracted = await model.qwenExtractor(windows[index], {max_length: 80000});
    const mel = Float32Array.from(extracted.input_features.data);
    const output = await language.audio.run({input_features: new ort.Tensor('float32', mel, [128, 500])});
    features.set(output.pooled_features.data, index * 2048);
  }
  return {features, count: windows.length};
}

function clapStarts(length) {
  const count = Math.min(3, Math.max(1, Math.ceil(length / 480000)));
  const last = Math.max(0, length - 480000);
  if (count === 1) return [0];
  return Array.from({length: count}, (_, index) => Math.trunc(index * last / (count - 1)));
}

async function clapScores(samples48k) {
  const sound = await loadSound();
  const starts = clapStarts(samples48k.length); const embeddings = [];
  for (let index = 0; index < starts.length; index += 1) {
    progress(`端末内で音を分析中 ${index + 1}/${starts.length}`);
    const wave = samples48k.slice(starts[index], Math.min(starts[index] + 480000, samples48k.length));
    const extracted = await model.clapExtractor(wave);
    const mel = Float32Array.from(extracted.input_features.data);
    const output = await sound.audio.run({input_features: new ort.Tensor('float32', mel, [1, 1, 1001, 64])});
    embeddings.push(Float32Array.from(output.audio_embeds.data));
  }
  const feature = new Float32Array(1024);
  for (let dimension = 0; dimension < 512; dimension += 1) {
    let sum = 0, maximum = -Infinity;
    for (const embedding of embeddings) {sum += embedding[dimension]; maximum = Math.max(maximum, embedding[dimension]);}
    feature[dimension] = sum / embeddings.length; feature[512 + dimension] = maximum;
  }
  const output = await sound.heads.run({features: new ort.Tensor('float32', feature, [1, 1024])});
  const sigmoid = values => Array.from(values, value => 1 / (1 + Math.exp(-value)));
  const broadValues = sigmoid(output.broad_logits.data);
  const animalValues = sigmoid(output.animal_logits.data);
  return {broad: Object.fromEntries(['animal', 'nature', 'music', 'environment', 'dog', 'cat', 'bird', 'insect']
      .map((label, index) => [label, broadValues[index]])),
    animal8: Object.fromEntries(ANIMAL8.map((label, index) => [label, animalValues[index]])),
    windows: starts.length};
}

function allSoundLabels(text) { return [...labelsIn(text, BROAD), ...labelsIn(text, ANIMAL8)]; }
function labelScore(label, scores) {
  const value = Number(scores[ANIMAL8.includes(label) ? 'animal8' : 'broad'][label]);
  if (!Number.isFinite(value)) throw new Error('音の判定値が不正です。');
  return Math.max(0, Math.min(1, value));
}
function combineProbabilities(values, operator) {
  if (!values.length) throw new Error('対応する音の種類を質問に指定してください。');
  return operator === 'or' ? 1 - values.reduce((product, value) => product * (1 - value), 1)
    : values.reduce((product, value) => product * value, 1);
}
function soundProbability(text, labels, scores, operator) {
  const value = combineProbabilities(labels.map(label => labelScore(label, scores)), operator);
  return isNegative(text) ? 1 - value : value;
}
function poissonBinomial(values) {
  let distribution = [1];
  for (const value of values) {
    const next = new Array(distribution.length + 1).fill(0);
    distribution.forEach((probability, count) => {
      next[count] += probability * (1 - value); next[count + 1] += probability * value;
    });
    distribution = next;
  }
  return distribution;
}
function soundLikelihoods(question, pairs, scores) {
  const instructionLabels = allSoundLabels(question.instructions);
  const optionLabels = pairs.map(([, description]) => allSoundLabels(description));
  if (question.type === 'score' && instructionLabels.length && pairs.length === instructionLabels.length + 1) {
    return poissonBinomial(instructionLabels.map(label => labelScore(label, scores)));
  }
  const values = pairs.map(([, description], index) => optionLabels[index].length
    ? soundProbability(description, optionLabels[index], scores, operatorFor(description)) : null);
  if (values.some(value => value !== null)) return values.map(value => value === null ? 1 : value);
  const labels = instructionLabels.length ? instructionLabels
    : allSoundLabels([question.instructions, ...pairs.map(([, value]) => value)].join('\n'));
  if (!labels.length) throw new Error('質問に対応する音の種類を指定してください。');
  const probability = soundProbability(question.instructions, labels, scores, operatorFor(question.instructions));
  if (pairs.length === 2) return [1 - probability, probability];
  if (question.type === 'score') {
    const maximum = pairs.length - 1;
    return pairs.map((_, index) => (1 - probability) * (1 - index / maximum) + probability * (index / maximum));
  }
  throw new Error('音のChoiceでは、各選択肢が表す音を説明してください。');
}
function soundAnswer(question, pairs, scores) {
  if (question.type === 'noul') {
    const labels = allSoundLabels(question.instructions);
    const fallback = labels.length ? labels : allSoundLabels([question.instructions, ...pairs.map(([, value]) => value)].join('\n'));
    const value = soundProbability(question.instructions, fallback, scores, operatorFor(question.instructions));
    return typedAnswer(question, pairs, [1 - value, value]);
  }
  return typedAnswer(question, pairs, soundLikelihoods(question, pairs, scores));
}
function fuseAnswer(question, pairs, languageAnswer, scores, route) {
  const language = pairs.map(([key]) => Number(languageAnswer.probabilities[key]));
  const acoustic = soundLikelihoods(question, pairs, scores);
  const combined = pairs.map((_, index) => {
    const [speechWeight, nonverbalWeight, jointWeight] = route.optionWeights[index];
    const languageWeight = speechWeight + 0.5 * jointWeight;
    const soundWeight = nonverbalWeight + 0.5 * jointWeight;
    return languageWeight * Math.log(Math.max(language[index], 1e-8))
      + soundWeight * Math.log(Math.max(acoustic[index], 1e-8));
  });
  const answer = typedAnswer(question, pairs, softmax(combined));
  answer.fusion = {method: 'learned_option_weighted_logit_fusion_v1'};
  return answer;
}

async function analyze(samples, samples48k, context, questions) {
  if (!(samples instanceof Float32Array)) samples = new Float32Array(samples);
  if (!(samples48k instanceof Float32Array)) samples48k = new Float32Array(samples48k || []);
  if (samples.length < 4000 || samples.length > 480000 || samples.some(value => !Number.isFinite(value))) throw new Error('音声は0.25秒〜30秒の範囲で入力してください。');
  if (samples48k.length < 12000 || samples48k.length > 1440000 || samples48k.some(value => !Number.isFinite(value))) throw new Error('音声の端末内変換に失敗しました。音声を選び直してください。');
  if (typeof context !== 'string' || context.length > 2000) throw new Error('文脈は2000文字以内にしてください。');
  if (!questions || typeof questions !== 'object' || Array.isArray(questions)) throw new Error('質問セットにはJSONオブジェクトを指定してください。');
  const entries = Object.entries(questions);
  if (!entries.length || entries.length > 16) throw new Error('質問数は1〜16個にしてください。');
  const pairs = entries.map(([, question]) => optionsFor(question));
  await load();
  const started = performance.now(); const warnings = new Set(); const answers = Object.create(null);
  const contextVector = await textVector(context, warnings);
  const questionVectors = []; const optionVectorsByQuestion = []; const routes = [];
  for (let index = 0; index < entries.length; index += 1) {
    const [, question] = entries[index]; const choices = pairs[index];
    const questionVector = await textVector(question.instructions, warnings);
    const optionVectors = new Float32Array(choices.length * 256);
    for (let option = 0; option < choices.length; option += 1) {
      optionVectors.set(await textVector(choices[option][1], warnings), option * 256);
    }
    questionVectors.push(questionVector); optionVectorsByQuestion.push(optionVectors);
    routes.push(await learnedRoute(contextVector, questionVector, optionVectors, choices.length));
  }

  let qwenCount = 0, qwenMs = 0;
  if (routes.some(route => route.mode === 'language' || route.mode === 'joint')) {
    const qwenStarted = performance.now();
    const audio = await qwenFeatures(samples); qwenCount = audio.count;
    for (let index = 0; index < entries.length; index += 1) {
      if (routes[index].mode === 'sound') continue;
      const [name, question] = entries[index]; const choices = pairs[index];
      progress(`質問を判定中 ${index + 1}/${entries.length}`);
      const output = await model.head.run({
        audio: new ort.Tensor('float32', audio.features, [1, audio.count, 2048]),
        audio_mask: new ort.Tensor('bool', new Uint8Array(audio.count).fill(1), [1, audio.count]),
        context: new ort.Tensor('float32', contextVector, [1, 256]),
        question: new ort.Tensor('float32', questionVectors[index], [1, 256]),
        options: new ort.Tensor('float32', optionVectorsByQuestion[index], [1, choices.length, 256]),
        option_mask: new ort.Tensor('bool', new Uint8Array(choices.length).fill(1), [1, choices.length]),
        is_noul: new ort.Tensor('bool', new Uint8Array([question.type === 'noul' ? 1 : 0]), [1]),
      });
      const logits = Array.from(output.logits.data); const maximum = Math.max(...logits);
      const exponentials = logits.map(value => Math.exp(value - maximum)); const total = exponentials.reduce((sum, value) => sum + value, 0);
      answers[name] = typedAnswer(question, choices, exponentials.map(value => value / total));
    }
    qwenMs = performance.now() - qwenStarted;
  }

  let clapWindows = 0, clapMs = 0;
  if (routes.some(route => route.mode === 'sound' || route.mode === 'joint')) {
    const clapStarted = performance.now(); const scores = await clapScores(samples48k); clapWindows = scores.windows;
    for (let index = 0; index < entries.length; index += 1) {
      const route = routes[index]; if (route.mode === 'language') continue;
      const [name, question] = entries[index];
      answers[name] = route.mode === 'joint'
        ? fuseAnswer(question, pairs[index], answers[name], scores, route)
        : soundAnswer(question, pairs[index], scores);
    }
    clapMs = performance.now() - clapStarted;
  }

  return {model: 'mocovoice-audio-jev-edge-qwen-noulmatch-clap-learned-router-v6', execution: {...status(), audio_sent_to_server: false}, answers,
    evidence: {audio_duration_ms: samples.length / 16, qwen_windows: qwenCount, clap_windows: clapWindows,
      shared_audio_encoders: ['Qwen/Qwen3-ASR-0.6B', 'laion/clap-htsat-unfused'], audio_reused_across_questions: true,
      question_routing: Object.fromEntries(entries.map(([name], index) => [name, publicRoute(routes[index])]))},
    latency_ms: {qwen: qwenMs, clap: clapMs, total: performance.now() - started}, warnings: [...warnings],
    scope: '研究用試作。未知タスクへの正確さと確率校正は未検証。'};
}

let queue = Promise.resolve();
onmessage = ({data}) => {
  queue = queue.then(async () => {
    currentId = data.id;
    try {
      const result = data.type === 'load' ? await load(data.backend) : await analyze(data.samples, data.samples48k, data.context, data.questions);
      postMessage({id: data.id, result});
    } catch (error) {postMessage({id: data.id, error: error.message || String(error)});}
  });
};