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> {
  const { className, methodName } = parseQualifiedName(functionName)

  // Try tree-sitter parsing first
  const treeResult = await findWithTreeSitter(code, className, methodName)
  if (treeResult) {
    return treeResult
  }

  // Fall back to regex if tree-sitter fails
  return findFunctionWithRegex(code, methodName, className)
}

async function findWithTreeSitter(
  code: string,
  className: string | undefined,
  methodName: string
): Promise<FunctionLocation | null> {
  try {
    const parser = await getParser()
    if (!parser) return null

    const tree = parser.parse(code)
    if (!tree) return null

    const result = className
      ? findMethodInClass(tree.rootNode, className, methodName)
      : findFunctionNode(tree.rootNode, methodName)

    if (!result) return null

    return {
      startLine: result.startPosition.row + 1,
      endLine: result.endPosition.row + 1,
    }
  } catch (error) {
    console.warn("Tree-sitter parse failed:", error)
    return null
  }
}

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

  return className
    ? findMethodInClassWithRegex(lines, className, functionName)
    : findStandaloneFunctionWithRegex(lines, functionName)
}

function findMethodInClassWithRegex(
  lines: string[],
  className: string,
  methodName: string
): FunctionLocation | null {
  const classPattern = new RegExp(
    `^(\\s*)class\\s+${escapeRegex(className)}\\s*[:(]`
  )
  const methodPattern = new RegExp(
    `^(\\s*)(async\\s+)?def\\s+${escapeRegex(methodName)}\\s*\\(`
  )

  let classIndent = -1
  let inClass = false

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

    // Skip empty lines and comments
    if (!trimmed || trimmed.startsWith("#")) continue

    const currentIndent = getIndentLevel(line)

    // Check for class definition
    if (!inClass) {
      const classMatch = line.match(classPattern)
      if (classMatch) {
        inClass = true
        classIndent = classMatch[1].length
      }
      continue
    }

    // Check if we've left the class
    if (currentIndent <= classIndent) {
      inClass = false
      continue
    }

    // Look for the method within the class
    const methodMatch = line.match(methodPattern)
    if (methodMatch) {
      const methodIndent = methodMatch[1].length
      return findBlockEnd(lines, i + 1, methodIndent)
    }
  }

  return null
}

function findStandaloneFunctionWithRegex(
  lines: string[],
  functionName: string
): FunctionLocation | null {
  const functionPattern = new RegExp(
    `^(\\s*)(async\\s+)?def\\s+${escapeRegex(functionName)}\\s*\\(`
  )

  for (let i = 0; i < lines.length; i++) {
    const match = lines[i].match(functionPattern)
    if (match) {
      const indent = match[1].length
      return findBlockEnd(lines, i + 1, indent)
    }
  }

  return null
}

function findBlockEnd(
  lines: string[],
  startLine: number,
  startIndent: number
): FunctionLocation {
  for (let i = startLine; i < lines.length; i++) {
    const trimmed = lines[i].trim()
    if (!trimmed || trimmed.startsWith("#")) continue

    const currentIndent = getIndentLevel(lines[i])
    if (currentIndent <= startIndent) {
      return { startLine, endLine: i }
    }
  }

  return { startLine, endLine: lines.length }
}

function getIndentLevel(line: string): number {
  return line.length - line.trimStart().length
}

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

function parseQualifiedName(functionName: string): {
  className: string | undefined
  methodName: string
} {
  const dotIndex = functionName.lastIndexOf(".")
  if (dotIndex > 0) {
    return {
      className: functionName.substring(0, dotIndex),
      methodName: functionName.substring(dotIndex + 1),
    }
  }
  return { className: undefined, methodName: functionName }
}

function findMethodInClass(
  rootNode: Node,
  className: string,
  methodName: string
): Node | null {
  // Find all class definitions matching the class name
  function searchClasses(node: Node): Node | null {
    if (node.type === "class_definition") {
      const nameNode = node.childForFieldName("name")
      if (nameNode && nameNode.text === className) {
        // Search for the method within this class body
        const classBody = node.childForFieldName("body")
        if (classBody) {
          for (const child of classBody.children) {
            if (
              child.type === "function_definition" ||
              child.type === "async_function_definition"
            ) {
              const fnName = child.childForFieldName("name")
              if (fnName && fnName.text === methodName) {
                return child
              }
            }
          }
        }
      }
    }

    for (const child of node.children) {
      const result = searchClasses(child)
      if (result) return result
    }
    return null
  }

  return searchClasses(rootNode)
}

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
}