codetutor/frontend/src/lib/python-helpers.ts

/**
 * Python helper code and detection utilities for the test runner.
 *
 * This module provides Python code as string constants for data structure
 * classes, conversion functions, and utilities to detect problem types.
 */

// =============================================================================
// Data Structure Classes
// =============================================================================

export const TREE_NODE_CLASS = `
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
`.trim();

export const LIST_NODE_CLASS = `
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
`.trim();

// =============================================================================
// Conversion Functions
// =============================================================================

/**
 * Builds a binary tree from a level-order array representation.
 * Example: [1, 2, 3, null, 4] -> TreeNode(1, TreeNode(2, None, TreeNode(4)), TreeNode(3))
 */
export const BUILD_TREE = `
def __build_tree(arr):
    if not arr or arr[0] is None:
        return None
    root = TreeNode(arr[0])
    queue = [root]
    i = 1
    while queue and i < len(arr):
        node = queue.pop(0)
        if i < len(arr) and arr[i] is not None:
            node.left = TreeNode(arr[i])
            queue.append(node.left)
        i += 1
        if i < len(arr) and arr[i] is not None:
            node.right = TreeNode(arr[i])
            queue.append(node.right)
        i += 1
    return root
`.trim();

/**
 * Converts a binary tree to a level-order array representation.
 * Example: TreeNode(1, TreeNode(2), TreeNode(3)) -> [1, 2, 3]
 */
export const TREE_TO_ARRAY = `
def __tree_to_array(root):
    if root is None:
        return []
    result = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            result.append(None)
        else:
            result.append(node.val)
            queue.append(node.left)
            queue.append(node.right)
    # Trim trailing nulls
    while result and result[-1] is None:
        result.pop()
    return result
`.trim();

/**
 * Builds a linked list from an array.
 * Example: [1, 2, 3] -> ListNode(1, ListNode(2, ListNode(3)))
 */
export const BUILD_LIST = `
def __build_list(arr):
    if not arr:
        return None
    head = ListNode(arr[0])
    current = head
    for val in arr[1:]:
        current.next = ListNode(val)
        current = current.next
    return head
`.trim();

/**
 * Converts a linked list to an array.
 * Example: ListNode(1, ListNode(2, ListNode(3))) -> [1, 2, 3]
 */
export const LIST_TO_ARRAY = `
def __list_to_array(head):
    result = []
    current = head
    while current:
        result.append(current.val)
        current = current.next
    return result
`.trim();

// =============================================================================
// Problem Type Detection
// =============================================================================

/** Parameter names that indicate a tree-based problem */
export const TREE_PARAMS = ["root", "tree", "p", "q"];

/** Parameter names that indicate a linked-list problem */
export const LIST_PARAMS = ["head", "l1", "l2", "list1", "list2", "node"];

export type ProblemType = "simple" | "tree" | "linkedlist" | "class-based";

/**
 * Detects the problem type based on the function signature and input.
 *
 * @param signature - The function signature (e.g., "def two_sum(nums, target):")
 * @param input - The test case input object
 * @returns The detected problem type
 */
export function detectProblemType(
  signature: string,
  input: Record<string, unknown>
): ProblemType {
  // Class-based: check if input has "operations" key
  if ("operations" in input) {
    return "class-based";
  }

  // Check signature for type hints
  const hasTreeNode = signature.includes("TreeNode");
  const hasListNode = signature.includes("ListNode");

  if (hasTreeNode) {
    return "tree";
  }

  if (hasListNode) {
    return "linkedlist";
  }

  // Check parameter names in signature and input keys
  const inputKeys = Object.keys(input);

  const hasTreeParam =
    TREE_PARAMS.some((param) => signature.includes(param)) ||
    inputKeys.some((key) => TREE_PARAMS.includes(key));

  const hasListParam =
    LIST_PARAMS.some((param) => signature.includes(param)) ||
    inputKeys.some((key) => LIST_PARAMS.includes(key));

  if (hasTreeParam) {
    return "tree";
  }

  if (hasListParam) {
    return "linkedlist";
  }

  return "simple";
}

/**
 * Gets the Python helper code needed for a given problem type.
 */
export function getPythonHelpers(problemType: ProblemType): string {
  switch (problemType) {
    case "tree":
      return [TREE_NODE_CLASS, BUILD_TREE, TREE_TO_ARRAY].join("\n\n");
    case "linkedlist":
      return [LIST_NODE_CLASS, BUILD_LIST, LIST_TO_ARRAY].join("\n\n");
    case "class-based":
      return ""; // Class-based problems define their own class
    case "simple":
    default:
      return "";
  }
}

/**
 * Gets the list of parameter names that need tree conversion.
 */
export function getTreeParams(input: Record<string, unknown>): string[] {
  return Object.keys(input).filter((key) => TREE_PARAMS.includes(key));
}

/**
 * Gets the list of parameter names that need linked list conversion.
 */
export function getListParams(input: Record<string, unknown>): string[] {
  return Object.keys(input).filter((key) => LIST_PARAMS.includes(key));
}