import * as ort from "onnxruntime-web";
import type {
  FieldPrediction,
  FieldType,
  MLFieldFeatures,
} from "../shared/ml";

const LABELS: FieldType[] = ["username", "password", "email", "totp", "none"];

class BackgroundMLDetector {
  private session: ort.InferenceSession | null = null;
  private isLoading = false;
  private loadPromise: Promise<void> | null = null;

  async init(modelUrl?: string): Promise<void> {
    if (this.session) return;
    if (this.isLoading && this.loadPromise) {
      await this.loadPromise;
      return;
    }

    this.isLoading = true;
    this.loadPromise = this.loadModel(modelUrl);
    try {
      await this.loadPromise;
    } finally {
      this.isLoading = false;
    }
  }

  private async loadModel(modelUrl?: string): Promise<void> {
    ort.env.wasm.wasmPaths = chrome.runtime.getURL("assets/");
    ort.env.wasm.numThreads = 1;
    ort.env.wasm.simd = true;

    const url = modelUrl || chrome.runtime.getURL("models/form_detector.onnx");
    this.session = await ort.InferenceSession.create(url);
  }

  async predict(features: MLFieldFeatures): Promise<FieldPrediction> {
    if (!this.session) {
      await this.init();
    }
    if (!this.session) {
      throw new Error("ML model not initialized");
    }

    const vector = this.featuresToVector(features);
    const input = new ort.Tensor("float32", vector, [1, 45]);
    const inputName = this.session.inputNames[0];
    const outputName = this.session.outputNames[0];
    const results = await this.session.run({ [inputName]: input });
    const probs = results[outputName].data as Float32Array;

    let maxIndex = 0;
    let maxProb = probs[0];
    for (let i = 1; i < probs.length; i += 1) {
      if (probs[i] > maxProb) {
        maxProb = probs[i];
        maxIndex = i;
      }
    }

    return {
      type: LABELS[maxIndex],
      confidence: maxProb,
      probabilities: {
        username: probs[0],
        password: probs[1],
        email: probs[2],
        totp: probs[3],
        none: probs[4],
      },
    };
  }

  private featuresToVector(features: MLFieldFeatures): Float32Array {
    return new Float32Array([
      features.type_text,
      features.type_email,
      features.type_password,
      features.type_tel,
      features.type_number,
      features.type_search,
      features.type_url,
      features.type_other,
      features.auto_username,
      features.auto_email,
      features.auto_current_password,
      features.auto_new_password,
      features.auto_one_time_code,
      features.auto_off,
      features.auto_other,
      features.name_has_user,
      features.name_has_login,
      features.name_has_email,
      features.name_has_pass,
      features.name_length,
      features.id_has_user,
      features.id_has_login,
      features.id_has_email,
      features.id_has_pass,
      features.id_length,
      features.placeholder_has_user,
      features.placeholder_has_email,
      features.placeholder_has_pass,
      features.placeholder_length,
      features.aria_label_has_user,
      features.aria_label_has_email,
      features.aria_label_has_pass,
      features.aria_label_length,
      features.parent_is_form,
      features.parent_is_div,
      features.parent_is_section,
      features.sibling_count,
      features.has_password_sibling,
      features.has_email_sibling,
      features.form_has_submit,
      features.form_action_has_login,
      features.is_required,
      features.has_placeholder,
      features.has_aria_label,
      features.inputmode_numeric,
    ]);
  }
}

export const backgroundMLDetector = new BackgroundMLDetector();