worker.js

const HF_BASE = 'https://huggingface.co/idle-intelligence/pocket-tts-gguf/resolve/main';
const VOICE_BASE = 'https://huggingface.co/kyutai/pocket-tts-without-voice-cloning/resolve/main';

let model = null;
let tokenizer = null;
let voiceIndices = {};   // name → voice_index
let activeVoiceIndex = -1;

function post(type, data = {}, transferables = []) {
    self.postMessage({ type, ...data }, transferables);
}

// ---- Fetch with Cache API + progress ----
const CACHE_NAME = 'tts-model-v3';

async function cachedFetch(url, label) {
    const cache = await caches.open(CACHE_NAME);
    const cached = await cache.match(url);
    if (cached) {
        post('status', { text: `${label} (cached)` });
        return await cached.arrayBuffer();
    }
    const resp = await fetch(url);
    if (!resp.ok) throw new Error(`Failed to fetch ${url}: ${resp.status}`);
    const contentLength = parseInt(resp.headers.get('Content-Length') || '0', 10);
    const reader = resp.body.getReader();
    const chunks = [];
    let loaded = 0;
    while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        chunks.push(value);
        loaded += value.byteLength;
        post('status', { text: label, progress: { loaded, total: contentLength } });
    }
    const buf = new Uint8Array(loaded);
    let offset = 0;
    for (const chunk of chunks) { buf.set(chunk, offset); offset += chunk.byteLength; }
    try {
        await cache.put(url, new Response(buf.buffer, { headers: { 'Content-Type': 'application/octet-stream' } }));
    } catch (e) { console.warn('[worker] cache error:', e); }
    return buf.buffer;
}

// ---- Minimal protobuf decoder for sentencepiece .model files ----
function decodeSentencepieceModel(buffer) {
    const view = new DataView(buffer.buffer, buffer.byteOffset, buffer.byteLength);
    let pos = 0;

    function readVarint() {
        let result = 0, shift = 0;
        while (pos < buffer.length) {
            const b = buffer[pos++];
            result |= (b & 0x7f) << shift;
            shift += 7;
            if ((b & 0x80) === 0) return result;
        }
        return result;
    }

    function readBytes(n) {
        const data = buffer.slice(pos, pos + n);
        pos += n;
        return data;
    }

    function decodePiece(data) {
        let pPos = 0, piece = '', score = 0, type = 1;
        const pView = new DataView(data.buffer, data.byteOffset, data.byteLength);
        while (pPos < data.length) {
            const key = readVarIntFrom(data, pPos);
            pPos = key.pos;
            const fieldNum = key.val >>> 3;
            const wireType = key.val & 0x7;
            if (fieldNum === 1 && wireType === 2) {
                const len = readVarIntFrom(data, pPos);
                pPos = len.pos;
                piece = new TextDecoder().decode(data.slice(pPos, pPos + len.val));
                pPos += len.val;
            } else if (fieldNum === 2 && wireType === 5) {
                score = pView.getFloat32(pPos, true);
                pPos += 4;
            } else if (fieldNum === 3 && wireType === 0) {
                const v = readVarIntFrom(data, pPos);
                type = v.val;
                pPos = v.pos;
            } else {
                if (wireType === 0) { const v = readVarIntFrom(data, pPos); pPos = v.pos; }
                else if (wireType === 1) { pPos += 8; }
                else if (wireType === 2) { const len = readVarIntFrom(data, pPos); pPos = len.pos + len.val; }
                else if (wireType === 5) { pPos += 4; }
                else break;
            }
        }
        return { piece, score, type };
    }

    function readVarIntFrom(buf, p) {
        let result = 0, shift = 0;
        while (p < buf.length) {
            const b = buf[p++];
            result |= (b & 0x7f) << shift;
            shift += 7;
            if ((b & 0x80) === 0) return { val: result, pos: p };
        }
        return { val: result, pos: p };
    }

    const pieces = [];
    while (pos < buffer.length) {
        const key = readVarint();
        const fieldNum = key >>> 3;
        const wireType = key & 0x7;
        if (fieldNum === 1 && wireType === 2) {
            const len = readVarint();
            const data = readBytes(len);
            const p = decodePiece(data);
            pieces.push(p);
        } else {
            if (wireType === 0) { readVarint(); }
            else if (wireType === 1) { pos += 8; }
            else if (wireType === 2) { const len = readVarint(); pos += len; }
            else if (wireType === 5) { pos += 4; }
            else break;
        }
    }
    return pieces;
}

// ---- Unigram tokenizer (Viterbi) ----
class UnigramTokenizer {
    constructor(pieces) {
        this.pieces = pieces;
        this.vocab = new Map();
        this.unkId = 0;
        for (let i = 0; i < pieces.length; i++) {
            const p = pieces[i];
            if (p.type === 2) this.unkId = i;
            if (p.type === 1 || p.type === 4) {
                this.vocab.set(p.piece, { id: i, score: p.score });
            }
            if (p.type === 6) {
                this.vocab.set(p.piece, { id: i, score: p.score });
            }
        }
    }

    encode(text) {
        const normalized = '\u2581' + text.replace(/ /g, '\u2581');
        return this._viterbi(normalized);
    }

    _viterbi(text) {
        const n = text.length;
        const best = new Array(n + 1);
        best[0] = { score: 0, len: 0, id: -1 };
        for (let i = 1; i <= n; i++) {
            best[i] = { score: -Infinity, len: 0, id: -1 };
        }

        for (let i = 0; i < n; i++) {
            if (best[i].score === -Infinity) continue;
            for (let len = 1; len <= n - i && len <= 64; len++) {
                const sub = text.substring(i, i + len);
                const entry = this.vocab.get(sub);
                if (entry) {
                    const newScore = best[i].score + entry.score;
                    if (newScore > best[i + len].score) {
                        best[i + len] = { score: newScore, len: len, id: entry.id };
                    }
                }
            }
            if (best[i + 1].score === -Infinity) {
                const ch = text.charCodeAt(i);
                const byteStr = `<0x${ch.toString(16).toUpperCase().padStart(2, '0')}>`;
                const byteEntry = this.vocab.get(byteStr);
                const fallbackId = byteEntry ? byteEntry.id : this.unkId;
                const fallbackScore = byteEntry ? byteEntry.score : -100;
                best[i + 1] = { score: best[i].score + fallbackScore, len: 1, id: fallbackId };
            }
        }

        const ids = [];
        let p = n;
        while (p > 0) {
            ids.push(best[p].id);
            p -= best[p].len;
        }
        ids.reverse();
        return new Uint32Array(ids);
    }
}

const MAX_TOKENS_PER_CHUNK = 50;
const SPACE_MARKER = '\u2581';
const SPACE_MARKER_ID = 260; // ▁ in the SentencePiece vocab; identified empirically

// Build a boundary-token set from a punctuation string, excluding the
// SentencePiece space marker (▁). The encoder prepends ▁ to its input, so
// encoding e.g. '.!...?' yields [..., 260, ...] where 260 is the bare ▁
// token. If we include 260 in the boundary set, every standalone ▁ in the
// tokenized sentence (which appears whenever the Viterbi tokenizer decides
// to split a ▁word into ▁ + word) gets treated as a sentence boundary —
// producing phantom splits that then feed the model 1-token segments.
function boundaryIds(text) {
    return tokenizer.encode(text).filter(id => id !== SPACE_MARKER_ID);
}

function splitTextIntoSegments(text, temperature) {
    const [processedText, defaultFramesAfterEos] = model.prepare_text(text);
    const normalized = processedText.trim();
    const tokenIds = tokenizer.encode(normalized);
    const totalTokens = tokenIds.length;

    if (totalTokens <= MAX_TOKENS_PER_CHUNK) {
        return [{ text: normalized, framesAfterEos: defaultFramesAfterEos }];
    }

    // Find boundary token IDs for sentence punctuation
    const eosSet = new Set(boundaryIds('.!...?'));
    const commaSet = new Set(boundaryIds(',;:'));

    // Split on sentence boundaries
    const boundaries = [0];
    let prevWasBoundary = false;
    for (let i = 0; i < tokenIds.length; i++) {
        if (eosSet.has(tokenIds[i])) {
            prevWasBoundary = true;
        } else {
            if (prevWasBoundary) boundaries.push(i);
            prevWasBoundary = false;
        }
    }
    boundaries.push(tokenIds.length);

    // Build segments from boundaries
    const segments = [];
    for (let i = 0; i < boundaries.length - 1; i++) {
        const start = boundaries[i];
        const end = boundaries[i + 1];
        const segTokens = tokenIds.slice(start, end);
        segments.push(segTokens);
    }

    // Sub-split oversized segments on commas
    const refined = [];
    for (const seg of segments) {
        if (seg.length <= MAX_TOKENS_PER_CHUNK) {
            refined.push(seg);
        } else {
            const subBounds = [0];
            let prev = false;
            for (let i = 0; i < seg.length; i++) {
                if (commaSet.has(seg[i])) {
                    prev = true;
                } else {
                    if (prev) subBounds.push(i);
                    prev = false;
                }
            }
            subBounds.push(seg.length);

            for (let i = 0; i < subBounds.length - 1; i++) {
                const sub = seg.slice(subBounds[i], subBounds[i + 1]);
                if (sub.length > 0) refined.push(sub);
            }
        }
    }

    // Final pass: merge tiny segments back together (under 10 tokens).
    // NOTE: do NOT use `buffer.flatMap(x => x)` here. `refined` contains
    // Uint32Arrays, and `flatMap` only flattens Array instances — not typed
    // arrays. Using flatMap silently produces `[Uint32Array]` (a regular
    // array wrapping one Uint32Array), and the downstream `new Uint32Array
    // (tokenIds)` then sees a 1-element input and feeds the model a single
    // token, which is why "audios" used to produce 2.24s of gibberish.
    const merged = [];
    let buffer = [];
    let bufferLen = 0;
    for (const seg of refined) {
        buffer.push(seg);
        bufferLen += seg.length;
        if (bufferLen >= 10 || seg === refined[refined.length - 1]) {
            const totalLen = buffer.reduce((sum, s) => sum + s.length, 0);
            const combined = new Uint32Array(totalLen);
            let off = 0;
            for (const s of buffer) {
                combined.set(s, off);
                off += s.length;
            }
            merged.push(combined);
            buffer = [];
            bufferLen = 0;
        }
    }

    return merged.map(tokens => ({
        tokens,
        framesAfterEos: 1,
    }));
}

// ---- Handlers ----
async function handleLoad(config) {
    const base = (config.baseUrl || '').replace(/\/+$/, '');

    // 1. Import WASM
    post('status', { text: 'Loading WASM module...' });
    const wasmJsUrl = base ? (base + '/pkg/tts_wasm.js') : new URL('../pkg/tts_wasm.js', import.meta.url).href;
    const wasmBgUrl = base ? (base + '/pkg/tts_wasm_bg.wasm') : new URL('../pkg/tts_wasm_bg.wasm', import.meta.url).href;
    const wasmModule = await import(wasmJsUrl);
    await wasmModule.default(wasmBgUrl);

    // 2. Download and load tokenizer
    const tokUrl = config.tokenizerUrl || `${HF_BASE}/tokenizer.model`;
    const tokBuf = await cachedFetch(tokUrl, 'Downloading tokenizer');
    const pieces = decodeSentencepieceModel(new Uint8Array(tokBuf));
    tokenizer = new UnigramTokenizer(pieces);
    post('status', { text: `Tokenizer loaded (${pieces.length} pieces)` });

    // 3. Download and init model
    const modelUrl = config.modelUrl || `${HF_BASE}/pocket-tts-q8_0.gguf`;
    const modelBuf = await cachedFetch(modelUrl, 'Downloading model');
    post('status', { text: 'Initializing model...' });
    model = new wasmModule.Model(new Uint8Array(modelBuf));

    // 4. Ready (voice loaded separately)
    const sampleRate = model.sample_rate();
    post('status', { text: 'Select a voice', ready: true });
    post('loaded', { sampleRate });
}

async function handleLoadVoice(name) {
    if (name in voiceIndices) {
        activeVoiceIndex = voiceIndices[name];
        post('voice_loaded', { name, voiceIndex: activeVoiceIndex });
        return;
    }

    const url = `${VOICE_BASE}/embeddings_v2/${name}.safetensors`;
    const voiceBuf = await cachedFetch(url, `Downloading voice: ${name}`);
    post('status', { text: `Loading voice: ${name}...` });
    const voiceIndex = model.add_voice(new Uint8Array(voiceBuf));
    voiceIndices[name] = voiceIndex;
    activeVoiceIndex = voiceIndex;
    post('status', { text: 'Ready', ready: true });
    post('voice_loaded', { name, voiceIndex });
}

async function handleGenerate(text, temperature) {
    const segments = splitTextIntoSegments(text, temperature);
    const totalTokens = segments.reduce((sum, s) => sum + (s.tokens ? s.tokens.length : 0), 0);

    post('gen_start', { numTokens: totalTokens, totalSegments: segments.length });

    let step = 0;
    for (let segIdx = 0; segIdx < segments.length; segIdx++) {
        const seg = segments[segIdx];
        const tokenIds = seg.tokens || tokenizer.encode(seg.text);
        const framesAfterEos = seg.framesAfterEos;

        model.start_generation(activeVoiceIndex, new Uint32Array(tokenIds), framesAfterEos, temperature);

        while (true) {
            const chunk = model.generation_step();
            if (!chunk) break;
            post('chunk', { data: chunk, step, segment: segIdx }, [chunk.buffer]);
            step++;
        }

        // Tiny pause between segments so the audio worklet doesn't glitch
        if (segIdx < segments.length - 1) {
            await new Promise(r => setTimeout(r, 20));
        }
    }

    post('done', { totalSteps: step });
}

self.onmessage = async (e) => {
    const { type, ...data } = e.data;
    try {
        if (type === 'load') {
            await handleLoad(data.config || {});
        } else if (type === 'load_voice') {
            await handleLoadVoice(data.name);
        } else if (type === 'generate') {
            await handleGenerate(data.text, data.temperature || 0.7);
        }
    } catch (err) {
        post('error', { message: err.message || String(err) });
        console.error(err);
    }
};