/* API CONFIGURATION */
let workspaceID = "<WORKSPACE_ID>"; // Replace with your workspace ID
let BASE_API_URL = "/api";

/* EXTRA CONFIGURATION FOR MODEL TRAINING FIELDS */

/*
  additional_traits:
  An array of strings representing additional traits.
  Example: ["AI_advisor", "AI_answering", "pro", "starter"]
*/
let ADDITIONAL_TRAITS = [];

/*
  known_positive_factors:
  An array of strings representing factors known to contribute positively.
  Example: ["activity:6_mo_cutoff:Had Online Orders", "activity:6_mo_cutoff:Hit 1K Followers"]
*/
let KNOWN_POSITIVE_FACTORS = [];

/*
  known_negative_factors:
  An array of strings representing factors known to have negative effects.
  Example: ["activity:6_mo_cutoff:Negative Reviews", "trait:open_churn_case=true"]
*/
let KNOWN_NEGATIVE_FACTORS = [];

/*
  exclude_factors:
  An array of strings representing factors to exclude from training.
  Example: ["property:last_month_impressions", "activity:6_mo_cutoff:Created User"]
*/
let EXCLUDE_FACTORS = [];

/* 
  Create safe fallbacks so that if any of the above are missing or not arrays,
  they default to an empty array.
*/
let safeAdditionalTraits = Array.isArray(ADDITIONAL_TRAITS)
  ? ADDITIONAL_TRAITS
  : [];
let safeKnownPositiveFactors = Array.isArray(KNOWN_POSITIVE_FACTORS)
  ? KNOWN_POSITIVE_FACTORS
  : [];
let safeKnownNegativeFactors = Array.isArray(KNOWN_NEGATIVE_FACTORS)
  ? KNOWN_NEGATIVE_FACTORS
  : [];
let safeExcludeFactors = Array.isArray(EXCLUDE_FACTORS)
  ? EXCLUDE_FACTORS
  : [];

// Helper functions
let getCombinedConsideredIds = (audiences, type) => {
  let churnAudience = audiences.find((aud) => aud.name === "Churn");
  let retentionAudience = audiences.find((aud) => aud.name === "Retention");

  if (type === "user") {
    let churnUserIds = churnAudience?.data?.matching_fs_user_ids || [];
    let retentionUserIds = retentionAudience?.data?.matching_fs_user_ids || [];
    return [...new Set([...churnUserIds, ...retentionUserIds])];
  } else {
    let churnAccountIds = churnAudience?.data?.matching_account_ids || [];
    let retentionAccountIds = retentionAudience?.data?.matching_account_ids || [];
    return [...new Set([...churnAccountIds, ...retentionAccountIds])];
  }
};

let buildProductTargets = (audiences, modelType) => {
  let productTargets = {};
  let allProductAccountIds = new Set();

  for (let audience of audiences) {
    if (!audience.hidden) continue;

    let productFSID = audience.product_fs_id;
    if (!productFSID) continue;

    let isTargetModel = (modelType === "churn" && audience.name.endsWith(" Churn")) ||
                       (modelType === "retention" && audience.name.endsWith(" Retention"));
    if (!isTargetModel || !audience.data?.matching_account_ids) continue;

    productTargets[productFSID] = audience.data.matching_account_ids;
    audience.data.matching_account_ids.forEach(id => allProductAccountIds.add(id));
  }

  return Object.keys(productTargets).length > 0 ? {
    productTargets,
    productConsideredIds: [...allProductAccountIds]
  } : null;
};

let trainModel = async (audience, model, consideredIds, productTargetsData, productFsIds, type) => {
  isUser = type === "user";
  let url = `${BASE_API_URL}/internal/prediction/models/${isUser ? "user_" : ""}${model}/train`;
  let targetKey = isUser ? "target_user_fsids" : "target_account_ids";
  let consideredKey = isUser ? "considered_user_fsids" : "considered_account_ids";

  let finalConsideredIds = productTargetsData?.productConsideredIds || consideredIds;

  let body = {
    [targetKey]: audience.data[targetKey.replace("target_", "matching_")],
    [consideredKey]: finalConsideredIds,
    params: {
      num_trees: 5,
      tree_depth: 50,
      learning_rate: 0.1,
      iterations: 100,
      lambda: 0.0,
      regularizer: "ridge",
      ensemble_weights: { lr: 0.5, rf: 0.5 }
    },
    additional_traits: ADDITIONAL_TRAITS,
    known_positive_factors: KNOWN_POSITIVE_FACTORS,
    known_negative_factors: KNOWN_NEGATIVE_FACTORS,
    exclude_factors: EXCLUDE_FACTORS
  };

  if (type === "account" && productFsIds.length > 0) {
    body.product_fs_ids = productFsIds;
    if (productTargetsData?.productTargets) {
      body.product_targets = productTargetsData.productTargets;
    }
  }

  await fetch(url, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "fs-workspace-id": workspaceID
    },
    body: JSON.stringify(body)
  });
};

// Main training function
let trainWithProductAudiences = async (type) => {
  // Fetch data
  let [audiencesRes, productsRes] = await Promise.all([
    fetch(`${BASE_API_URL}/audiences?include_hidden=true`, {
      headers: { "Content-Type": "application/json", "fs-workspace-id": workspaceID }
    }),
    type === "account" ? fetch(`${BASE_API_URL}/products`, {
      headers: { "Content-Type": "application/json", "fs-workspace-id": workspaceID }
    }) : Promise.resolve(null)
  ]);

  let audiencesData = await audiencesRes.json();
  let audiences = audiencesData.response.audiences;
  let products = productsRes ? (await productsRes.json()).response?.products || [] : [];
  let productFsIds = products.map(p => p.fs_id);

  // Get considered IDs from global Churn/Retention audiences
  let consideredIds = getCombinedConsideredIds(audiences, type);

  // Filter training audiences (global Churn/Retention)
  let trainingAudiences = audiences.filter(aud =>
    aud.hidden && (aud.name === "Churn" || aud.name === "Retention")
  );

  // Train models concurrently
  let trainPromises = trainingAudiences.map(async (audience) => {
    let model = audience.name.toLowerCase();
    let productTargetsData = buildProductTargets(audiences, model);
    return trainModel(audience, model, consideredIds, productTargetsData, productFsIds, type);
  });

  await Promise.all(trainPromises);

  // Update scores
  let updateUrl = `${BASE_API_URL}/internal/prediction/scores/${type === "user" ? "users" : "accounts"}/update`;
  await fetch(updateUrl, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "fs-workspace-id": workspaceID
    }
  });
};