import { SyntaxNode } from '@lezer/common'
import { TextDocument } from 'vscode-languageserver-textdocument'
import * as Terms from '../../../src/parser/shrimp.terms'
import { globals } from '../../../src/prelude'

/**
 * Tracks variables in scope at a given position in the parse tree.
 * Used to distinguish identifiers (in scope) from words (not in scope).
 */
export class EditorScopeAnalyzer {
  private document: TextDocument
  private scopeCache = new Map<number, Set<string>>()

  constructor(document: TextDocument) {
    this.document = document
    const preludeKeys = Object.keys(globals)
    this.scopeCache.set(0, new Set(preludeKeys))
  }

  /**
   * Check if a name is in scope at the given node's position.
   */
  isInScope(name: string, node: SyntaxNode): boolean {
    const scope = this.getScopeAt(node)
    return scope.has(name)
  }

  /**
   * Get all variables in scope at the given node's position.
   */
  private getScopeAt(node: SyntaxNode): Set<string> {
    const position = node.from

    // Check cache first
    if (this.scopeCache.has(position)) {
      return this.scopeCache.get(position)!
    }

    const scope = new Set<string>()

    // Find all containing function definitions
    const containingFunctions = this.findContainingFunctions(node)

    // Collect scope from each containing function (inner to outer)
    for (const fnNode of containingFunctions) {
      this.collectParams(fnNode, scope)
      this.collectAssignments(fnNode, position, scope)
    }

    // Collect top-level assignments
    const root = this.getRoot(node)
    this.collectAssignments(root, position, scope)

    this.scopeCache.set(position, scope)
    return scope
  }

  /**
   * Find all function definitions that contain the given node.
   */
  private findContainingFunctions(node: SyntaxNode): SyntaxNode[] {
    const functions: SyntaxNode[] = []
    let current = node.parent

    while (current) {
      if (current.type.id === Terms.FunctionDef) {
        functions.unshift(current) // Add to beginning for outer-to-inner order
      }
      current = current.parent
    }

    return functions
  }

  /**
   * Get the root node of the tree.
   */
  private getRoot(node: SyntaxNode): SyntaxNode {
    let current = node
    while (current.parent) {
      current = current.parent
    }
    return current
  }

  /**
   * Collect parameter names from a function definition.
   */
  private collectParams(fnNode: SyntaxNode, scope: Set<string>) {
    let child = fnNode.firstChild
    while (child) {
      if (child.type.id === Terms.Params) {
        let param = child.firstChild
        while (param) {
          if (param.type.id === Terms.Identifier) {
            const text = this.document.getText({
              start: this.document.positionAt(param.from),
              end: this.document.positionAt(param.to),
            })
            scope.add(text)
          }
          param = param.nextSibling
        }
        break
      }
      child = child.nextSibling
    }
  }

  /**
   * Collect assignment names from a scope node that occur before the given position.
   */
  private collectAssignments(scopeNode: SyntaxNode, beforePosition: number, scope: Set<string>) {
    const cursor = scopeNode.cursor()

    cursor.iterate((node) => {
      // Stop if we've passed the position we're checking
      if (node.from >= beforePosition) return false

      if (node.type.id === Terms.Assign) {
        const assignNode = node.node
        const child = assignNode.firstChild
        if (child?.type.id === Terms.AssignableIdentifier) {
          const text = this.document.getText({
            start: this.document.positionAt(child.from),
            end: this.document.positionAt(child.to),
          })
          scope.add(text)
        }
      }

      // Don't descend into nested functions unless it's the current scope
      if (node.type.id === Terms.FunctionDef && node.node !== scopeNode) {
        return false
      }
    })
  }
}