#!/usr/bin/swift

import Foundation

let inputFile = "input.yaml"
let outputFile = "output.yaml"

var output: [String] = []
var operationIds: [String: Int] = [:]
var scanning = false
var skipCount = 0

private struct AnyOfChild {
    var line: String
    var indentation: String {
        return line.components(separatedBy: .alphanumerics.union(.init(charactersIn: "-"))).first!
    }

    var key: String {
        return line
            .drop(while: { $0.isWhitespace || $0 == "-" })
            .components(separatedBy: ":").first!
    }
    var value: String? {
        let value = line
            .lastIndex(of: ":")
            .flatMap { line.index($0, offsetBy: 1, limitedBy: line.index(before: line.endIndex)) }
            .map({ String(line[$0...].drop(while: \.isWhitespace)) })

        return value
    }
}

extension [AnyOfChild] {
    var key: String { contains(where: { $0.key == "$ref" }) ? "$ref" : "type" }
}

do {
    let lines = try String(contentsOfFile: inputFile, encoding: .utf8).components(separatedBy: .newlines)
    for index in lines.indices {
        if skipCount > 0 {
            skipCount -= 1
            continue
        }

        var line = lines[index]

        if line.contains("anyOf:") {
            let indentation = line.components(separatedBy: "anyOf:").first!
            var children: [AnyOfChild] = []

            var currentIndex = index
            while let nextIndex = lines.index(currentIndex, offsetBy: 1, limitedBy: lines.index(before: lines.endIndex)) {
                let child = AnyOfChild(line: lines[nextIndex])
                guard child.indentation.count > indentation.count else { break }
                children.append(child)
                currentIndex = nextIndex
            }

            // these are special cases we dont handle yet
            let abortKeys = ["items", "enum"]
            if children.contains(where: { abortKeys.contains($0.key) }) {
                output.append(line)
                continue
            }

            // for now, only process items where there are `null`s
            let containsNull = children
                .filter { $0.value != nil }
                .contains(where: { $0.value!.contains("null") })
            if !containsNull {
                output.append(line)
                continue
            }

            var extracts: [AnyOfChild] = []
            var squashes: [AnyOfChild] = []

            for child in children {
                if ["type", "$ref"].contains(child.key) {
                    // this will be squashed into an array like:
                    // type: [string, 'null']
                    // or
                    // $ref: ['#/foo/bar', 'null']
                    squashes.append(child)
                } else {
                    // this will be pulled back to the 'parent'
                    extracts.append(child)
                }
            }

            let newLine = "\(indentation)\(squashes.key): [\(squashes.compactMap(\.value).joined(separator: ", "))]"
            output.append(newLine)

            for extract in extracts where extract.value != nil {
                let newLine = "\(indentation)\(extract.key): \(extract.value!)"
                output.append(newLine)
            }

            skipCount = children.count

            continue
        }

        if line.contains("operationId:") {
            let id = line.components(separatedBy: "operationId:").last!.trimmingCharacters(in: .whitespacesAndNewlines)
            let count = operationIds[id, default: 0]
            operationIds[id] = count + 1
            line.append(count > 0 ? "\(count)" : "")
            output.append(line)
            continue
        }

        // No processing needed, output normal line
        output.append(line)
    }

    let data = output.joined(separator: "\n").data(using: .utf8)!
    try data.write(to: URL(filePath: outputFile))

} catch let error as NSError {
    print("Ooops! Something went wrong: \(error)")
}
