import type { NameBinding, Node } from '../contract'
import type { BindKind, Scope } from './bindings'
import { BindingCollector } from './bindings'
import { named } from './shared'

/** 推导式、生成器表达式：自成一层作用域，第一个 for 的可迭代对象在外层求值。 */
const COMPREHENSIONS = new Set(['list_comprehension', 'set_comprehension', 'dictionary_comprehension', 'generator_expression'])

/** 赋值目标里一层层往里拆的写法。 */
const TARGET_LISTS = new Set(['pattern_list', 'tuple_pattern', 'list_pattern', 'tuple', 'list', 'expression_list', 'parenthesized_expression', 'list_splat_pattern', 'list_splat', 'as_pattern_target'])

/**
 * Python 的作用域规则：模块、函数（def、lambda）、类、推导式各是一层。函数里赋过值、当参数的名字是函数自己的，
 * 用到的名字从里往外找（往外找时跳过类主体），global / nonlocal 声明的照声明找。参数的注解和默认值、装饰器、父类在外层求值。
 * 模式匹配捕获的、类型参数（PEP 695）、`type X = …` 拿不准；模块里有 `from m import *` 时模块级的名字都拿不准。
 */
class PythonScopes {
  readonly collector = new BindingCollector()
  readonly module: Scope
  /** 各层用 global / nonlocal 声明的名字 */
  private readonly declared = new Map<Scope, Map<string, 'global' | 'nonlocal'>>()
  private star = false

  constructor() {
    this.module = this.collector.scope('module', undefined)
  }

  visit(node: Node, scope: Scope): void {
    switch (node.type) {
      case 'identifier':
        this.collector.see(scope, node)
        return
      case 'function_definition':
        this.function(node, scope)
        return
      case 'lambda': {
        const inner = this.collector.scope('function', scope)
        this.parameters(node.childForFieldName('parameters'), scope, inner)
        this.visitField(node, 'body', inner)
        return
      }
      case 'class_definition':
        this.class(node, scope)
        return
      case 'assignment':
      case 'augmented_assignment':
      case 'for_statement':
        this.withTarget(node, 'left', scope)
        return
      case 'as_pattern':
        this.withTarget(node, 'alias', scope)
        return
      case 'named_expression': {
        const name = node.childForFieldName('name')
        if (name)
          this.declare(name, this.functionScope(scope), 'local', scope)
        this.visitField(node, 'value', scope)
        return
      }
      case 'delete_statement':
        for (const target of named(node))
          this.target(target, scope)
        return
      case 'import_statement':
      case 'import_from_statement':
        this.import(node, scope)
        return
      case 'future_import_statement':
        return
      case 'global_statement':
      case 'nonlocal_statement':
        for (const name of named(node)) {
          this.declarations(scope).set(name.text, node.type === 'global_statement' ? 'global' : 'nonlocal')
          this.collector.see(scope, name)
        }
        return
      case 'keyword_argument':
        this.visitField(node, 'value', scope)
        return
      case 'attribute':
        this.visitField(node, 'object', scope)
        return
      case 'match_statement':
        this.match(node, scope)
        return
      case 'type_alias_statement':
        this.collector.unsure(scope, node)
        return
    }
    if (COMPREHENSIONS.has(node.type)) {
      this.comprehension(node, scope)
      return
    }
    for (const child of named(node))
      this.visit(child, scope)
  }

  bindings(): NameBinding[] {
    return this.collector.bindings((scope, name) => this.resolve(scope, name))
  }

  private function(node: Node, scope: Scope): void {
    const name = node.childForFieldName('name')
    if (name)
      this.declare(name, scope, 'local', scope)
    const inner = this.collector.scope('function', scope)
    const types = node.childForFieldName('type_parameters')
    if (types) {
      this.collector.unsure(scope, types)
      this.collector.unsure(inner, types)
    }
    this.parameters(node.childForFieldName('parameters'), scope, inner)
    this.visitField(node, 'return_type', scope)
    this.visitField(node, 'body', inner)
  }

  private class(node: Node, scope: Scope): void {
    const name = node.childForFieldName('name')
    if (name)
      this.declare(name, scope, 'local', scope)
    const inner = this.collector.scope('class', scope)
    const types = node.childForFieldName('type_parameters')
    if (types) {
      this.collector.unsure(scope, types)
      this.collector.unsure(inner, types)
    }
    this.visitField(node, 'superclasses', scope)
    this.visitField(node, 'body', inner)
  }

  /** 参数名是函数这一层的；注解、默认值在外层求值。 */
  private parameters(list: Node | null, outer: Scope, inner: Scope): void {
    for (const parameter of named(list)) {
      switch (parameter.type) {
        case 'identifier':
        case 'list_splat_pattern':
        case 'dictionary_splat_pattern':
          this.parameter(parameter, inner)
          break
        case 'typed_parameter':
          this.parameter(named(parameter)[0], inner)
          this.visitField(parameter, 'type', outer)
          break
        case 'default_parameter':
        case 'typed_default_parameter':
          this.parameter(parameter.childForFieldName('name'), inner)
          this.visitField(parameter, 'type', outer)
          this.visitField(parameter, 'value', outer)
          break
        case 'keyword_separator':
        case 'positional_separator':
          break
        default:
          this.collector.unsure(inner, parameter)
      }
    }
  }

  private parameter(node: Node | null | undefined, scope: Scope): void {
    const name = node?.type === 'list_splat_pattern' || node?.type === 'dictionary_splat_pattern' ? named(node)[0] : node
    if (name?.type === 'identifier')
      this.declare(name, scope, 'local', scope)
    else if (node)
      this.collector.unsure(scope, node)
  }

  /** 推导式自成一层：变量是它的，第一个 for 的可迭代对象在外层求值，后面的在里面。 */
  private comprehension(node: Node, scope: Scope): void {
    const inner = this.collector.scope('comprehension', scope)
    let first = true
    for (const child of named(node)) {
      if (child.type !== 'for_in_clause') {
        this.visit(child, inner)
        continue
      }
      for (let i = 0; i < child.childCount; i++) {
        const part = child.child(i)
        if (!part?.isNamed)
          continue
        const field = child.fieldNameForChild(i)
        if (field === 'left')
          this.target(part, inner)
        else
          this.visit(part, field === 'right' && first ? scope : inner)
      }
      first = false
    }
  }

  /** 模式匹配：被匹配的照常算；case 里的模式（捕获的名字、值）拿不准，守卫和分支照常算。 */
  private match(node: Node, scope: Scope): void {
    for (const child of named(node)) {
      if (child.type !== 'block') {
        this.visit(child, scope)
        continue
      }
      for (const clause of named(child)) {
        for (const part of named(clause)) {
          if (part.type === 'case_pattern')
            this.collector.unsure(scope, part)
          else
            this.visit(part, scope)
        }
      }
    }
  }

  /** 导入的名字：`import a.b` 绑定 a，`as` 后面的名字、`from m import x` 的 x；模块路径、`as` 前面的原名不算。 */
  private import(node: Node, scope: Scope): void {
    for (let i = 0; i < node.childCount; i++) {
      const child = node.child(i)
      if (!child)
        continue
      if (child.type === 'wildcard_import')
        this.star = true
      if (node.fieldNameForChild(i) !== 'name')
        continue
      const name = child.type === 'aliased_import' ? child.childForFieldName('alias') : named(child)[0]
      if (name?.type === 'identifier')
        this.declare(name, scope, 'import', scope)
    }
  }

  /** field 那个子节点是赋值目标，别的照常算。 */
  private withTarget(node: Node, field: string, scope: Scope): void {
    for (let i = 0; i < node.childCount; i++) {
      const child = node.child(i)
      if (!child?.isNamed)
        continue
      if (node.fieldNameForChild(i) === field)
        this.target(child, scope)
      else
        this.visit(child, scope)
    }
  }

  /** 赋值目标：名字在这一层绑定；属性、下标是用到对象；认不出的写法拿不准。 */
  private target(node: Node, scope: Scope): void {
    if (node.type === 'identifier')
      this.declare(node, scope, 'local', scope)
    else if (TARGET_LISTS.has(node.type))
      named(node).forEach(child => this.target(child, scope))
    else if (node.type === 'attribute' || node.type === 'subscript')
      this.visit(node, scope)
    else
      this.collector.unsure(scope, node)
  }

  /** name 在 owner 这一层绑定，出现在 scope 里。 */
  private declare(name: Node, owner: Scope, kind: BindKind, scope: Scope): void {
    owner.bind(name.text, kind)
    this.collector.see(scope, name)
  }

  private visitField(node: Node, field: string, scope: Scope): void {
    for (const child of node.childrenForFieldName(field)) {
      if (child?.isNamed)
        this.visit(child, scope)
    }
  }

  private declarations(scope: Scope): Map<string, 'global' | 'nonlocal'> {
    let found = this.declared.get(scope)
    if (!found) {
      found = new Map()
      this.declared.set(scope, found)
    }
    return found
  }

  /** 海象运算符绑定的那一层：往外第一层不是推导式的。 */
  private functionScope(scope: Scope): Scope {
    let found = scope
    while (found.kind === 'comprehension' && found.parent)
      found = found.parent
    return found
  }

  /** scope 里用到的 name 落在哪个绑定上；拿不准的没有。 */
  private resolve(scope: Scope, name: string): string | undefined {
    for (let current: Scope | undefined = scope; current; current = current.parent) {
      if (current !== scope && current.kind === 'class')
        continue
      if (current.unsure.has(name))
        return undefined
      const declared = this.declared.get(current)?.get(name)
      if (declared === 'global')
        return this.global(name)
      if (declared === 'nonlocal')
        return this.nonlocal(current, name)
      if (current.binds.has(name))
        return current === this.module ? this.global(name) : `${current.id}:${name}`
    }
    return this.global(name)
  }

  /** scope 里 nonlocal 声明的 name：往外找最近一层绑定它的函数，找不到拿不准。 */
  private nonlocal(scope: Scope, name: string): string | undefined {
    for (let current = scope.parent; current && current !== this.module; current = current.parent) {
      if (current.kind === 'class')
        continue
      if (current.unsure.has(name))
        return undefined
      const declared = this.declared.get(current)?.get(name)
      if (declared === 'nonlocal')
        return this.nonlocal(current, name)
      if (declared === undefined && current.binds.has(name))
        return `${current.id}:${name}`
      if (declared !== undefined)
        return undefined
    }
    return undefined
  }

  /** 模块级的名字（含哪儿都没绑定的内置名字）；有 `from m import *` 时拿不准。 */
  private global(name: string): string | undefined {
    return this.star ? undefined : `:${name}`
  }
}

/** 一份 Python 源码里各个名字的绑定（见 NameBinding）。 */
export function pythonBindings(root: Node): NameBinding[] {
  const scopes = new PythonScopes()
  for (const child of named(root))
    scopes.visit(child, scopes.module)
  return scopes.bindings()
}