import type { Tree } from '@vscode/tree-sitter-wasm'
import type { IndexedFile, IndexEntry, Node, RuleLocation, RulePosition, RuleWorkspace, WorkspaceIndex } from './contract'
import * as path from 'node:path'

/**
 * 工作区索引（WorkspaceIndex：库规则的，和宿主自己认中文名的）的宿主部分：找文件、预筛、解析、缓存，跟着文件改动只重看改了的文件。
 * 规则只说要什么；这里不碰 VS Code，工作区的能力由 IndexHost 给（VS Code 里是 workspaceFiles.ts，测试里是内存里的文件）。
 */

/** 工作区里的一份文件：uri，它所在的工作区文件夹（uri）和相对这个文件夹的路径（/ 分隔）。 */
export interface WorkspaceFile { uri: string, folder: string, path: string }

/** 文件改动：改了、新建了、删了（删的可能是目录，里面的文件一起算删了）。 */
export interface FileEvent { kind: 'change' | 'create' | 'delete', file: WorkspaceFile }

export interface IndexHost {
  /** include 在各工作区文件夹里找到的文件，排除的不算。 */
  list: (include: string) => Promise<WorkspaceFile[]>
  read: (file: WorkspaceFile) => Promise<string>
  /** 盯着文件改动：include 以外和被排除的文件改了、新建了不报，删除都报。 */
  watch: (include: string, listener: (event: FileEvent) => void) => { dispose: () => void }
  /** 用 languageId 的语言插件解析；返回的语法树用完要 delete。 */
  parse: (languageId: string, source: string) => Promise<Tree>
  report: (message: string) => void
}

/** 一次读几份文件：建索引时不一口气把整个工作区的文件都读进内存。 */
const READ_BATCH = 32

function keyOf(folder: string, relative: string): string {
  return `${folder}\n${path.posix.normalize(relative)}`
}

/** 每一行开头在源码里的下标。 */
function lineStarts(source: string): number[] {
  const starts = [0]
  for (let i = source.indexOf('\n'); i >= 0; i = source.indexOf('\n', i + 1))
    starts.push(i + 1)
  return starts
}

function positionAt(starts: number[], index: number): RulePosition {
  let low = 0
  let high = starts.length - 1
  while (low < high) {
    const middle = (low + high + 1) >> 1
    if (starts[middle] <= index)
      low = middle
    else
      high = middle - 1
  }
  return { line: low, character: index - starts[low] }
}

/** 认出的一项在哪：定义处名字的位置，连同认的时候给的这一项（宿主自己的索引在项里多带东西，例如中文名）。 */
export interface IndexedEntry<E extends IndexEntry = IndexEntry> { location: RuleLocation, entry: E }

/** 认出的项是 E 的索引。 */
export interface WorkspaceIndexOf<E extends IndexEntry> extends WorkspaceIndex {
  entries: (root: Node, file: IndexedFile) => E[]
}

/** 一份看过的文件：认出的各项，和建索引时问过在不在的文件（它们新建、删掉时这份要重看）。 */
interface Indexed<E extends IndexEntry> {
  file: WorkspaceFile
  entries: IndexedEntry<E>[]
  asked: Set<string>
}

function compareLocations(a: RuleLocation, b: RuleLocation): number {
  return a.uri.localeCompare(b.uri) || a.start.line - b.start.line || a.start.character - b.start.character
}

/** 一个 WorkspaceIndex 建出来的索引：建好之前的查询、文件改动都排在建索引后面，一件件来。 */
class BuiltIndex<E extends IndexEntry> {
  /** include 找到的全部文件（按 keyOf），IndexedFile.exists 查的就是它 */
  private readonly files = new Map<string, WorkspaceFile>()
  /** 预筛过了、解析过的文件 */
  private readonly indexed = new Map<string, Indexed<E>>()
  private readonly watcher: { dispose: () => void }
  private queue: Promise<void>

  constructor(private readonly index: WorkspaceIndexOf<E>, private readonly host: IndexHost) {
    // 先盯着再找文件：找文件期间的改动排在建好之后补上
    this.watcher = host.watch(index.include, event => this.enqueue(() => this.apply(event)))
    this.queue = this.guard(() => this.build())
  }

  async lookup(key: string): Promise<IndexedEntry<E>[]> {
    await this.queue
    return [...this.indexed.values()]
      .flatMap(indexed => indexed.entries.filter(({ entry }) => entry.key === key))
      .sort((a, b) => compareLocations(a.location, b.location))
  }

  dispose(): void {
    this.watcher.dispose()
  }

  private enqueue(task: () => Promise<void>): void {
    this.queue = this.queue.then(() => this.guard(task))
  }

  private async guard(task: () => Promise<void>): Promise<void> {
    try {
      await task()
    }
    catch (err) {
      this.host.report(`索引 ${this.index.id} 更新失败：${err instanceof Error ? err.message : String(err)}`)
    }
  }

  private async build(): Promise<void> {
    const files = await this.host.list(this.index.include)
    for (const file of files)
      this.files.set(keyOf(file.folder, file.path), file)
    for (let i = 0; i < files.length; i += READ_BATCH)
      await Promise.all(files.slice(i, i + READ_BATCH).map(file => this.visit(file)))
  }

  /** 重看一份文件：读、预筛、解析、交给规则认。 */
  private async visit(file: WorkspaceFile): Promise<void> {
    const key = keyOf(file.folder, file.path)
    this.indexed.delete(key)
    let source: string
    let tree: Tree
    try {
      source = await this.host.read(file)
      if (!this.index.mentions(source))
        return
      tree = await this.host.parse(this.index.languageId, source)
    }
    catch (err) {
      this.host.report(`索引 ${this.index.id} 读不了 ${file.uri}：${err instanceof Error ? err.message : String(err)}`)
      return
    }
    try {
      const asked = new Set<string>()
      const exists = (relative: string): boolean => {
        const other = keyOf(file.folder, relative)
        asked.add(other)
        return this.files.has(other)
      }
      const starts = lineStarts(source)
      const entries = this.index.entries(tree.rootNode, { path: file.path, source, exists }).map(entry => ({
        entry,
        location: { uri: file.uri, start: positionAt(starts, entry.name.startIndex), end: positionAt(starts, entry.name.endIndex) },
      }))
      this.indexed.set(key, { file, entries, asked })
    }
    catch (err) {
      this.host.report(`索引 ${this.index.id} 看 ${file.uri} 时出错：${err instanceof Error ? err.message : String(err)}`)
    }
    finally {
      tree.delete()
    }
  }

  private async apply({ kind, file }: FileEvent): Promise<void> {
    const key = keyOf(file.folder, file.path)
    if (kind === 'change') {
      if (this.files.has(key))
        await this.visit(file)
      return
    }
    const touched = kind === 'create' ? this.add(file) : this.remove(file)
    // 问过这些文件在不在的（Python 包路径看 __init__.py）跟着重看
    const dependents = [...this.indexed.values()]
      .filter(indexed => touched.some(other => indexed.asked.has(other)))
      .map(indexed => indexed.file)
      .filter(dependent => keyOf(dependent.folder, dependent.path) !== key)
    await Promise.all([...(kind === 'create' ? [file] : []), ...dependents].map(target => this.visit(target)))
  }

  private add(file: WorkspaceFile): string[] {
    const key = keyOf(file.folder, file.path)
    this.files.set(key, file)
    return [key]
  }

  /** 删掉文件，或目录连同里面的文件。 */
  private remove(file: WorkspaceFile): string[] {
    const prefix = `${path.posix.normalize(file.path)}/`
    const gone = [...this.files.entries()]
      .filter(([, known]) => known.folder === file.folder && (known.path === file.path || known.path.startsWith(prefix)))
      .map(([key]) => key)
    for (const key of gone) {
      this.files.delete(key)
      this.indexed.delete(key)
    }
    return gone
  }
}

/** 各库规则的工作区索引：第一次查某个索引时才建，同一个 id 只建一份。 */
export class WorkspaceIndexes implements RuleWorkspace {
  private readonly built = new Map<string, BuiltIndex<IndexEntry>>()

  constructor(private readonly host: IndexHost) {}

  async lookup(index: WorkspaceIndex, key: string): Promise<RuleLocation[]> {
    return (await this.find(index, key)).map(({ location }) => location)
  }

  /** 同 lookup，连同认出的那一项（第一次查时建，同一个 id 只建一份）。 */
  find<E extends IndexEntry>(index: WorkspaceIndexOf<E>, key: string): Promise<IndexedEntry<E>[]> {
    let built = this.built.get(index.id) as BuiltIndex<E> | undefined
    if (!built) {
      built = new BuiltIndex(index, this.host)
      this.built.set(index.id, built)
    }
    return built.lookup(key)
  }

  dispose(): void {
    for (const built of this.built.values())
      built.dispose()
    this.built.clear()
  }
}