function byteLength(value) {
  return Buffer.byteLength(String(value || ''), 'utf8');
}

function unique(values) {
  return Array.from(new Set(values.filter(Boolean)));
}

function countMatches(text, regex) {
  return (text.match(regex) || []).length;
}

function maxParenDepth(text) {
  let depth = 0;
  let maxDepth = 0;
  for (const ch of text) {
    if (ch === '(') {
      depth += 1;
      maxDepth = Math.max(maxDepth, depth);
    } else if (ch === ')') {
      depth = Math.max(0, depth - 1);
    }
  }
  return maxDepth;
}

function numericConstants(text) {
  return text.match(/(?<![A-Za-z_])-?\d+(?:\.\d+)?/g) || [];
}

function constantBitLength(constants) {
  return constants.reduce((sum, value) => {
    const n = Math.abs(Number(value));
    if (!Number.isFinite(n)) return sum + 64;
    return sum + Math.max(1, Math.ceil(Math.log2(n + 1)));
  }, 0);
}

function functionNames(text) {
  return unique((text.match(/\b[A-Za-z_][A-Za-z0-9_]*\s*\(/g) || []).map((item) => item.replace(/\s*\($/, '')));
}

function directCoordinateEqualityCount(text) {
  return (
    countMatches(text, /\b[xy]\s*(?:==|===)\s*-?\d+/g) +
    countMatches(text, /-?\d+\s*(?:==|===)\s*\b[xy]\b/g)
  );
}

function estimateDegree(text) {
  const normalized = text.replace(/\s+/g, '');
  let degree = 1;
  if (/(x\*x|y\*y|x\*y|y\*x|\*\*2|\^2)/.test(normalized)) degree = Math.max(degree, 2);
  if (/(\*\*3|\^3|x\*x\*x|y\*y\*y)/.test(normalized)) degree = Math.max(degree, 3);
  if (/(\*\*4|\^4)/.test(normalized)) degree = Math.max(degree, 4);
  return degree;
}

function tableValueCount(text) {
  let total = 0;
  const calls = text.match(/\btable\s*\(([^)]*)\)/g) || [];
  for (const call of calls) {
    const inner = call.replace(/^\s*table\s*\(/, '').replace(/\)\s*$/, '');
    const parts = inner.split(',').map((part) => part.trim()).filter(Boolean);
    total += Math.max(0, parts.length - 1);
  }
  return total;
}

function allowedStructuresForStep(step) {
  return unique([
    step.primary_structure,
    ...(Array.isArray(step.structure_bundle) ? step.structure_bundle : []),
    ...(Array.isArray(step.secondary_structure) ? step.secondary_structure : []),
    ...(Array.isArray(step.secondary_structures) ? step.secondary_structures : []),
    ...(Array.isArray(step.components)
      ? step.components.flatMap((component) => [
          component.primary_structure,
          ...(Array.isArray(component.secondary_structure) ? component.secondary_structure : []),
        ])
      : []),
    ...(Array.isArray(step.formula_template?.helpers) && step.formula_template.helpers.length ? ['HELPER_EXPRESSIONS'] : []),
    ...(Array.isArray(step.components) &&
    step.components.some((component) => Array.isArray(component.formula_template?.helpers) && component.formula_template.helpers.length)
      ? ['HELPER_EXPRESSIONS']
      : []),
    'S05_MOD_RESIDUE',
  ]);
}

function outputPipelineFunctions(step) {
  const functions = new Set();
  const pipeline = step.output_pipeline || {};
  for (const stage of pipeline.static_stages || []) {
    const id = stage.operator_id;
    if (
      [
        'absolute',
        'divide',
        'floor',
        'ceil',
        'round',
        'mod',
        'sqrt',
        'power',
        'norm',
        'integral',
        'derivative',
        'partial_derivative',
        'gradient',
        'rank',
        'series',
        'laplace_transform',
        'fourier_transform',
        'inverse',
        'eigenvalue',
        'laplacian_operator',
        'limit',
      ].includes(id)
    ) {
      functions.add('abs');
    }
    if (['divide', 'floor', 'integral', 'laplace_transform', 'inverse', 'eigenvalue', 'limit'].includes(id)) {
      functions.add('floor');
    }
    if (id === 'ceil') functions.add('ceil');
    if (id === 'round') functions.add('round');
    if (id === 'sqrt' || id === 'norm') functions.add('sqrt');
    if (id === 'fourier_transform') functions.add('table');
    if (['multiply', 'derivative', 'partial_derivative', 'rank'].includes(id)) functions.add('if');
  }
  for (const stage of pipeline.dynamic_stages || []) {
    if (stage.parameter?.family === 'TIME_CLOCK') functions.add('if');
    if (stage.parameter?.family === 'HALF_LIFE_DECAY') {
      functions.add('floor');
      functions.add('min');
    }
    if (stage.parameter?.family === 'CHAIN_STATE') functions.add('abs');
  }
  return Array.from(functions);
}

function budgetForDifficulty(difficulty, policy) {
  const d = Number(difficulty || 1);
  const match = (policy.by_difficulty || []).find(
    (row) => d >= row.difficulty_min && d <= row.difficulty_max,
  );
  return {
    ...(policy.default_budget || {}),
    ...((match && match.budget) || {}),
  };
}

function answerBudgetForStep(step, rules) {
  const policy = rules.answer_budget_policy || {};
  const budget = budgetForDifficulty(step.target_difficulty, policy);
  const compositionCount = Math.max(1, Number(step.composition_count || 1));
  const structures = allowedStructuresForStep(step);
  const highOrderCount = structures.filter((structure) =>
    [
      'S09_LOOKUP_PERMUTATION',
      'S11_DECAY_TIME',
      'S12_GRAPH_RULE',
      'S13_BOARD_RULE',
      'S14_ENCODING_HASH',
      'S15_PROBABILITY_SEEDED',
      'S16_STATE_MEMORY',
      'S17_FSM_DYNAMIC',
      'S18_RECURRENCE_SEQUENCE',
      'S19_CONVOLUTION_LOCAL',
      'S20_OPTIMIZATION_OBJECTIVE',
      'S21_LANGUAGE_SYMBOL',
      'S22_BIO_REACTION_GROWTH',
      'S23_ECON_GAME_SYSTEM',
      'S24_COMPOSITE_BOSS',
      'S25_CHAIN_STATE_DEPENDENT',
      'S26_SPECIAL_CONSTANT_FIELD',
      'S27_PRIME_ARITHMETIC',
      'S28_SPECTRAL_TRANSFORM',
      'S29_GROUP_SYMMETRY',
      'S30_COMPLEX_ANALYTIC',
      'S31_TENSOR_GEOMETRY',
      'S32_VARIATIONAL_HAMILTONIAN',
      'S33_FIELD_EQUATION',
      'S34_QUANTUM_OPERATOR',
      'S35_FLUID_NAVIER_STOKES',
      'S36_STATISTICAL_MECHANICS',
      'S37_GAUGE_LATTICE_FIELD',
    ].includes(structure),
  ).length;
  const compositionExtra = Math.max(0, compositionCount - 1);
  const highExtra = Math.max(0, highOrderCount);
  const pipeline = step.output_pipeline || {};
  const pipelineExtra = Math.max(0, (pipeline.static_stages || []).length + (pipeline.dynamic_stages || []).length);
  if (compositionExtra === 0 && highExtra === 0 && pipelineExtra === 0) return budget;
  return {
    ...budget,
    bytes: budget.bytes + compositionExtra * 1800 + highExtra * 2600 + pipelineExtra * 520,
    tokens: budget.tokens + compositionExtra * 720 + highExtra * 760 + pipelineExtra * 180,
    constants: budget.constants + compositionExtra * 260 + highExtra * 240 + pipelineExtra * 60,
    constant_bits: budget.constant_bits + compositionExtra * 2400 + highExtra * 2200 + pipelineExtra * 640,
    branches: budget.branches + compositionExtra * 18 + highExtra * 18 + pipelineExtra * 4,
    max_depth: budget.max_depth + compositionCount * 4 + highExtra * 6 + pipelineExtra * 3,
    table_values: budget.table_values + compositionExtra * 64 + highExtra * 128 + pipelineExtra * 12,
    total_cost: budget.total_cost + compositionExtra * 3600 + highExtra * 3600 + pipelineExtra * 900,
  };
}

function allowedFunctionsForStep(step, rules) {
  const policy = rules.answer_budget_policy?.feature_policy || {};
  const allowed = new Set(policy.always_allowed_functions || []);
  const structures = allowedStructuresForStep(step);
  for (const structure of structures) {
    for (const fn of policy.structure_functions?.[structure] || []) {
      allowed.add(fn);
    }
  }
  for (const fn of outputPipelineFunctions(step)) allowed.add(fn);
  return Array.from(allowed);
}

function analyzeExpression(expr) {
  const text = String(expr || '');
  const constants = numericConstants(text);
  const functions = functionNames(text).filter((name) => name !== 'mod');
  const identifiers = text.match(/\b[A-Za-z_][A-Za-z0-9_]*\b/g) || [];
  const operators = countMatches(text, /(\+|-|\*|\/|%|\bmod\b|\^|==|>=|<=|>|<|&&|\|\|)/gi);

  return {
    bytes: byteLength(text),
    tokens: identifiers.length + constants.length + operators,
    operators,
    functions,
    constants: constants.length,
    constant_bits: constantBitLength(constants),
    branches: countMatches(text, /\bif\s*\(/gi),
    max_depth: maxParenDepth(text),
    degree: estimateDegree(text),
    table_values: tableValueCount(text),
    direct_coordinate_equalities: directCoordinateEqualityCount(text.replace(/\s+/g, '')),
  };
}

function mergeMetrics(uMetrics, vMetrics, answer) {
  const bytes = byteLength(answer.u) + byteLength(answer.v);
  return {
    bytes,
    formulas: 2,
    tokens: uMetrics.tokens + vMetrics.tokens,
    operators: uMetrics.operators + vMetrics.operators,
    functions: unique([...uMetrics.functions, ...vMetrics.functions]),
    constants: uMetrics.constants + vMetrics.constants,
    constant_bits: uMetrics.constant_bits + vMetrics.constant_bits,
    branches: uMetrics.branches + vMetrics.branches,
    max_depth: Math.max(uMetrics.max_depth, vMetrics.max_depth),
    degree: Math.max(uMetrics.degree, vMetrics.degree),
    table_values: uMetrics.table_values + vMetrics.table_values,
    direct_coordinate_equalities:
      uMetrics.direct_coordinate_equalities + vMetrics.direct_coordinate_equalities,
  };
}

function complexityCost(metrics, rules) {
  const weights = rules.answer_budget_policy?.cost_weights || {};
  return Math.ceil(
    metrics.tokens * (weights.token || 1) +
      metrics.operators * (weights.operator || 1) +
      metrics.constants * (weights.constant || 1) +
      metrics.constant_bits * (weights.constant_bit || 0) +
      metrics.branches * (weights.branch || 1) +
      metrics.max_depth * (weights.depth || 1) +
      Math.max(0, metrics.degree - 1) * (weights.degree || 1) +
      metrics.table_values * (weights.table_value || 1) +
      metrics.functions.length * (weights.function || 1),
  );
}

function validateAnswerBudget({ answer, step, rules }) {
  if (!answer || typeof answer.u !== 'string' || typeof answer.v !== 'string') {
    return { ok: false, stage: 'budget', errors: ['Answer must include string expressions for u and v'] };
  }

  const budget = answerBudgetForStep(step, rules);
  const allowedFunctions = allowedFunctionsForStep(step, rules);
  const allowedFunctionSet = new Set(allowedFunctions);
  const uMetrics = analyzeExpression(answer.u);
  const vMetrics = analyzeExpression(answer.v);
  const metrics = mergeMetrics(uMetrics, vMetrics, answer);
  metrics.total_cost = complexityCost(metrics, rules);

  const errors = [];
  const checkMax = (key, label) => {
    const limit = budget[key];
    if (Number.isFinite(Number(limit)) && metrics[key] > Number(limit)) {
      errors.push(`${label} exceeds budget: ${metrics[key]} > ${limit}`);
    }
  };

  checkMax('bytes', 'byte count');
  checkMax('formulas', 'formula count');
  checkMax('tokens', 'token count');
  checkMax('constants', 'constant count');
  checkMax('constant_bits', 'total constant bits');
  checkMax('branches', 'conditional branch count');
  checkMax('max_depth', 'nesting depth');
  checkMax('degree', 'maximum degree');
  checkMax('table_values', 'table value count');
  checkMax('total_cost', 'complexity cost');

  const hardReject = rules.answer_budget_policy?.hard_reject || {};
  const directLimit = Number(hardReject.direct_coordinate_equality_limit ?? 0);
  if (metrics.direct_coordinate_equalities > directLimit) {
    errors.push(`Direct coordinate equality checks exceed limit: ${metrics.direct_coordinate_equalities} > ${directLimit}`);
  }

  const disabled = new Set(rules.answer_budget_policy?.feature_policy?.disabled_functions || []);
  for (const fn of metrics.functions) {
    if (disabled.has(fn) && !allowedFunctionSet.has(fn)) {
      errors.push(`Function ${fn} is currently disabled`);
    } else if (!allowedFunctionSet.has(fn)) {
      errors.push(`Function ${fn} is not allowed for the current Step structure`);
    }
  }

  return {
    ok: errors.length === 0,
    stage: 'budget',
    errors,
    budget,
    metrics,
    allowed_structures: allowedStructuresForStep(step),
    allowed_functions: allowedFunctions,
  };
}

module.exports = {
  answerBudgetForStep,
  allowedFunctionsForStep,
  allowedStructuresForStep,
  analyzeExpression,
  validateAnswerBudget,
};
