// A ShockAI Model API host: Node 24, no dependencies. Run it with
//
//   PUBLIC_ORIGIN=https://your-host.example.com node server.mts
//
// It serves your model (model.mts, described by manifest.json) under
// /shockai/v0, checks the platform's signature on every request but
// /healthz, runs each case on a thread of its own, answers a run within 25
// seconds or accepts it and calls back until the platform takes the
// callback, and keeps runs for 24 hours under data/. Guide:
// https://app.shockai.co/developers/guides/host-your-model
//
// Your key signs your callbacks. While manifest.json has no signing block,
// the host makes data/host-key.pem on first start, serves its public half as
// signing.callbackKeys and prints the block to copy into manifest.json. Once
// manifest.json lists keys, the host serves them as they are, and will not
// start without its key or with a key they do not list. To rotate, list both
// keys in a new version and connect it, sign with the new key, then drop the
// old one in the version after.
//
// Environment:
//   PUBLIC_ORIGIN       required: the origin the platform calls, as it spells it
//   PORT                default 8080
//   BASE_PATH           default /shockai/v0
//   PLATFORM_KEYS_URL   default https://app.shockai.co/developers/keys/v0.json
//   DATA_DIR            default ./data (runs, your key, the request log)
//   BUILD_HASH          what is deployed, e.g. a commit; reported with every result
//   HOST_SIGNING_KEY    your Ed25519 key as a PKCS#8 PEM; default data/host-key.pem
//   SYNC_BUDGET_MS      default 25000: how long a run may take before it is answered 202
//   MODEL_DELAY_MS      default 0: how long every run computes first, its thread kept busy, to try the 202 path

import type { ModelAnswer } from './model.mts';

import {
  type KeyObject,
  createHash,
  createPrivateKey,
  createPublicKey,
  generateKeyPairSync,
  randomUUID,
  sign,
  verify,
} from 'node:crypto';
import {
  appendFileSync,
  existsSync,
  mkdirSync,
  readFileSync,
  readdirSync,
  renameSync,
  rmSync,
  writeFileSync,
} from 'node:fs';
import {
  type IncomingMessage,
  type ServerResponse,
  createServer,
} from 'node:http';
import { dirname, join } from 'node:path';
import { fileURLToPath } from 'node:url';
import { Worker } from 'node:worker_threads';

const here = dirname(fileURLToPath(import.meta.url));
const env = process.env;
if (!env.PUBLIC_ORIGIN)
  throw new Error(
    'Set PUBLIC_ORIGIN to the origin the platform calls, e.g. https://your-host.example.com',
  );

const config = {
  port: Number(env.PORT ?? 8080),
  origin: new URL(env.PUBLIC_ORIGIN).origin,
  basePath: (env.BASE_PATH ?? '/shockai/v0').replace(/\/$/, ''),
  platformKeysUrl:
    env.PLATFORM_KEYS_URL ?? 'https://app.shockai.co/developers/keys/v0.json',
  dataDir: env.DATA_DIR ?? join(here, 'data'),
  buildHash: env.BUILD_HASH ?? 'dev',
  syncBudgetMs: Number(env.SYNC_BUDGET_MS ?? 25_000),
  modelDelayMs: Number(env.MODEL_DELAY_MS ?? 0),
};

const SCHEME = 'shockai-model-api-v0';
const TOLERANCE_SECONDS = 300;
const MAX_BODY_BYTES = 16 * 1024 * 1024;
const KEEP_RUNS_MS = 24 * 3600 * 1000;
const CALLBACK_ATTEMPTS = 8;

type Jwk = { kty: string; crv: string; x: string };
type RunStatus = 'queued' | 'running' | 'succeeded' | 'failed' | 'cancelled';
type Run = Record<string, unknown> & {
  id: string;
  status: RunStatus;
  caseId: string;
  updatedAt: string;
};
type Stored = {
  run: Run;
  idempotencyKey: string;
  callbackUrl: string | null;
  /** Answered 202, to its request or a repeat of it: it owes a callback. */
  answeredAsync: boolean;
  /** The platform answered the callback for good: took it, or refused it. */
  delivered?: boolean;
};

// ---------------------------------------------------------------------------
// Keys and signatures (the scheme: /developers/guides/signing)
// ---------------------------------------------------------------------------

mkdirSync(join(config.dataDir, 'runs'), { recursive: true });
mkdirSync(join(config.dataDir, 'keys'), { recursive: true });

const manifestFile = JSON.parse(
  readFileSync(join(here, 'manifest.json'), 'utf8'),
) as Record<string, unknown> & {
  id: string;
  version: string;
  signing?: { callbackKeys?: Jwk[] };
};
/** The keys manifest.json lists. Once it lists any, no key is ever made. */
const pinnedKeys = manifestFile.signing?.callbackKeys ?? [];

function loadHostKey(): KeyObject {
  if (env.HOST_SIGNING_KEY)
    return createPrivateKey(env.HOST_SIGNING_KEY.replace(/\\n/g, '\n'));
  const path = join(config.dataDir, 'host-key.pem');
  if (!existsSync(path)) {
    if (pinnedKeys.length > 0)
      throw new Error(
        `manifest.json lists signing.callbackKeys, but ${path} is missing. Restore it, or set HOST_SIGNING_KEY. A new key needs a new manifest version that lists it.`,
      );
    const { privateKey } = generateKeyPairSync('ed25519');
    writeFileSync(path, privateKey.export({ format: 'pem', type: 'pkcs8' }), {
      mode: 0o600,
    });
  }
  return createPrivateKey(readFileSync(path));
}

function publicJwk(key: KeyObject): Jwk {
  const { kty, crv, x } = createPublicKey(key).export({ format: 'jwk' });
  return { kty: kty!, crv: crv!, x: x! };
}

function thumbprint(jwk: Jwk): string {
  return createHash('sha256')
    .update(JSON.stringify({ crv: jwk.crv, kty: jwk.kty, x: jwk.x }))
    .digest('base64url');
}

const hostKey = loadHostKey();
const hostKid = thumbprint(publicJwk(hostKey));
if (pinnedKeys.length === 0) {
  console.log(
    `manifest.json lists no signing key, so the host serves its own. To pin it, add to manifest.json: ${JSON.stringify({ signing: { callbackKeys: [publicJwk(hostKey)] } })}`,
  );
} else if (!pinnedKeys.some((key) => thumbprint(key) === hostKid)) {
  throw new Error(
    `The host's key (kid ${hostKid}) is not among manifest.json's signing.callbackKeys. Sign with a key it lists, or list this one in a new version.`,
  );
}

function signatureBase(
  timestamp: string,
  method: string,
  target: string,
  idempotencyKey: string,
  body: Buffer | string,
): string {
  const bodyHash = createHash('sha256').update(body).digest('hex');
  return [
    SCHEME,
    timestamp,
    method.toUpperCase(),
    target,
    idempotencyKey,
    bodyHash,
  ].join('\n');
}

let platformKeys: Jwk[] = [];
let keysFetchedAt = 0;

let keysFetching: Promise<void> | null = null;

/**
 * Fetch the platform's keys, at most once a minute. A request that arrives
 * while a fetch is on its way waits for it, rather than being refused.
 */
function refreshPlatformKeys(): Promise<void> {
  if (keysFetching) return keysFetching;
  if (Date.now() - keysFetchedAt < 60_000) return Promise.resolve();
  keysFetching = (async () => {
    try {
      const response = await fetch(config.platformKeysUrl);
      if (response.ok)
        platformKeys = ((await response.json()) as { keys: Jwk[] }).keys;
    } catch {
      // Keep the keys we had; the next unknown key id asks again.
    } finally {
      keysFetchedAt = Date.now();
      keysFetching = null;
    }
  })();
  return keysFetching;
}

/** Why a request's signature fails, or null when it holds. */
async function signatureProblem(
  request: IncomingMessage,
  body: Buffer,
): Promise<string | null> {
  const parts = String(request.headers['shockai-signature'] ?? '')
    .split(',')
    .map((part) => part.trim());
  const field = (name: string) => {
    const values = parts
      .filter((part) => part.startsWith(`${name}=`))
      .map((part) => part.slice(name.length + 1));
    return values.length === 1 ? values[0]! : '';
  };
  const [timestamp, kid, signature] = [
    field('t'),
    field('kid'),
    field('ed25519'),
  ];
  if (
    !/^[0-9]{1,12}$/.test(timestamp) ||
    !/^[A-Za-z0-9_-]{86}$/.test(signature)
  )
    return 'bad_signature';
  if (Math.abs(Date.now() / 1000 - Number(timestamp)) > TOLERANCE_SECONDS)
    return 'expired_signature';

  let jwk = platformKeys.find((key) => thumbprint(key) === kid);
  if (!jwk) {
    await refreshPlatformKeys();
    jwk = platformKeys.find((key) => thumbprint(key) === kid);
  }
  if (!jwk) return 'unknown_key';

  const base = signatureBase(
    timestamp,
    request.method ?? 'GET',
    config.origin + (request.url ?? '/'),
    String(request.headers['idempotency-key'] ?? ''),
    body,
  );
  const publicKey = createPublicKey({
    key: { kty: jwk.kty, crv: jwk.crv, x: jwk.x },
    format: 'jwk',
  });
  return verify(
    null,
    Buffer.from(base),
    publicKey,
    Buffer.from(signature, 'base64url'),
  )
    ? null
    : 'bad_signature';
}

function signOutgoing(method: string, url: string, body: string): string {
  const timestamp = String(Math.floor(Date.now() / 1000));
  const target = new URL(url);
  const base = signatureBase(
    timestamp,
    method,
    target.origin + target.pathname + target.search,
    '',
    body,
  );
  const signature = sign(null, Buffer.from(base), hostKey).toString(
    'base64url',
  );
  return `t=${timestamp},kid=${hostKid},ed25519=${signature}`;
}

// ---------------------------------------------------------------------------
// The manifest: manifest.json, with your build, and your public key unless
// it lists its own
// ---------------------------------------------------------------------------

const manifest = {
  ...manifestFile,
  signing:
    pinnedKeys.length > 0
      ? manifestFile.signing
      : { callbackKeys: [publicJwk(hostKey)] },
  build: { hash: config.buildHash },
};

// ---------------------------------------------------------------------------
// Runs, kept on disk for 24 hours
// ---------------------------------------------------------------------------

const now = () => new Date().toISOString().replace(/\.\d{3}Z$/, 'Z');
const runPath = (id: string) => join(config.dataDir, 'runs', `${id}.json`);
const keyPath = (key: string) =>
  join(
    config.dataDir,
    'keys',
    `${createHash('sha256').update(key).digest('hex')}.json`,
  );

function save(stored: Stored): void {
  // Write then rename, so a crash never leaves half a file.
  const path = runPath(stored.run.id);
  writeFileSync(`${path}.tmp`, JSON.stringify(stored));
  renameSync(`${path}.tmp`, path);
}

function load(id: string): Stored | null {
  if (!/^[A-Za-z0-9_-]{1,128}$/.test(id) || !existsSync(runPath(id)))
    return null;
  return JSON.parse(readFileSync(runPath(id), 'utf8')) as Stored;
}

function isFinal(status: RunStatus): boolean {
  return (
    status === 'succeeded' || status === 'failed' || status === 'cancelled'
  );
}

function sweep(): void {
  for (const dir of ['runs', 'keys']) {
    for (const file of readdirSync(join(config.dataDir, dir))) {
      const path = join(config.dataDir, dir, file);
      let at = '';
      try {
        const record = JSON.parse(readFileSync(path, 'utf8')) as {
          run?: Run;
          at?: string;
        };
        at = record.run?.updatedAt ?? record.at ?? '';
      } catch {
        // A write cut off by a crash: nothing to keep.
      }
      if (!(Date.now() - Date.parse(at) <= KEEP_RUNS_MS))
        rmSync(path, { force: true });
    }
  }
}

function markDelivered(id: string): void {
  const current = load(id);
  if (current) save({ ...current, delivered: true });
}

/**
 * Send a run's callback until the platform takes it. A callback the host
 * could not deliver is sent again on its next start, so it arrives at least
 * once; a 4xx other than 408 or 429 is a refusal, and not sent again.
 */
async function deliver(stored: Stored): Promise<void> {
  if (!stored.answeredAsync || !stored.callbackUrl || stored.delivered) return;
  const { run, callbackUrl } = stored;
  // The event is the run reaching its state, the same on every delivery,
  // so the platform can tell a repeat.
  const body = JSON.stringify({
    id: `evt_${run.id}.${run.status}`,
    createdAt: run.updatedAt,
    type: `run.${run.status}`,
    data: { run },
  });
  for (let attempt = 0; attempt < CALLBACK_ATTEMPTS; attempt++) {
    if (attempt > 0)
      await new Promise((resolve) =>
        setTimeout(resolve, 1000 * 2 ** (attempt - 1)),
      );
    let status = 0;
    try {
      // Signed afresh for every attempt, so a retry carries a current time.
      // A redirect is not followed: a 3xx is a failed delivery.
      const response = await fetch(callbackUrl, {
        method: 'POST',
        headers: {
          'content-type': 'application/json',
          'shockai-signature': signOutgoing('POST', callbackUrl, body),
        },
        body,
        redirect: 'manual',
        signal: AbortSignal.timeout(30_000),
      });
      status = response.status;
    } catch {
      // Try again; the platform also polls GET /runs/{id}.
    }
    if (status >= 200 && status < 300) return markDelivered(run.id);
    if (status >= 400 && status < 500 && status !== 408 && status !== 429) {
      console.error('callback refused', run.id, status);
      return markDelivered(run.id);
    }
  }
  console.error(
    `callback not delivered after ${CALLBACK_ATTEMPTS} attempts; the next start tries again`,
    run.id,
  );
}

function finish(stored: Stored, result: Partial<Run>): Stored {
  const current = load(stored.run.id) ?? stored;
  if (isFinal(current.run.status)) return current; // cancelled meanwhile
  const done = {
    ...current,
    run: { ...current.run, ...result, updatedAt: now() } as Run,
  };
  save(done);
  return done;
}

// Each run's model works on a thread of its own, so a model that computes
// for minutes holds up neither /healthz nor the 202. runModel may be async.
const MODEL_THREAD = `
const { parentPort, workerData } = require('node:worker_threads');
import(workerData.model)
  .then(({ runModel }) => {
    const until = Date.now() + workerData.delayMs;
    while (Date.now() < until); // MODEL_DELAY_MS: a model computing
    return runModel(workerData.blastCase);
  })
  .then((answer) => parentPort.postMessage(answer));
`;

function runOnThread(blastCase: Record<string, unknown>): Promise<ModelAnswer> {
  return new Promise((resolve, reject) => {
    const worker = new Worker(MODEL_THREAD, {
      eval: true,
      workerData: {
        model: new URL('./model.mts', import.meta.url).href,
        delayMs: config.modelDelayMs,
        blastCase,
      },
    });
    worker.once('message', resolve);
    worker.once('error', reject); // the model threw
    worker.once('exit', (code) =>
      reject(new Error(`the model's thread exited with ${code}`)),
    );
  });
}

async function execute(
  stored: Stored,
  blastCase: Record<string, unknown> & { caseId: string },
): Promise<Stored> {
  const startedAt = now();
  const current = load(stored.run.id) ?? stored;
  if (isFinal(current.run.status)) return current; // cancelled before it began
  save({
    ...current,
    run: { ...current.run, status: 'running', updatedAt: startedAt },
  });
  let done: Stored;
  try {
    const answer = await runOnThread(blastCase);
    done = finish(stored, {
      status: 'succeeded',
      result: {
        specVersion: '0.1',
        caseId: blastCase.caseId,
        outputs: answer.outputs,
        unavailable: answer.unavailable,
        warnings: answer.warnings,
        provenance: {
          modelId: manifest.id,
          modelVersion: manifest.version,
          caseId: blastCase.caseId,
          startedAt,
          finishedAt: now(),
          inputsUsed: answer.inputsUsed,
          buildHash: config.buildHash,
        },
      },
    });
  } catch (error) {
    console.error('run failed', stored.run.id, error);
    done = finish(stored, {
      status: 'failed',
      error: {
        type: 'provider_error',
        code: 'model_error',
        message: 'The model failed on this case.',
      },
    });
  }
  // Answered 202, to its request or a repeat of it: call back.
  void deliver(done);
  return done;
}

// ---------------------------------------------------------------------------
// HTTP
// ---------------------------------------------------------------------------

function send(response: ServerResponse, status: number, body: unknown): void {
  response.writeHead(status, { 'content-type': 'application/json' });
  response.end(JSON.stringify(body));
}

function apiError(
  response: ServerResponse,
  status: number,
  type: string,
  code: string,
  message: string,
): void {
  send(response, status, { error: { type, code, message } });
}

/** The body, which the checks in the handler bound by its Content-Length. */
async function readBody(request: IncomingMessage): Promise<Buffer> {
  const chunks: Buffer[] = [];
  for await (const chunk of request) chunks.push(chunk as Buffer);
  return Buffer.concat(chunks);
}

async function createRun(
  request: IncomingMessage,
  response: ServerResponse,
  body: Buffer,
): Promise<void> {
  const idempotencyKey = String(request.headers['idempotency-key'] ?? '');
  if (!idempotencyKey)
    return apiError(
      response,
      400,
      'invalid_request',
      'missing_idempotency_key',
      'Send an Idempotency-Key.',
    );

  // A repeated key is the original run, never a second one. Answered 202,
  // it calls back when it ends, whatever its first request is told.
  if (existsSync(keyPath(idempotencyKey))) {
    const { runId } = JSON.parse(
      readFileSync(keyPath(idempotencyKey), 'utf8'),
    ) as { runId: string };
    const stored = load(runId);
    if (stored && isFinal(stored.run.status))
      return send(response, 200, stored.run);
    if (stored) {
      save({ ...stored, answeredAsync: true });
      return send(response, 202, stored.run);
    }
  }

  let parsed: {
    case?: Record<string, unknown> & { caseId?: unknown };
    callbackUrl?: unknown;
  };
  try {
    parsed = JSON.parse(body.toString('utf8'));
  } catch {
    return apiError(
      response,
      400,
      'invalid_request',
      'invalid_json',
      'The body is not JSON.',
    );
  }
  const blastCase = parsed.case;
  if (!blastCase || typeof blastCase.caseId !== 'string') {
    return apiError(
      response,
      400,
      'invalid_request',
      'missing_case',
      'The body carries no case.',
    );
  }

  const at = now();
  const stored: Stored = {
    run: {
      id: `run_${randomUUID()}`,
      status: 'queued',
      modelId: manifest.id,
      modelVersion: manifest.version,
      caseId: blastCase.caseId,
      createdAt: at,
      updatedAt: at,
    },
    idempotencyKey,
    callbackUrl:
      typeof parsed.callbackUrl === 'string' ? parsed.callbackUrl : null,
    answeredAsync: false,
  };
  save(stored);
  writeFileSync(
    keyPath(idempotencyKey),
    JSON.stringify({ runId: stored.run.id, at }),
  );

  const running = execute(
    stored,
    blastCase as Record<string, unknown> & { caseId: string },
  );
  const timer = new Promise<null>((resolve) =>
    setTimeout(() => resolve(null), config.syncBudgetMs),
  );
  const finished = await Promise.race([running, timer]);
  if (finished) return send(response, 200, finished.run);

  // Too slow to answer now: accept it, and call back when it is done.
  const current = load(stored.run.id) ?? stored;
  if (isFinal(current.run.status)) return send(response, 200, current.run);
  save({ ...current, answeredAsync: true });
  send(response, 202, current.run);
}

function log(
  request: IncomingMessage,
  status: number,
  startedAt: number,
): void {
  // What was asked and how it went; never a body.
  const line = {
    at: now(),
    method: request.method,
    path: (request.url ?? '').split('?')[0],
    status,
    ms: Date.now() - startedAt,
  };
  appendFileSync(
    join(config.dataDir, 'requests.jsonl'),
    `${JSON.stringify(line)}\n`,
  );
}

const server = createServer(async (request, response) => {
  const startedAt = Date.now();
  response.on('finish', () => log(request, response.statusCode, startedAt));
  try {
    const path = (request.url ?? '/').split('?')[0]!;
    if (!path.startsWith(`${config.basePath}/`))
      return apiError(
        response,
        404,
        'invalid_request',
        'not_found',
        'No such endpoint.',
      );
    const route = path.slice(config.basePath.length);

    if (route === '/healthz' && request.method === 'GET')
      return send(response, 200, { ok: true });

    // A body comes with a Content-Length, of 16 MiB at most. Anything else
    // is refused before it is read, and the connection closed.
    if (request.headers['transfer-encoding']) {
      response.setHeader('connection', 'close');
      return apiError(
        response,
        411,
        'invalid_request',
        'length_required',
        'Send the body with a Content-Length.',
      );
    }
    if (Number(request.headers['content-length'] ?? 0) > MAX_BODY_BYTES) {
      response.setHeader('connection', 'close');
      return apiError(
        response,
        413,
        'invalid_request',
        'too_large',
        'The body is over 16 MiB.',
      );
    }
    const body = await readBody(request);
    const problem = await signatureProblem(request, body);
    if (problem)
      return apiError(
        response,
        401,
        'authentication',
        problem,
        'The signature check failed.',
      );

    if (route === '/manifest' && request.method === 'GET')
      return send(response, 200, manifest);
    if (route === '/runs' && request.method === 'POST')
      return await createRun(request, response, body);

    const match = /^\/runs\/([A-Za-z0-9_-]{1,128})$/.exec(route);
    const stored = match ? load(match[1]!) : null;
    if (match && !stored)
      return apiError(
        response,
        404,
        'invalid_request',
        'run_not_found',
        'No such run.',
      );
    if (stored && request.method === 'GET')
      return send(response, 200, stored.run);
    if (stored && request.method === 'DELETE') {
      if (isFinal(stored.run.status)) return send(response, 200, stored.run);
      const cancelled = {
        ...stored,
        run: { ...stored.run, status: 'cancelled' as const, updatedAt: now() },
      };
      save(cancelled);
      return send(response, 200, cancelled.run);
    }
    return apiError(
      response,
      404,
      'invalid_request',
      'not_found',
      'No such endpoint.',
    );
  } catch (error) {
    console.error(error);
    if (!response.headersSent)
      apiError(
        response,
        500,
        'provider_error',
        'internal',
        'Something went wrong on the host.',
      );
  }
});

for (const file of readdirSync(join(config.dataDir, 'runs'))) {
  if (!file.endsWith('.json')) continue;
  let stored = load(file.slice(0, -'.json'.length));
  if (!stored) continue;
  // A run the host was working on when it stopped will not finish: say so.
  if (!isFinal(stored.run.status))
    stored = finish(stored, {
      status: 'failed',
      error: {
        type: 'provider_error',
        code: 'host_restarted',
        message: 'The host restarted during the run.',
      },
    });
  // And a callback the platform has not yet taken goes again.
  void deliver(stored);
}
sweep();
setInterval(sweep, 3600 * 1000).unref();
void refreshPlatformKeys();

server.listen(config.port, () => {
  const address = server.address();
  console.log(
    `listening on ${typeof address === 'object' && address ? address.port : config.port}`,
  );
});