// SPDX-FileCopyrightText: 2025 The Project Pterodactyl Developers
//
// SPDX-License-Identifier: MPL-2.0

import Foundation

fileprivate struct TokenLocation {
	let startLine: Int
	let startColumn: Int
}

fileprivate struct LocatedToken {
	let token: Token
	let location: TokenLocation
	let lines: [String.SubSequence]

	init(token: Token, location: TokenLocation) {
		self.token = token
		self.location = location
		self.lines = token.text.split(separator: "\n", omittingEmptySubsequences: false)
	}

	var nextLocation: TokenLocation {
		let startLine = location.startLine + lines.count - 1
		let startColumn = if lines.count > 1 { lines.last!.utf16.count } else { lines.last?.utf16.count ?? location.startColumn }
		return TokenLocation(startLine: startLine, startColumn: startColumn)
	}
}

public struct BlockLayoutProcessor {
	private let locatedTokens: [LocatedToken]

	public init(tokens: [Token]) {
		var locatedTokens: [LocatedToken] = []
		var location = TokenLocation(startLine: 0, startColumn: 0)
		for token in tokens {
			let locatedToken = LocatedToken(token: token, location: location)
			locatedTokens.append(locatedToken)
			location = locatedToken.nextLocation
		}

		self.locatedTokens = locatedTokens
	}


	public func layout() -> [Token] {
		var result: [Token] = []
		var indentStack: [Int] = [0]
		var previousLine = 0
		var firstTokenInBlock = false

		for (index, locatedToken) in locatedTokens.enumerated() {
			guard locatedToken.token.kind != .eof else { break }
			guard locatedToken.token.kind.canDetermineLayoutColumn else {
				result.append(locatedToken.token)
				continue
			}

			if locatedToken.location.startLine > previousLine {
				while indentStack.count > 1 && locatedToken.location.startColumn < indentStack.last! {
					indentStack.removeLast()
					result.append(Token(kind: .blockEnd, text: ""))
				}

				if !firstTokenInBlock && indentStack.count > 1 && locatedToken.location.startColumn == indentStack.last! {
					result.append(Token(kind: .blockSep, text: ""))
				}
			}

			result.append(locatedToken.token)

			if locatedToken.token.kind.isBlockHerald {
				firstTokenInBlock = true

				if let nextToken = locatedTokens[index...].first(where: { $0.location.startLine > locatedToken.location.startLine && $0.token.kind.canDetermineLayoutColumn }) {
					result.append(Token(kind: .blockBegin, text: ""))
					indentStack.append(nextToken.location.startColumn)
				} else {
					result.append(Token(kind: .blockBegin, text: ""))
					result.append(Token(kind: .blockEnd, text: ""))
				}
			} else {
				firstTokenInBlock = false
			}

			previousLine = locatedToken.location.startLine
		}

		while indentStack.count > 1 {
			indentStack.removeLast()
			result.append(Token(kind: .blockEnd, text: ""))
		}

		if let eof = locatedTokens.last, eof.token.kind == .eof {
			result.append(eof.token)
		}

		return result
	}
}