import type { Node, Parser as ParserType } from "web-tree-sitter"

export interface FunctionLocation {
  startLine: number
  endLine: number
}

let parserPromise: Promise<ParserType | null> | null = null

async function getParser(): Promise<ParserType | null> {
  if (typeof window === "undefined") {
    return null
  }

  if (!parserPromise) {
    parserPromise = (async () => {
      try {
        const { Parser, Language } = await import("web-tree-sitter")
        await Parser.init({
          locateFile: (scriptName: string) => `/${scriptName}`,
        })
        const parser = new Parser()
        const Python = await Language.load("/tree-sitter-python.wasm")
        parser.setLanguage(Python)
        return parser
      } catch (error) {
        console.error("Tree-sitter initialization failed:", error)
        parserPromise = null
        throw error
      }
    })()
  }
  return parserPromise
}

export async function findFunctionInCode(
  code: string,
  functionName: string
): Promise<FunctionLocation | null> {
  try {
    const parser = await getParser()
    if (parser) {
      const tree = parser.parse(code)
      if (tree) {
        const result = findFunctionNode(tree.rootNode, functionName)
        if (result) {
          return {
            startLine: result.startPosition.row + 1,
            endLine: result.endPosition.row + 1,
          }
        }
      }
    }
  } catch (error) {
    console.warn("Tree-sitter parse failed, trying regex fallback:", error)
  }

  return findFunctionWithRegex(code, functionName)
}

function findFunctionWithRegex(
  code: string,
  functionName: string
): FunctionLocation | null {
  const lines = code.split("\n")

  const defPattern = new RegExp(
    `^(\\s*)(async\\s+)?def\\s+${escapeRegex(functionName)}\\s*\\(`
  )

  let startLine = -1
  let startIndent = -1

  for (let i = 0; i < lines.length; i++) {
    const line = lines[i]

    if (startLine === -1) {
      const match = line.match(defPattern)
      if (match) {
        startLine = i + 1
        startIndent = match[1].length
      }
    } else {
      const trimmed = line.trim()
      if (trimmed === "" || trimmed.startsWith("#")) {
        continue
      }

      const currentIndent = line.length - line.trimStart().length
      if (currentIndent <= startIndent) {
        return { startLine, endLine: i }
      }
    }
  }

  if (startLine !== -1) {
    return { startLine, endLine: lines.length }
  }

  return null
}

function escapeRegex(str: string): string {
  return str.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")
}

function findFunctionNode(node: Node, functionName: string): Node | null {
  if (
    node.type === "function_definition" ||
    node.type === "async_function_definition"
  ) {
    const nameNode = node.childForFieldName("name")
    if (nameNode && nameNode.text === functionName) {
      return node
    }
  }

  if (node.type === "class_definition") {
    const classBody = node.childForFieldName("body")
    if (classBody) {
      for (const child of classBody.children) {
        const result = findFunctionNode(child, functionName)
        if (result) return result
      }
    }
  }

  for (const child of node.children) {
    const result = findFunctionNode(child, functionName)
    if (result) return result
  }

  return null
}