import { performance } from "node:perf_hooks";
import { initTracePrompt as initCfgAsync, ConfigManager } from "./config";
import { encryptBuffer } from "./crypto/encryptor";
import { computeLeaf } from "./crypto/hasher";
import { countTokens } from "./utils/tokenCounter";
import { PersistentBatcher as Batcher } from "./queue/persistentBatcher";
const stringify = require("json-stable-stringify") as (v: any) => string;
import type { TracePromptInit, WrapOpts, EncryptedBundle } from "./types";

import { registry } from "./metrics";
import { Histogram } from "prom-client";

const wrapperLatencyHist = new Histogram({
  name: "traceprompt_llm_wrapper_latency_ms",
  help: "End‑to‑end latency from prompt send to response receive in the SDK wrapper (ms)",
  buckets: [50, 100, 250, 500, 1000, 2000, 5000],
  registers: [registry],
});

export async function initTracePrompt(
  cfg?: Partial<TracePromptInit>
): Promise<void> {
  await initCfgAsync(cfg);
}

export function wrapLLM<P extends Record<string, any>, R>(
  originalFn: (prompt: string, params?: P) => Promise<R>,
  meta: WrapOpts
): (prompt: string, params?: P) => Promise<R> {
  const staticMeta = ConfigManager.cfg.staticMeta;

  return async function wrapped(prompt: string, params?: P): Promise<R> {
    const t0 = performance.now();
    const result = await originalFn(prompt, params);
    const t1 = performance.now();

    wrapperLatencyHist.observe(t1 - t0);

    const plaintextJson = JSON.stringify({
      prompt,
      response: result,
    });

    const enc = await encryptBuffer(Buffer.from(plaintextJson, "utf8"));

    const payload = {
      ...staticMeta,
      orgId: ConfigManager.cfg.orgId,
      modelVendor: meta.modelVendor,
      modelName: meta.modelName,
      userId: meta.userId,
      ts_client: new Date().toISOString(),
      latency_ms: +(t1 - t0).toFixed(2),
      prompt_tokens: countTokens(prompt),
      response_tokens: countTokens(
        typeof result === "string" ? result : JSON.stringify(result)
      ),
      enc,
    };

    const leafHash = computeLeaf(stringify(payload));
    Batcher.enqueue({ payload, leafHash });

    return result;
  };
}