import Foundation

struct DiscoveredServer: Sendable {
	var hostname: String
	var port: Int
	var socketType: String
}

enum AutoDiscovery {
	static func discoverIMAP(for email: String) async -> DiscoveredServer? {
		guard let domain = email.split(separator: "@").last.map(String.init) else { return nil }

		if let server = await queryISPDB(domain: domain) {
			return server
		}

		if let server = await querySRV(domain: domain) {
			return server
		}

		return nil
	}

	private static func queryISPDB(domain: String) async -> DiscoveredServer? {
		let url = URL(string: "https://autoconfig.thunderbird.net/v1.1/\(domain)")!
		guard let (data, response) = try? await URLSession.shared.data(from: url),
			  let httpResponse = response as? HTTPURLResponse,
			  httpResponse.statusCode == 200,
			  let xml = String(data: data, encoding: .utf8)
		else { return nil }

		return parseISPDBXML(xml)
	}

	private static func parseISPDBXML(_ xml: String) -> DiscoveredServer? {
		guard let imapRange = xml.range(of: "<incomingServer type=\"imap\">"),
			  let endRange = xml.range(of: "</incomingServer>", range: imapRange.upperBound..<xml.endIndex)
		else { return nil }

		let section = String(xml[imapRange.upperBound..<endRange.lowerBound])

		func extractTag(_ tag: String) -> String? {
			guard let start = section.range(of: "<\(tag)>"),
				  let end = section.range(of: "</\(tag)>", range: start.upperBound..<section.endIndex)
			else { return nil }
			return String(section[start.upperBound..<end.lowerBound])
		}

		guard let hostname = extractTag("hostname"),
			  let portStr = extractTag("port"),
			  let port = Int(portStr)
		else { return nil }

		let socketType = extractTag("socketType") ?? "SSL"
		return DiscoveredServer(hostname: hostname, port: port, socketType: socketType)
	}

	private static func querySRV(domain: String) async -> DiscoveredServer? {
		let candidates = [
			"imap.\(domain)",
			"mail.\(domain)",
		]
		for candidate in candidates {
			if await testConnection(host: candidate, port: 993) {
				return DiscoveredServer(hostname: candidate, port: 993, socketType: "SSL")
			}
		}
		return nil
	}

	private static func testConnection(host: String, port: Int) async -> Bool {
		do {
			return try await withThrowingTaskGroup(of: Bool.self) { group in
				group.addTask {
					let task = URLSession.shared.streamTask(withHostName: host, port: port)
					task.resume()
					let (data, _) = try await task.readData(ofMinLength: 1, maxLength: 1024, timeout: 3)
					task.cancel()
					return data.map { !$0.isEmpty } ?? false
				}
				group.addTask {
					try await Task.sleep(for: .seconds(3))
					return false
				}
				let result = try await group.next() ?? false
				group.cancelAll()
				return result
			}
		} catch {
			return false
		}
	}
}
