/*
 * darwin-apk © 2024, 2025 Gay Pizza Specifications
 * SPDX-License-Identifier: Apache-2.0
 */

import Foundation

public struct ApkPackageGraph: ~Copyable {
  private var _nodes = [ApkPackageGraphNode]()

  public var nodes: [ApkPackageGraphNode] { self._nodes }
  public var shallowIsolates: [ApkPackageGraphNode] { self._nodes.filter(\.isShallow) }
  public var deepIsolates: [ApkPackageGraphNode] { self._nodes.filter(\.isDeep) }

  public init() {}

  public mutating func buildGraphNode(index pkgIndex: ApkIndex, providers: ApkIndexProviderCache) {
    for (packageID, package) in pkgIndex.packages.enumerated() {
      let children: [ApkPackageGraphNode.ChildRef] = package.dependencies.compactMap { dependency in
        guard let providerID = providers.resolve(index: pkgIndex, requirement: dependency.requirement) else {
          return nil
        }
        return .init(constraint: .dependency, packageID: providerID, versionSpec: dependency.requirement.versionSpec)
      } /* + package.installIf.compactMap { installIf in
        guard let prvID = providers.resolve(index: pkgIndex, requirement: installIf.requirement) else {
          return nil
        }
        return .init(constraint: .installIf, packageID: prvID, versionSpec: installIf.requirement.versionSpec)
      } */
      self._nodes.append(.init(
        id: packageID,
        children: children
      ))
    }

    var reverseDependencies = [ApkIndex.Index: [ApkIndex.Index]]()
    for (index, node) in self._nodes.enumerated() {
      for child in node.children {
        reverseDependencies[child.packageID, default: []].append(index)
      }
    }

    for (ref, parents) in reverseDependencies {
      self._nodes[ref].parentIDs = parents
    }
  }
}

extension ApkPackageGraph {
  public func orderSort(breakCycles: Bool = true) throws(SortError) -> [ApkPackageGraphNode] {
    var stack = [ApkPackageGraphNode]()
    var resolving = Set<ApkIndex.Index>()
    var visited = Set<ApkIndex.Index>()
    var ignoring = Set<ApkIndex.Index>()
    for node in self.shallowIsolates {
      try orderSort(node, &stack, &resolving, &visited, &ignoring, breakCycles)
    }
    return stack
  }

  internal func orderSort(
    _ node: ApkPackageGraphNode,
    _ stack: inout [ApkPackageGraphNode],
    _ resolving: inout Set<ApkIndex.Index>,
    _ visited: inout Set<ApkIndex.Index>,
    _ ignoring: inout Set<ApkIndex.Index>,
    _ breakCycles: Bool
  ) throws(SortError) {
    for dep in node.children {
      let depID = dep.packageID
      guard !ignoring.contains(depID) else {
        continue
      }

      guard !resolving.contains(depID) else {
        throw .cyclicDependency(cycles: "\(node) -> \(dep)")
      }

      if !visited.contains(depID) {
        resolving.insert(depID)
        let depNode = self._nodes[depID]
        do {
          try orderSort(depNode, &stack, &resolving, &visited, &ignoring, breakCycles)
        } catch {
          guard breakCycles else {
            throw error
          }

          stack.append(depNode)
          ignoring.insert(depID)
          try orderSort(depNode, &stack, &resolving, &visited, &ignoring, breakCycles)
          ignoring.remove(depID)
        }
        resolving.remove(depID)
        visited.insert(depID)
      }
    }

    if !stack.contains(node) {
      stack.append(node)
    }
  }
}

extension ApkPackageGraph {
  func findDependencyCycle(node: ApkPackageGraphNode) -> (ApkPackageGraphNode, ApkPackageGraphNode)? {
    var resolving = Set<Int>()
    var visited = Set<Int>()
    return self.findDependencyCycle(node: node, &resolving, &visited)
  }

  func findDependencyCycle(
    node: ApkPackageGraphNode,
    _ resolving: inout Set<ApkIndex.Index>,
    _ visited: inout Set<ApkIndex.Index>
  ) -> (ApkPackageGraphNode, ApkPackageGraphNode)? {
    for dependency in node.children {
      let depNode = self._nodes[dependency.packageID]
      if resolving.contains(depNode.packageID) {
        return (node, depNode)
      }

      if !visited.contains(depNode.packageID) {
        resolving.insert(depNode.packageID)
        if let cycle = findDependencyCycle(node: depNode, &resolving, &visited) {
          return cycle
        }

        resolving.remove(depNode.packageID)
        visited.insert(depNode.packageID)
      }
    }

    return nil
  }

  public func parallelOrderSort(breakCycles: Bool = true) throws(SortError) -> [[ApkPackageGraphNode]] {
    var results = [[ApkPackageGraphNode]]()

    // Map all nodes to all of their children, remove any self dependencies
    var working = self._nodes.reduce(into: [ApkPackageGraphNode: Set<ApkPackageGraphNode>]()) { d, node in
      d[node] = Set(node.children.filter { child in
        if case .dependency = child.constraint {
          !child.versionSpec.isConflict && child.packageID != node.packageID
        } else { false }
      }.map { self._nodes[$0.packageID] })
    }

    // Collect all child nodes that aren't already in the map
    // This should be empty every time
    let extras = working.values.reduce(Set<ApkPackageGraphNode>()) { a, b in
      a.union(b)
    }.subtracting(working.keys)
    assert(extras.isEmpty, "Dangling nodes in the graph")

    // Put all extra nodes into the working map, with an empty set
    extras.forEach {
      working[$0] = []
    }

    while !working.isEmpty {
      // Set of all nodes now with satisfied dependencies
      var set = Set(working
        .filter { _, children in children.isEmpty }
        .map(\.key))

      // If nothing was satisfied in this loop, check for cycles
      // If no cycles exist and the working set is empty, resolve is complete
      if set.isEmpty {
        if working.isEmpty {
          break
        }

        let cycles = working.keys.compactMap { node in
          self.findDependencyCycle(node: node)
        }

        // Error if cycle breaking is turned off
        if !breakCycles {
          throw .cyclicDependency(cycles: cycles.map { node, dependency in
              "\(node) -> \(dependency)"
            }.joined(separator: "\n"))
        }

        // Break cycles by setting the new resolution set to dependencies that cycled
        set = Set(cycles.flatMap { [$0.0, $0.1] })
      }

      // Add installation set to list of installation sets
      results.append(Array(set))

      // Filter the working set for anything that wasn't dealt with this iteration
      working = working.filter { node, _ in
        !set.contains(node)
      }.reduce(into: [ApkPackageGraphNode: Set<ApkPackageGraphNode>]()) { d, node in
        d[node.key] = node.value.subtracting(set)
      }
    }

    return results
  }

  public enum SortError: Error, LocalizedError {
    case cyclicDependency(cycles: String)

    var errorDescription: String {
      switch self {
      case .cyclicDependency(let cycles): "Dependency cycles found:\n\(cycles)"
      }
    }
  }
}