import fs from "node:fs";
import path from "node:path";
import vm from "node:vm";
import Ajv2020 from "ajv/dist/2020.js";

export const root = path.resolve(import.meta.dirname, "..");
export const profileDir = path.join(root, "profiles", "models");
export const profileSchemaPath = path.join(
  root,
  "schemas",
  "model-profile.schema.json",
);

export function profileIdForRepo(repo) {
  return String(repo)
    .split("/")
    .map((part) =>
      part
        .toLowerCase()
        .replace(/[^a-z0-9]+/g, "-")
        .replace(/^-+|-+$/g, ""),
    )
    .join("--");
}

export function loadModelCatalog() {
  const dataPath = path.join(root, "assets", "local-frontier-model-data.js");
  const context = { window: {} };
  vm.runInNewContext(fs.readFileSync(dataPath, "utf8"), context, {
    filename: dataPath,
  });
  return context.window.LOCAL_FRONTIER_MODEL_DATA;
}

export function listProfileFiles() {
  if (!fs.existsSync(profileDir)) return [];
  return fs
    .readdirSync(profileDir)
    .filter((file) => file.endsWith(".json"))
    .sort()
    .map((file) => path.join(profileDir, file));
}

export function readJson(filePath) {
  return JSON.parse(fs.readFileSync(filePath, "utf8"));
}

export function loadProfiles() {
  return listProfileFiles().map((filePath) => ({
    filePath,
    profile: readJson(filePath),
  }));
}

export function validateProfiles() {
  const schema = readJson(profileSchemaPath);
  const ajv = new Ajv2020({
    allErrors: true,
    strict: true,
    strictRequired: false,
  });
  const validate = ajv.compile(schema);
  const errors = [];
  const catalog = loadModelCatalog();
  const catalogIds = new Set((catalog.models || []).map((model) => model.id));
  const seenIds = new Set();
  const seenRepos = new Set();
  const profiles = loadProfiles();

  for (const { filePath, profile } of profiles) {
    const relativePath = path.relative(root, filePath);
    if (!validate(profile)) {
      for (const error of validate.errors || []) {
        errors.push(`${relativePath}${error.instancePath}: ${error.message}`);
      }
      continue;
    }

    const stem = path.basename(filePath, ".json");
    if (stem !== profile.id) {
      errors.push(`${relativePath}: file stem must match id ${profile.id}`);
    }

    const expectedId = profileIdForRepo(profile.repo);
    if (profile.id !== expectedId) {
      errors.push(
        `${relativePath}: id must be ${expectedId} for ${profile.repo}`,
      );
    }

    if (seenIds.has(profile.id)) {
      errors.push(`${relativePath}: duplicate profile id ${profile.id}`);
    }
    seenIds.add(profile.id);

    if (seenRepos.has(profile.repo)) {
      errors.push(`${relativePath}: duplicate repo profile ${profile.repo}`);
    }
    seenRepos.add(profile.repo);

    if (!catalogIds.has(profile.repo)) {
      errors.push(
        `${relativePath}: repo is not present in scraped model catalog`,
      );
    }

    checkAuditedProfile(relativePath, profile, errors);
  }

  return {
    errors,
    profiles: profiles.map(({ profile }) => profile),
    catalogCount: (catalog.models || []).length,
  };
}

function checkAuditedProfile(relativePath, profile, errors) {
  if (profile.status !== "audited") return;

  if (!Number.isInteger(profile.architecture.max_context_tokens)) {
    errors.push(`${relativePath}: audited profiles require max context tokens`);
  }

  if (profile.serving.weight_format === "unknown") {
    errors.push(`${relativePath}: audited profiles cannot use unknown weights`);
  }
  if (
    profile.serving.kv_store_format === "unknown" ||
    profile.serving.kv_read_format === "unknown"
  ) {
    errors.push(
      `${relativePath}: audited profiles cannot use unknown KV formats`,
    );
  }
  if (profile.architecture.kv_adapter.kind === "unknown") {
    errors.push(`${relativePath}: audited profiles cannot use unknown KV`);
  }

  const weightAdapter = profile.architecture.weight_adapter;
  if (weightAdapter.kind === "dense_resident_swept") {
    if (weightAdapter.swept_params_b > weightAdapter.resident_params_b) {
      errors.push(`${relativePath}: swept params exceed resident params`);
    }
    if (
      weightAdapter.swept_weight_gb !== undefined &&
      weightAdapter.resident_weight_gb !== undefined &&
      weightAdapter.swept_weight_gb > weightAdapter.resident_weight_gb
    ) {
      errors.push(`${relativePath}: swept weight exceeds resident weight`);
    }
    if (weightAdapter.auxiliary_resident_params_b !== undefined) {
      const expectedAuxiliary =
        weightAdapter.resident_params_b - weightAdapter.swept_params_b;
      if (
        Math.abs(
          weightAdapter.auxiliary_resident_params_b - expectedAuxiliary,
        ) > 1e-9
      ) {
        errors.push(
          `${relativePath}: auxiliary resident params must equal resident minus swept params`,
        );
      }
    }
    if (
      weightAdapter.auxiliary_resident_weight_gb !== undefined &&
      weightAdapter.resident_weight_gb !== undefined &&
      weightAdapter.swept_weight_gb !== undefined
    ) {
      const expectedAuxiliary =
        weightAdapter.resident_weight_gb - weightAdapter.swept_weight_gb;
      if (
        Math.abs(
          weightAdapter.auxiliary_resident_weight_gb - expectedAuxiliary,
        ) > 1e-9
      ) {
        errors.push(
          `${relativePath}: auxiliary resident weight must equal resident minus swept weight`,
        );
      }
    }
    return;
  }
  if (weightAdapter.kind === "moe_distinct_experts_exact") {
    if (weightAdapter.routed_experts_per_token > weightAdapter.routed_experts) {
      errors.push(
        `${relativePath}: routed experts per token exceed expert count`,
      );
    }
    const activeWeight =
      weightAdapter.fixed_weight_gb +
      weightAdapter.routed_experts_per_token *
        weightAdapter.routed_expert_weight_gb;
    if (activeWeight > weightAdapter.resident_weight_gb + 1e-9) {
      errors.push(`${relativePath}: active traffic exceeds resident footprint`);
    }
    if (
      weightAdapter.main_resident_weight_gb !== undefined &&
      weightAdapter.auxiliary_resident_weight_gb !== undefined
    ) {
      const expectedResident =
        weightAdapter.main_resident_weight_gb +
        weightAdapter.auxiliary_resident_weight_gb;
      if (
        Math.abs(expectedResident - weightAdapter.resident_weight_gb) > 1e-9
      ) {
        errors.push(
          `${relativePath}: main plus auxiliary resident weight must equal resident weight`,
        );
      }
    }
    return;
  }
  if (weightAdapter.kind !== "moe_distinct_experts") return;

  if (weightAdapter.active_params_b > weightAdapter.total_params_b) {
    errors.push(`${relativePath}: active params exceed total params`);
  }
  if (weightAdapter.routed_experts_per_token > weightAdapter.routed_experts) {
    errors.push(
      `${relativePath}: routed experts per token exceed expert count`,
    );
  }

  const denominator =
    weightAdapter.routed_experts - weightAdapter.routed_experts_per_token;
  const expertParamB =
    weightAdapter.expert_param_b ??
    (weightAdapter.total_params_b - weightAdapter.active_params_b) /
      Math.max(1, denominator);
  const fixedParamB =
    weightAdapter.fixed_param_b ??
    Math.max(
      0,
      weightAdapter.active_params_b -
        weightAdapter.routed_experts_per_token * expertParamB,
    );
  const residentGb =
    weightAdapter.total_params_b * profile.serving.weight_bytes_per_param;
  const activeGb =
    (fixedParamB + weightAdapter.routed_experts_per_token * expertParamB) *
    profile.serving.weight_bytes_per_param;

  if (expertParamB <= 0) {
    errors.push(`${relativePath}: derived expert params must be positive`);
  }
  if (activeGb > residentGb + 1e-9) {
    errors.push(`${relativePath}: active traffic exceeds resident footprint`);
  }
}