const http = require('http'), fs = require('fs'), path = require('path');
const PORT = 8140;
const ROOT = __dirname;
const HF = path.join(process.env.USERPROFILE, '.cache', 'huggingface', 'hub',
  'models--tiiuae--falcon-mamba-7b-instruct', 'snapshots',
  'b250fc9399d14f56aca18e9ea70bbfb1f73479eb');

const MIME = {
  '.html': 'text/html', '.js': 'text/javascript', '.wgsl': 'text/plain',
  '.json': 'application/json', '.safetensors': 'application/octet-stream'
};

const { execFile } = require('child_process');

function handleAPI(req, res) {
  let body = '';
  req.on('data', c => body += c);
  req.on('end', () => {
    const data = JSON.parse(body);
    const p = req.url;
    const script = p === '/tokenize'
      ? `import sys,json;from transformers import AutoTokenizer;t=AutoTokenizer.from_pretrained('tiiuae/falcon-mamba-7b-instruct');d=json.loads(sys.stdin.read());text=d['text'];system=d.get('system','');msgs=[];
if system:msgs.append({'role':'system','content':system})
msgs.append({'role':'user','content':text});templated=t.apply_chat_template(msgs,tokenize=False,add_generation_prompt=True);print(json.dumps(t.encode(templated)))`
      : `import sys,json;from transformers import AutoTokenizer;t=AutoTokenizer.from_pretrained('tiiuae/falcon-mamba-7b-instruct');print(json.dumps(t.decode(json.loads(sys.stdin.read())['tokens'])))`;
    const child = execFile('python', ['-c', script], { maxBuffer: 10*1024*1024 }, (err, stdout) => {
      res.writeHead(200, { 'Content-Type': 'application/json', 'Access-Control-Allow-Origin': '*' });
      if (err) { res.end(JSON.stringify({ error: err.message })); return; }
      res.end(JSON.stringify({ result: JSON.parse(stdout.trim()) }));
    });
    child.stdin.write(JSON.stringify(data));
    child.stdin.end();
  });
}

http.createServer((req, res) => {
  if (req.method === 'OPTIONS') {
    res.writeHead(200, { 'Access-Control-Allow-Origin': '*', 'Access-Control-Allow-Methods': 'GET,POST', 'Access-Control-Allow-Headers': 'Content-Type' });
    return res.end();
  }
  if (req.method === 'POST' && (req.url === '/tokenize' || req.url === '/detokenize')) {
    return handleAPI(req, res);
  }

  let p = decodeURIComponent(req.url.split('?')[0]);
  if (p === '/') p = '/index.html';

  let fp;
  if (p.startsWith('/weights/')) {
    fp = path.join(HF, p.slice('/weights/'.length));
  } else {
    fp = path.join(ROOT, p);
  }

  fs.stat(fp, (e, st) => {
    if (e) { res.writeHead(404); return res.end('not found: ' + fp); }
    const total = st.size;
    const range = req.headers.range;
    const ct = MIME[path.extname(fp)] || 'application/octet-stream';

    if (range) {
      const m = range.match(/bytes=(\d+)-(\d*)/);
      if (m) {
        const start = parseInt(m[1]);
        const end = m[2] ? parseInt(m[2]) : total - 1;
        res.writeHead(206, {
          'Content-Type': ct,
          'Content-Range': 'bytes ' + start + '-' + end + '/' + total,
          'Content-Length': end - start + 1,
          'Accept-Ranges': 'bytes',
          'Access-Control-Allow-Origin': '*'
        });
        fs.createReadStream(fp, { start, end }).pipe(res);
        return;
      }
    }

    res.writeHead(200, {
      'Content-Type': ct, 'Content-Length': total,
      'Accept-Ranges': 'bytes', 'Access-Control-Allow-Origin': '*'
    });
    fs.createReadStream(fp).pipe(res);
  });
}).listen(PORT, () => {
  console.log('Mamba WebGPU on :' + PORT);
  console.log('Weights from: ' + HF);
  console.log('HF exists: ' + fs.existsSync(HF));
});