shrimp/src/compiler/utils.ts

import { CompilerError } from '#compiler/compilerError.ts'
import * as terms from '#parser/shrimp.terms'
import type { SyntaxNode, Tree } from '@lezer/common'

export const checkTreeForErrors = (tree: Tree): CompilerError[] => {
  const errors: CompilerError[] = []
  tree.iterate({
    enter: (node) => {
      if (node.type.isError) {
        errors.push(new CompilerError(`Unexpected syntax.`, node.from, node.to))
      }
    },
  })

  return errors
}

export const getAllChildren = (node: SyntaxNode): SyntaxNode[] => {
  const children: SyntaxNode[] = []
  let child = node.firstChild
  while (child) {
    children.push(child)
    child = child.nextSibling
  }

  return children.filter((n) => n.type.id !== terms.Comment)
}

export const getBinaryParts = (node: SyntaxNode) => {
  const children = getAllChildren(node)
  const [left, op, right] = children

  if (!left || !op || !right) {
    throw new CompilerError(`BinOp expected 3 children, got ${children.length}`, node.from, node.to)
  }

  return { left, op, right }
}

export const getAssignmentParts = (node: SyntaxNode) => {
  const children = getAllChildren(node)
  const [left, equals, right] = children

  if (!equals || !right) {
    throw new CompilerError(
      `Assign expected 3 children, got ${children.length}`,
      node.from,
      node.to
    )
  }

  // array destructuring
  if (left && left.type.id === terms.Array) {
    const identifiers = getAllChildren(left).filter((child) => child.type.id === terms.Identifier)
    return { arrayPattern: identifiers, right }
  }

  if (!left || left.type.id !== terms.AssignableIdentifier) {
    throw new CompilerError(
      `Assign left child must be an AssignableIdentifier or Array, got ${
        left ? left.type.name : 'none'
      }`,
      node.from,
      node.to
    )
  }

  return { identifier: left, right }
}

export const getCompoundAssignmentParts = (node: SyntaxNode) => {
  const children = getAllChildren(node)
  const [left, operator, right] = children

  if (!left || left.type.id !== terms.AssignableIdentifier) {
    throw new CompilerError(
      `CompoundAssign left child must be an AssignableIdentifier, got ${
        left ? left.type.name : 'none'
      }`,
      node.from,
      node.to
    )
  } else if (!operator || !right) {
    throw new CompilerError(
      `CompoundAssign expected 3 children, got ${children.length}`,
      node.from,
      node.to
    )
  }

  return { identifier: left, operator, right }
}

export const getFunctionDefParts = (node: SyntaxNode, input: string) => {
  const children = getAllChildren(node)
  const [fnKeyword, paramsNode, colon, ...rest] = children

  if (!fnKeyword || !paramsNode || !colon || !rest) {
    throw new CompilerError(
      `FunctionDef expected at least 4 children, got ${children.length}`,
      node.from,
      node.to
    )
  }

  const paramNames = getAllChildren(paramsNode).map((param) => {
    if (param.type.id !== terms.Identifier && param.type.id !== terms.NamedParam) {
      throw new CompilerError(
        `FunctionDef params must be Identifier or NamedParam, got ${param.type.name}`,
        param.from,
        param.to
      )
    }
    return input.slice(param.from, param.to)
  })

  // Separate body nodes from catch/finally/end
  const bodyNodes: SyntaxNode[] = []
  let catchExpr: SyntaxNode | undefined
  let catchVariable: string | undefined
  let catchBody: SyntaxNode | undefined
  let finallyExpr: SyntaxNode | undefined
  let finallyBody: SyntaxNode | undefined

  for (const child of rest) {
    if (child.type.id === terms.CatchExpr) {
      catchExpr = child
      const catchChildren = getAllChildren(child)
      const [_catchKeyword, identifierNode, _colon, body] = catchChildren
      if (!identifierNode || !body) {
        throw new CompilerError(
          `CatchExpr expected identifier and body, got ${catchChildren.length} children`,
          child.from,
          child.to
        )
      }
      catchVariable = input.slice(identifierNode.from, identifierNode.to)
      catchBody = body
    } else if (child.type.id === terms.FinallyExpr) {
      finallyExpr = child
      const finallyChildren = getAllChildren(child)
      const [_finallyKeyword, _colon, body] = finallyChildren
      if (!body) {
        throw new CompilerError(
          `FinallyExpr expected body, got ${finallyChildren.length} children`,
          child.from,
          child.to
        )
      }
      finallyBody = body
    } else if (child.type.name === 'keyword' && input.slice(child.from, child.to) === 'end') {
      // Skip the end keyword
    } else {
      bodyNodes.push(child)
    }
  }

  return { paramNames, bodyNodes, catchVariable, catchBody, finallyBody }
}

export const getFunctionCallParts = (node: SyntaxNode, input: string) => {
  const [identifierNode, ...args] = getAllChildren(node)

  if (!identifierNode) {
    throw new CompilerError(`FunctionCall expected at least 1 child, got 0`, node.from, node.to)
  }

  const namedArgs = args.filter((arg) => arg.type.id === terms.NamedArg)
  const positionalArgs = args
    .filter((arg) => arg.type.id === terms.PositionalArg)
    .map((arg) => {
      const child = arg.firstChild
      if (!child) throw new CompilerError(`PositionalArg has no child`, arg.from, arg.to)

      return child
    })

  return { identifierNode, namedArgs, positionalArgs }
}

export const getNamedArgParts = (node: SyntaxNode, input: string) => {
  const children = getAllChildren(node)
  const [namedArgPrefix, valueNode] = getAllChildren(node)

  if (!namedArgPrefix || !valueNode) {
    const message = `NamedArg expected 2 children, got ${children.length}`
    throw new CompilerError(message, node.from, node.to)
  }

  const name = input.slice(namedArgPrefix.from, namedArgPrefix.to - 1) // Remove the trailing =
  return { name, valueNode }
}

export const getIfExprParts = (node: SyntaxNode, input: string) => {
  const children = getAllChildren(node)

  const [ifKeyword, conditionNode, _colon, thenBlock, ...rest] = children
  if (!ifKeyword || !conditionNode || !thenBlock) {
    throw new CompilerError(
      `IfExpr expected at least 4 children, got ${children.length}`,
      node.from,
      node.to
    )
  }

  let elseIfBlocks: { conditional: SyntaxNode; thenBlock: SyntaxNode }[] = []
  let elseThenBlock: SyntaxNode | undefined
  rest.forEach((child) => {
    const parts = getAllChildren(child)

    if (child.type.id === terms.ElseExpr) {
      if (parts.length !== 3) {
        const message = `ElseExpr expected 1 child, got ${parts.length}`
        throw new CompilerError(message, child.from, child.to)
      }
      elseThenBlock = parts.at(-1)
    } else if (child.type.id === terms.ElseIfExpr) {
      const [_else, _if, conditional, _colon, thenBlock] = parts
      if (!conditional || !thenBlock) {
        const names = parts.map((p) => p.type.name).join(', ')
        const message = `ElseIfExpr expected conditional and thenBlock, got ${names}`
        throw new CompilerError(message, child.from, child.to)
      }

      elseIfBlocks.push({ conditional, thenBlock })
    }
  })

  return { conditionNode, thenBlock, elseThenBlock, elseIfBlocks }
}

export const getPipeExprParts = (node: SyntaxNode) => {
  const [pipedFunctionCall, operator, ...rest] = getAllChildren(node)
  if (!pipedFunctionCall || !operator || rest.length === 0) {
    const message = `PipeExpr expected at least 3 children, got ${getAllChildren(node).length}`
    throw new CompilerError(message, node.from, node.to)
  }

  const pipeReceivers = rest.filter((child) => child.name !== 'operator')

  return { pipedFunctionCall, pipeReceivers }
}

export const getStringParts = (node: SyntaxNode, input: string) => {
  const children = getAllChildren(node)

  // String nodes always have at least 2 children (the quote tokens)
  // For simple strings like 'hello' with no interpolation, there are no child nodes
  // The text is just between the quotes
  const parts = children.filter((child) => {
    return (
      child.type.id === terms.StringFragment ||
      child.type.id === terms.Interpolation ||
      child.type.id === terms.EscapeSeq
    )
  })

  // Validate each part is the expected type
  parts.forEach((part) => {
    if (
      part.type.id !== terms.StringFragment &&
      part.type.id !== terms.Interpolation &&
      part.type.id !== terms.EscapeSeq
    ) {
      throw new CompilerError(
        `String child must be StringFragment, Interpolation, or EscapeSeq, got ${part.type.name}`,
        part.from,
        part.to
      )
    }
  })

  // hasInterpolation means the string has interpolation ($var) or escape sequences (\n)
  // A simple string like 'hello' has one StringFragment but no interpolation
  const hasInterpolation = parts.some(
    (p) => p.type.id === terms.Interpolation || p.type.id === terms.EscapeSeq
  )
  return { parts, hasInterpolation }
}

export const getDotGetParts = (node: SyntaxNode, input: string) => {
  const children = getAllChildren(node)
  const [object, property] = children

  if (!object || !property) {
    throw new CompilerError(
      `DotGet expected 2 identifier children, got ${children.length}`,
      node.from,
      node.to
    )
  }

  if (object.type.id !== terms.IdentifierBeforeDot) {
    throw new CompilerError(
      `DotGet object must be an IdentifierBeforeDot, got ${object.type.name}`,
      object.from,
      object.to
    )
  }

  if (![terms.Identifier, terms.Number, terms.ParenExpr].includes(property.type.id)) {
    throw new CompilerError(
      `DotGet property must be an Identifier or Number, got ${property.type.name}`,
      property.from,
      property.to
    )
  }

  const objectName = input.slice(object.from, object.to)

  return { objectName, property }
}

export const getTryExprParts = (node: SyntaxNode, input: string) => {
  const children = getAllChildren(node)

  // First child is always 'try' keyword, second is colon, third is Block
  const [tryKeyword, _colon, tryBlock, ...rest] = children

  if (!tryKeyword || !tryBlock) {
    throw new CompilerError(
      `TryExpr expected at least 3 children, got ${children.length}`,
      node.from,
      node.to
    )
  }

  let catchExpr: SyntaxNode | undefined
  let catchVariable: string | undefined
  let catchBody: SyntaxNode | undefined
  let finallyExpr: SyntaxNode | undefined
  let finallyBody: SyntaxNode | undefined

  rest.forEach((child) => {
    if (child.type.id === terms.CatchExpr) {
      catchExpr = child
      const catchChildren = getAllChildren(child)
      const [_catchKeyword, identifierNode, _colon, body] = catchChildren
      if (!identifierNode || !body) {
        throw new CompilerError(
          `CatchExpr expected identifier and body, got ${catchChildren.length} children`,
          child.from,
          child.to
        )
      }
      catchVariable = input.slice(identifierNode.from, identifierNode.to)
      catchBody = body
    } else if (child.type.id === terms.FinallyExpr) {
      finallyExpr = child
      const finallyChildren = getAllChildren(child)
      const [_finallyKeyword, _colon, body] = finallyChildren
      if (!body) {
        throw new CompilerError(
          `FinallyExpr expected body, got ${finallyChildren.length} children`,
          child.from,
          child.to
        )
      }
      finallyBody = body
    }
  })

  return {
    tryBlock,
    catchExpr,
    catchVariable,
    catchBody,
    finallyExpr,
    finallyBody,
  }
}