import { env, PreTrainedModel, AutoTokenizer, Tensor } from '@huggingface/transformers';

export const MODEL_ID = 'raincandy-u/MacroStories';
export const REVISION = 'dc1ef998bf5331f7f70a77c815ed3e4f643a0ba7';
export const MODEL_PATH = 'macrostories';

env.allowRemoteModels = false;
env.allowLocalModels = true;
// Keep this an origin-relative path: Transformers.js metadata discovery treats
// absolute HTTP paths as remote even when they refer to our own static assets.
env.localModelPath = new URL('../models/', import.meta.url).pathname;
env.backends.onnx.wasm.wasmPaths = new URL('../vendor/', import.meta.url).href;
// A static HF Space iframe doesn't need SharedArrayBuffer or COOP/COEP headers.
env.backends.onnx.wasm.numThreads = 1;
env.backends.onnx.wasm.proxy = false;

export async function loadRuntime(device = 'wasm', progress = () => {}) {
  if (!['wasm', 'webgpu'].includes(device)) throw new Error('Choose WebAssembly or WebGPU.');
  if (device === 'webgpu') {
    if (!globalThis.navigator?.gpu || !await navigator.gpu.requestAdapter()) {
      throw new Error('No WebGPU adapter available.');
    }
  }
  const [tokenizer, vocabulary, provenance] = await Promise.all([
    AutoTokenizer.from_pretrained(MODEL_PATH, { local_files_only: true }),
    fetch(new URL('../models/macrostories/tokenizer.json', import.meta.url)).then(r => { if (!r.ok) throw new Error('Tokenizer file missing.'); return r.json(); }),
    fetch(new URL('../models/macrostories/provenance.json', import.meta.url)).then(r => { if (!r.ok) throw new Error('Model provenance missing.'); return r.json(); }),
  ]);
  // Fallback belongs to BackendController: never reuse ORT after GPU failure.
  const runtime = await createRuntime(device, progress);
  return { ...runtime, tokenizer, vocabulary: vocabulary.model.vocab, provenance };
}

async function createRuntime(device, progress) {
  progress({ status: 'runtime', device });
  // Ouro isn't a registered AutoModelForCausalLM architecture. The public base
  // class loads the exported graph; our loop supplies its custom four-pass KV cache.
  const model = await PreTrainedModel.from_pretrained(MODEL_PATH, {
    device, dtype: 'fp32', local_files_only: true, progress_callback: progress,
    // Keep the growing KV cache on the GPU. Only the 378 logits need readback
    // for sampling; copying all past keys/values out and back per token is wasteful.
    session_options: device === 'webgpu' ? { preferredOutputLocation: { present_key_values: 'gpu-buffer', logits: 'cpu' } } : {},
  });
  const decoder = new CachedDecoder(model);
  try {
    const logits = await decoder.forward(1);
    if (logits.length !== 378 || !logits.every(Number.isFinite)) throw new Error('Runtime produced invalid logits.');
    decoder.reset();
  } catch (error) {
    await model.dispose();
    throw error;
  }
  return { model, decoder, device };
}

export class CachedDecoder {
  constructor(model) { this.model = model; this.reset(); }
  reset() {
    this.past?.dispose();
    this.position = 0;
    this.past = new Tensor('float32', new Float32Array(0), [4, 2, 0, 16]);
  }
  async forward(token) {
    if (!Number.isInteger(token) || token < 0 || token >= 378) throw new Error('Invalid token ID.');
    if (this.position >= 2048) throw new Error('The 2,048-token context is full.');
    const input = new Tensor('int64', BigInt64Array.of(BigInt(token)), [1]);
    const position = new Tensor('int64', BigInt64Array.of(BigInt(this.position)), [1]);
    let output;
    try {
      output = await this.model({ input_ids: input, position, past_key_values: this.past });
      const logits = Float32Array.from(output.logits.data);
      this.past.dispose();
      this.past = output.present_key_values;
      this.position++;
      return logits;
    } finally {
      input.dispose();
      position.dispose();
      output?.logits.dispose();
    }
  }
}

export function encodePrompt(tokenizer, text) {
  const canonical = text.replace(/\b(Alex|Robin|Casey)\b/g, name => `[PERSON_NAME_${['Alex', 'Robin', 'Casey'].indexOf(name) + 1}]`);
  const ids = Array.from(tokenizer(canonical).input_ids.data, Number);
  return { ids, unknown: ids.filter(id => id === 3).length, canonical };
}

export function renderStory(tokenizer, ids) {
  const display = ids.filter(id => ![0, 1, 2].includes(id));
  if (!display.length) return { raw: '', text: '' };
  const raw = tokenizer.decode(display, { skip_special_tokens: false }).trim();
  const text = raw.replace(/\[PERSON_NAME_(\d+)\]/g, (match, n) => ['Alex', 'Robin', 'Casey'][Number(n) - 1] ?? match);
  return { raw, text };
}