import { env, pipeline, type TextClassificationPipeline } from "@huggingface/transformers";
import { hashText } from "../shared/hash";
import type { InferenceRequest, InferenceResult } from "../shared/types";
import { loadModelManifest, type ModelManifest } from "./model-metadata";

let classifier: TextClassificationPipeline | undefined;
let manifest: ModelManifest | undefined;

// Transformers.js disables local files in browser workers by default. Our model
// is packaged inside the extension, so never fall back to the network.
env.allowLocalModels = true;
env.allowRemoteModels = false;
env.useBrowserCache = false;

export async function classify(requests: InferenceRequest[], modelBaseUrl: string): Promise<InferenceResult[]> {
  const modelManifest = manifest ??= await loadModelManifest(modelBaseUrl);
  classifier ??= await loadClassifier(modelBaseUrl);
  const invoke = classifier as unknown as (texts: string[], options: Record<string, unknown>) => Promise<Array<Array<{ label: string; score: number }> | { label: string; score: number }>>;
  const startedAt = performance.now();
  // Transformers.js defaults `top_k` to 1. Request both scores because the
  // filtering threshold is based on the toxic class probability, including
  // when the non-toxic class is the model's highest-confidence prediction.
  const outputs = await invoke(requests.map((request) => request.text), { top_k: 2, max_length: modelManifest.maxLength });
  const modelDurationMs = performance.now() - startedAt;
  return requests.map((request, index) => {
    const output = Array.isArray(outputs[index]) ? outputs[index] : [outputs[index]];
    const toxic = output.find((item) => resolveOutputIndex(item?.label, modelManifest) === modelManifest.labels.toxic);
    if (!toxic) throw new Error(`Classifier output did not contain toxic label index ${modelManifest.labels.toxic}`);
    const probability = Math.max(0, Math.min(1, Number(toxic?.score ?? 0)));
    return {
      requestId: request.requestId,
      id: request.id,
      textHash: hashText(request.text),
      label: probability >= 0.5 ? "toxic" : "not_toxic",
      probability,
      navigationId: request.navigationId,
      modelRevision: modelManifest.revision,
      modelDurationMs
    };
  });
}

async function loadClassifier(modelBaseUrl: string): Promise<TextClassificationPipeline> {
  if (!env.backends.onnx.wasm) throw new Error("ONNX WASM backend is unavailable.");
  env.backends.onnx.wasm.wasmPaths = { wasm: new URL("../../ort/ort-wasm-simd-threaded.jsep.wasm", modelBaseUrl).href };
  return (pipeline as unknown as (task: string, model: string, options: { device: "wasm"; dtype: "q8"; local_files_only: boolean }) => Promise<TextClassificationPipeline>)("text-classification", modelBaseUrl, { device: "wasm", dtype: "q8", local_files_only: true });
}

export function resolveOutputIndex(label: string | undefined, metadata: ModelManifest): number | undefined {
  if (!label) return undefined;
  const normalized = label.toLowerCase().replace(/[\s-]+/g, "_");
  if (normalized === "toxic") return metadata.labels.toxic;
  if (normalized === "not_toxic" || normalized === "non_toxic" || normalized === "non-toxic") return metadata.labels.nonToxic;
  const configured = metadata.labels.names[normalized] ?? metadata.labels.names[label];
  if (configured !== undefined) return configured;
  const match = normalized.match(/^label_(\d+)$/);
  return match ? Number(match[1]) : undefined;
}
