#!/usr/bin/env swift
//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2026 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

// Generate @implementation Swift stub from the AST extracted from clang files.
//
// This script parses the JSON output from -emit-polyglot-ast and generates
// a stub Swift source file with the API and Objective-C code commented out.

import Foundation

struct PolyglotAST: Codable {
  let formatVersion: Int
  let compilerVersion: String
  let topLevelDecls: [TopLevelDecl]
}

struct TopLevelDecl: Codable {
  let formatKind: String
  let target: String?
  let fullText: String?
  let declKind: String?
  let className: String?
  let properties: [Property]?
  let methods: [Method]?
  let swiftSignature: String?
  let body: String?
}

struct Property: Codable {
  let swiftSignature: String?
}

struct Method: Codable {
  let body: String?
  let swiftSignature: String?
}

let idePlaceholder = "<#code#>"

func generateStub(_ ast: PolyglotAST) -> [String] {
  var output: [String] = []

  output.append("// Stub generated by \(URL(fileURLWithPath: args[0]).lastPathComponent)")
  output.append("// Compiler version: \(ast.compilerVersion)")
  output.append("// Format version: \(ast.formatVersion)")

  // Group interfaces and implementations of a same class into one container
  var mergedContainers: [String: [TopLevelDecl]] = [:]
  for decl in ast.topLevelDecls {
    if decl.formatKind == "container" {
      guard let className = decl.className else { continue }
      if mergedContainers[className] == nil {
        mergedContainers[className] = [decl]
      } else {
        mergedContainers[className]!.append(decl)
      }
    }
  }

  // Output stub code
  for decl in ast.topLevelDecls {
    let formatKind = decl.formatKind

    switch formatKind {
    case "import":
      if let target = decl.target {
        output.append("import \(target)")
      }

    case "unhandled":
      if let fullText = decl.fullText {
        output.append("// Unhandled: \(fullText)")
      }

    case "container":
      guard let className = decl.className else {
        break
      }
      if let containers = mergedContainers[className] {
        let merged = processMergedContainers(containers, className: className)
        output.append(contentsOf: merged)
        mergedContainers.removeValue(forKey: className)
      }

    default:
      continue
    }
  }

  return output
}

func processMergedContainers(_ containers: [TopLevelDecl], className: String) -> [String] {
  var output: [String] = []

  // Extension
  output.append("@objc @implementation")
  output.append("extension \(className) {")

  // Properties
  for container in containers {
    if let properties = container.properties {
      for prop in properties {
        if let signature = prop.swiftSignature {
          output.append("    " + signature)
        }
      }
    }
  }

  // Methods
  for container in containers {
    guard let methods = container.methods else {
      continue
    }
    for method in methods {
      if let signature = method.swiftSignature {
        if let body = method.body {
          // We have the original source for the body.
          let bodyLines = formatBodyAsComment(body)
          let bodyText = bodyLines.joined(separator: "\n        ")
          if signature.contains(idePlaceholder) {
            // The compiler attached a stub body to the signature.
            var stub = signature.replacingOccurrences(of: idePlaceholder, with: bodyText)
            output.append(stub)
          } else {
            // The compiler returned a signature with no body.
            output.append(signature.dropLast() + " {")
            output.append("        " + bodyText)
            output.append("    }")
          }
        } else {
          output.append(signature.dropLast() + " {")
          output.append("    }")
        }
      }
    }
  }

  output.append("}")

  return output
}

func formatBodyAsComment(_ body: String) -> [String] {
  guard !body.isEmpty else { return [] }

  var lines = body.split(separator: "\n", omittingEmptySubsequences: false)
  // Strip surrounding { } lines if both are present
  if let first = lines.first, let last = lines.last,
     first.trimmingCharacters(in: .whitespaces) == "{",
     last.trimmingCharacters(in: .whitespaces) == "}" {
    lines = lines.dropFirst().dropLast()
  }

  // Strip the first line's leading whitespace from all lines
  let firstNonEmpty = lines.first { !$0.trimmingCharacters(in: .whitespaces).isEmpty }
  let commonIndent = firstNonEmpty?.prefix(while: { $0 == " " || $0 == "\t" }).count ?? 0

  return lines.map { line in
    let trimmed = line.trimmingCharacters(in: .whitespaces).isEmpty
      ? "" : String(line.dropFirst(commonIndent))
    return "  // \(trimmed)"
  }
}

let args = CommandLine.arguments
guard args.count == 3 else {
  fputs("Usage: \(args[0]) <input.json> <output.swift>\n", stderr)
  fputs("Generate stuf Swift file for an @implemention class\n", stderr)
  exit(1)
}
let inputPath = args[1]
let outputPath = args[2]

// Read and parse JSON AST
guard let ast = try? Data(contentsOf: URL(fileURLWithPath: inputPath)) else {
  fputs("Error: Cannot read input file: \(inputPath)\n", stderr)
  exit(1)
}

let polyglotAST: PolyglotAST
do {
  let decoder = JSONDecoder()
  polyglotAST = try decoder.decode(PolyglotAST.self, from: ast)
} catch {
  fputs("Error: Cannot parse JSON: \(error.localizedDescription)\n", stderr)
  exit(1)
}

// Main logic
let apiLines = generateStub(polyglotAST)

// Write output
let outputContent = apiLines.joined(separator: "\n") + "\n"
guard let outputData = outputContent.data(using: .utf8) else {
  fputs("Error: Cannot encode output as UTF-8\n", stderr)
  exit(1)
}
do {
  try outputData.write(to: URL(fileURLWithPath: outputPath))
} catch {
  fputs("Error: Cannot write output file: \(error.localizedDescription)\n", stderr)
  exit(1)
}
