import type { FieldPrediction, FieldType, MLFieldFeatures } from '../shared/ml';

// Output classes in the order the weight rows are written below.
const CLASSES: FieldType[] = ['password', 'email', 'username', 'totp', 'none'];

// Feature vector order — must stay in sync with ContentMLDetector.extractFeatures().
export const ML_FEATURE_KEYS: readonly (keyof MLFieldFeatures)[] = [
  // input type (8)
  'type_text', 'type_email', 'type_password', 'type_tel', 'type_number',
  'type_search', 'type_url', 'type_other',
  // autocomplete (7)
  'auto_username', 'auto_email', 'auto_current_password', 'auto_new_password',
  'auto_one_time_code', 'auto_off', 'auto_other',
  // name attribute (5)
  'name_has_user', 'name_has_login', 'name_has_email', 'name_has_pass', 'name_length',
  // id attribute (5)
  'id_has_user', 'id_has_login', 'id_has_email', 'id_has_pass', 'id_length',
  // placeholder (4)
  'placeholder_has_user', 'placeholder_has_email', 'placeholder_has_pass', 'placeholder_length',
  // aria-label (4)
  'aria_label_has_user', 'aria_label_has_email', 'aria_label_has_pass', 'aria_label_length',
  // DOM context (3)
  'parent_is_form', 'parent_is_div', 'parent_is_section',
  // siblings (3)
  'sibling_count', 'has_password_sibling', 'has_email_sibling',
  // form (2)
  'form_has_submit', 'form_action_has_login',
  // misc (3)
  'is_required', 'has_placeholder', 'has_aria_label',
  // input mode (1)
  'inputmode_numeric',
  // total: 45
];

// Biases — prior log-odds for each class.
// Login fields are rarer than generic inputs, so we start skeptical.
const B: number[] = [
  -1.5, // password
  -1.5, // email
  -2.0, // username
  -3.0, // totp
   0.0, // none
];

// W[class][feature] — one row per output class, 45 weights each.
// Groups within each row follow the FEATURES order above.
const W: number[][] = [
  // ── password ─────────────────────────────────────────────────────────────
  [
    // type
     0.0,  -1.0,  8.0, -0.5, -0.5, -2.0, -2.0, -1.0,
    // auto
    -1.0,  -1.0,  8.0,  7.0, -1.0,  0.5,  0.0,
    // name
    -0.5,   0.0, -0.5,  5.0,  0.0,
    // id
    -0.5,   0.0, -0.5,  5.0,  0.0,
    // placeholder
    -0.5,  -0.5,  4.0,  0.0,
    // aria-label
    -0.5,  -0.5,  4.0,  0.0,
    // parent
     0.3,   0.0,  0.0,
    // siblings
     0.2,  -0.5,  0.5,
    // form
     0.5,   0.5,
    // misc
     0.3,   0.2,  0.2,
    // inputmode
    -0.5,
  ],
  // ── email ─────────────────────────────────────────────────────────────────
  [
    // type
     0.3,   6.0, -2.0, -0.5, -1.0, -2.0, -1.0, -1.0,
    // auto
     1.0,   7.0, -2.0, -2.0, -2.0,  0.0,  0.0,
    // name
     0.0,   0.5,  4.5, -1.0,  0.0,
    // id
     0.0,   0.5,  4.5, -1.0,  0.0,
    // placeholder
     0.0,   4.0, -1.0,  0.0,
    // aria-label
     0.0,   4.0, -1.0,  0.0,
    // parent
     0.3,   0.0,  0.0,
    // siblings
     0.2,   0.5, -0.3,
    // form
     0.5,   0.5,
    // misc
     0.3,   0.2,  0.2,
    // inputmode
    -1.0,
  ],
  // ── username ──────────────────────────────────────────────────────────────
  [
    // type
     1.5,   0.5, -2.0, -0.5, -1.0, -2.0, -1.0, -1.0,
    // auto
     7.0,   1.0, -2.0, -2.0, -2.0,  0.0,  0.0,
    // name
     4.0,   3.5,  0.5, -1.0,  0.0,
    // id
     4.0,   3.5,  0.5, -1.0,  0.0,
    // placeholder
     3.5,   0.5, -0.5,  0.0,
    // aria-label
     3.5,   0.5, -0.5,  0.0,
    // parent
     0.3,   0.0,  0.0,
    // siblings
     0.2,   2.0,  0.5,
    // form
     0.5,   1.0,
    // misc
     0.3,   0.2,  0.2,
    // inputmode
    -1.0,
  ],
  // ── totp ──────────────────────────────────────────────────────────────────
  [
    // type
     0.0,  -1.0, -2.0,  2.0,  3.0, -2.0, -2.0, -1.0,
    // auto
    -2.0,  -2.0, -2.0, -2.0,  9.0,  0.5,  0.0,
    // name
    -0.5,  -0.5, -0.5, -0.5,  0.0,
    // id
    -0.5,  -0.5, -0.5, -0.5,  0.0,
    // placeholder
    -0.5,  -0.5, -0.5,  0.0,
    // aria-label
    -0.5,  -0.5, -0.5,  0.0,
    // parent
     0.2,   0.0,  0.0,
    // siblings
     0.2,  -0.3, -0.3,
    // form
     0.3,   0.0,
    // misc
     0.2,   0.1,  0.1,
    // inputmode
     4.0,
  ],
  // ── none ──────────────────────────────────────────────────────────────────
  [
    // type
     0.5,  -2.0, -3.0,  1.0,  1.5,  4.0,  4.0,  2.0,
    // auto
    -2.0,  -2.0, -3.0, -3.0, -3.0,  0.5,  0.2,
    // name
    -1.0,  -1.0, -1.0, -1.0,  0.1,
    // id
    -1.0,  -1.0, -1.0, -1.0,  0.1,
    // placeholder
    -0.5,  -0.5, -0.5,  0.0,
    // aria-label
    -0.5,  -0.5, -0.5,  0.0,
    // parent
    -0.5,   0.0,  0.0,
    // siblings
    -0.3,  -1.0, -0.5,
    // form
    -0.5,  -1.0,
    // misc
    -0.3,   0.0,  0.0,
    // inputmode
     1.0,
  ],
];

function softmax(logits: number[]): number[] {
  const max = Math.max(...logits);
  const exps = logits.map(x => Math.exp(x - max));
  const sum = exps.reduce((a, b) => a + b, 0);
  return exps.map(x => x / sum);
}

function predict(features: MLFieldFeatures): FieldPrediction {
  const vec = ML_FEATURE_KEYS.map(k => features[k]);
  const logits = W.map((row, ci) =>
    B[ci] + row.reduce((s, w, fi) => s + w * vec[fi], 0)
  );
  const probs = softmax(logits);
  const maxIdx = probs.indexOf(Math.max(...probs));
  return {
    type: CLASSES[maxIdx],
    confidence: probs[maxIdx],
    probabilities: Object.fromEntries(
      CLASSES.map((t, i) => [t, probs[i]])
    ) as Record<FieldType, number>,
  };
}

export function isValidMLFeatureBatch(value: unknown, maxFields: number): value is MLFieldFeatures[] {
  return Array.isArray(value) && value.length <= maxFields && value.every((item) =>
    typeof item === "object" && item !== null && ML_FEATURE_KEYS.every((key) =>
      typeof (item as Record<string, unknown>)[key] === "number" &&
      Number.isFinite((item as Record<string, unknown>)[key])
    )
  );
}

export function predictBatch(featuresArray: MLFieldFeatures[]): FieldPrediction[] {
  return featuresArray.map(predict);
}