// PRE Web GUI — GBNF grammar builder for llama.cpp
// Converts OpenAI-format tool definitions (JSON Schema) into GBNF grammars
// that constrain the sampler to produce valid JSON matching the schema.

/**
 * Convert a JSON Schema to a GBNF grammar string.
 * Supports: string, integer, number, boolean, null, object, array, enum.
 *
 * @param {Object} schema - JSON Schema object
 * @param {string} [ruleName='root'] - Name of the root rule
 * @returns {string} GBNF grammar string
 */
function jsonSchemaToGbnf(schema, ruleName = 'root') {
  const rules = new Map();
  let ruleCounter = 0;

  function nextRuleName(hint) {
    return `${hint || 'r'}-${ruleCounter++}`;
  }

  function addRule(name, definition) {
    rules.set(name, definition);
    return name;
  }

  // Common primitive rules
  addRule('ws', '"" | " " | "\\n" | "\\n" " "');
  addRule('json-string', '"\\"" json-string-chars "\\"" ');
  addRule('json-string-chars', '("" | json-string-char json-string-chars)');
  addRule('json-string-char', '[^"\\\\\\x00-\\x1f] | "\\\\" ["\\\\/bfnrt] | "\\\\u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]');
  addRule('json-integer', '("-"? ([0-9] | [1-9] [0-9]*))');
  addRule('json-number', 'json-integer ("." [0-9]+)? (("e" | "E") ("+" | "-")? [0-9]+)?');
  addRule('json-boolean', '("true" | "false")');
  addRule('json-null', '"null"');

  function processSchema(s, hint) {
    if (!s || typeof s !== 'object') return 'json-string';

    // Enum: literal string alternatives
    if (s.enum && Array.isArray(s.enum)) {
      const name = nextRuleName(hint || 'enum');
      const alts = s.enum.map(v => `"\\"${String(v).replace(/"/g, '\\"')}\\""` ).join(' | ');
      return addRule(name, alts);
    }

    // Const: single literal
    if (s.const !== undefined) {
      const name = nextRuleName(hint || 'const');
      return addRule(name, `"\\"${String(s.const).replace(/"/g, '\\"')}\\""` );
    }

    switch (s.type) {
      case 'string':
        return 'json-string';

      case 'integer':
        return 'json-integer';

      case 'number':
        return 'json-number';

      case 'boolean':
        return 'json-boolean';

      case 'null':
        return 'json-null';

      case 'array': {
        const name = nextRuleName(hint || 'arr');
        const itemRule = s.items ? processSchema(s.items, `${name}-item`) : 'json-string';
        addRule(name, `"[" ws "]" | "[" ws ${itemRule} ("," ws ${itemRule})* ws "]"`);
        return name;
      }

      case 'object': {
        const name = nextRuleName(hint || 'obj');
        const props = s.properties || {};
        const required = new Set(s.required || []);
        const propNames = Object.keys(props);

        if (propNames.length === 0) {
          // Empty object or additionalProperties: accept any JSON object
          addRule(name, '"{" ws "}"');
          return name;
        }

        // Build rules for each property
        const propRules = [];
        for (const propName of propNames) {
          const propSchema = props[propName];
          const valueRule = processSchema(propSchema, `${name}-${propName}`);
          const propRuleName = nextRuleName(`${name}-p-${propName}`);
          addRule(propRuleName, `"\\"${propName}\\"" ws ":" ws ${valueRule}`);
          propRules.push({ name: propName, rule: propRuleName, required: required.has(propName) });
        }

        // Build the object rule with required/optional properties
        const requiredProps = propRules.filter(p => p.required);
        const optionalProps = propRules.filter(p => !p.required);

        if (requiredProps.length === 0 && optionalProps.length === 0) {
          addRule(name, '"{" ws "}"');
          return name;
        }

        // Simple approach: required properties in order, then optional ones
        let parts = [];
        for (const p of requiredProps) {
          parts.push(p.rule);
        }
        for (const p of optionalProps) {
          parts.push(`(${p.rule})?`);
        }

        const body = parts.join(' ("," ws) ');
        addRule(name, `"{" ws ${body} ws "}"`);
        return name;
      }

      default:
        // Unknown or missing type — accept any JSON value
        return 'json-string';
    }
  }

  const rootRule = processSchema(schema, ruleName);
  if (rootRule !== ruleName) {
    addRule(ruleName, rootRule);
  }

  // Build final grammar string
  let grammar = '';
  for (const [name, def] of rules) {
    grammar += `${name} ::= ${def}\n`;
  }
  return grammar;
}

/**
 * Build a GBNF grammar for a single tool's arguments.
 * @param {Object} toolDef - OpenAI-format tool definition
 * @returns {string|null} GBNF grammar string, or null if no parameters
 */
function buildArgumentGrammar(toolDef) {
  const params = toolDef?.function?.parameters;
  if (!params) return null;
  return jsonSchemaToGbnf(params, 'root');
}

/**
 * Build a GBNF grammar that constrains output to a valid tool call.
 * The grammar allows calling any of the provided tools.
 *
 * @param {Array} toolDefs - Array of OpenAI-format tool definitions
 * @returns {string|null} GBNF grammar string
 */
function buildToolCallGrammar(toolDefs) {
  if (!toolDefs || toolDefs.length === 0) return null;

  const toolRules = [];
  let grammar = '';

  // Common rules
  grammar += 'ws ::= "" | " " | "\\n" | "\\n" " "\n';
  grammar += 'json-string ::= "\\"" json-string-chars "\\""\n';
  grammar += 'json-string-chars ::= ("" | json-string-char json-string-chars)\n';
  grammar += 'json-string-char ::= [^"\\\\\\x00-\\x1f] | "\\\\" ["\\\\/bfnrt] | "\\\\u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]\n';
  grammar += 'json-integer ::= ("-"? ([0-9] | [1-9] [0-9]*))\n';
  grammar += 'json-number ::= json-integer ("." [0-9]+)? (("e" | "E") ("+" | "-")? [0-9]+)?\n';
  grammar += 'json-boolean ::= ("true" | "false")\n';
  grammar += 'json-null ::= "null"\n\n';

  for (const def of toolDefs) {
    const name = def.function?.name;
    if (!name) continue;

    const safeName = name.replace(/[^a-zA-Z0-9_-]/g, '_');
    const argsGrammar = buildArgumentGrammar(def);
    if (argsGrammar) {
      // Extract the root rule from the argument grammar and rename it
      const argRuleName = `${safeName}-args`;
      // Use a simplified version — just the JSON schema object
      grammar += `${argRuleName} ::= ${jsonSchemaToGbnf(def.function.parameters, argRuleName).split('\n').find(l => l.startsWith(argRuleName))?.split('::=')[1]?.trim() || '"{" ws "}"'}\n`;
    }

    const argRef = argsGrammar ? `${safeName}-args` : '"{" ws "}"';
    toolRules.push(`${safeName}-call`);
    grammar += `${safeName}-call ::= "{\\"name\\":\\"${name}\\",\\"arguments\\":" ws ${argRef} ws "}"\n`;
  }

  // Root: one tool call
  grammar += `\nroot ::= ${toolRules.join(' | ')}\n`;

  return grammar;
}

module.exports = { jsonSchemaToGbnf, buildArgumentGrammar, buildToolCallGrammar };
