import Foundation

/// Drives a feed with URLSession and command execution: ask the core for requests, send them,
/// hand the responses back, repeat until it reports. The loop is the same one the TypeScript host
/// runs, so a provider with two dependent calls (teams, then billing) works here with no extra machinery.
public enum UsageDriver {
	public typealias CommandRunner = @Sendable (UsageRequest) async throws -> UsageResponse

	public static func fetch(
		provider: String,
		credentialJSON: String,
		nowMs: Int64 = Int64(Date().timeIntervalSince1970 * 1000),
		session: URLSession = URLSession(configuration: .ephemeral),
		runCommand: CommandRunner? = nil,
		maxRounds: Int = 4
	) async throws -> UsageReport {
		let state = UsageState()
		var responses: [UsageResponse] = []

		for _ in 0..<maxRounds {
			switch try UsageCore.advance(state: state, provider: provider, credentialJSON: credentialJSON, responses: responses, nowMs: nowMs) {
			case .report(let report):
				return report
			case .requests(let requests):
				responses = []
				for request in requests {
					if request.kind == "command" {
						if let runCommand {
							do {
								responses.append(try await runCommand(request))
							} catch {
								responses.append(UsageResponse(status: 126, body: error.localizedDescription))
							}
						} else {
							// The one answer a feed is allowed to read as "this host will not run things".
							responses.append(UsageResponse(status: 127, body: "this host does not run commands (\(request.command ?? ""))"))
						}
					} else {
						responses.append(try await send(request, session: session))
					}
				}
			}
		}
		return UsageReport(
			provider: provider,
			providerName: nil,
			accountId: nil,
			email: nil,
			planType: nil,
			limits: [],
			balance: nil,
			notes: nil,
			error: "usage feed for \(provider) did not finish in \(maxRounds) rounds"
		)
	}

	private final class PipeBuffer: @unchecked Sendable {
		private let lock = NSLock()
		private let maxBytes: Int
		private var data = Data()

		init(maxBytes: Int) {
			self.maxBytes = maxBytes
		}

		func append(_ chunk: Data) {
			if chunk.isEmpty { return }
			lock.withLock {
				if data.count < maxBytes {
					data.append(chunk.prefix(maxBytes - data.count))
				}
			}
		}

		func snapshot() -> Data {
			lock.withLock { data }
		}
	}

	/// Run a command request using Process with the environment allowlist, timeout, and stdout cap.
	@Sendable public static func runCommand(_ request: UsageRequest) async throws -> UsageResponse {
		guard let cmd = request.command, !cmd.isEmpty else {
			return UsageResponse(status: 127, body: "no command specified")
		}

		let process = Process()
		let stdoutPipe = Pipe()
		let stderrPipe = Pipe()
		let stdinPipe = Pipe()

		// Environment allowlist: PATH, HOME, LANG, LC_ALL, LC_CTYPE, TZ and proxy variables
		let allowedKeys: Set<String> = [
			"PATH", "HOME", "LANG", "LC_ALL", "LC_CTYPE", "TZ",
			"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY",
			"http_proxy", "https_proxy", "all_proxy", "no_proxy",
		]
		let systemEnv = ProcessInfo.processInfo.environment
		var filteredEnv: [String: String] = [:]
		for key in allowedKeys {
			if let val = systemEnv[key] {
				filteredEnv[key] = val
			}
		}
		process.environment = filteredEnv

		let args = request.args ?? []
		if cmd.contains("/") {
			process.executableURL = URL(fileURLWithPath: cmd)
			process.arguments = args
		} else {
			let pathVar = filteredEnv["PATH"] ?? "/usr/bin:/bin:/usr/local/bin"
			var foundURL: URL? = nil
			for dir in pathVar.split(separator: ":") {
				let candidate = URL(fileURLWithPath: String(dir)).appendingPathComponent(cmd)
				if FileManager.default.isExecutableFile(atPath: candidate.path) {
					foundURL = candidate
					break
				}
			}
			if let foundURL {
				process.executableURL = foundURL
				process.arguments = args
			} else {
				// Not falling back to `/usr/bin/env`: it would exit 127 with the same code the protocol
				// reserves for "this host does not run commands", and the feed would read a missing
				// binary as a refusal. 126 says what actually happened.
				return UsageResponse(status: 126, body: "\(cmd) is not on PATH")
			}
		}

		process.standardOutput = stdoutPipe
		process.standardError = stderrPipe
		if request.stdin != nil {
			process.standardInput = stdinPipe
		}

		let timeoutSec: TimeInterval = Double(request.timeoutMs ?? 15_000) / 1000.0
		let maxBytes = 64 * 1024
		let stdoutBuffer = PipeBuffer(maxBytes: maxBytes)
		let stderrBuffer = PipeBuffer(maxBytes: maxBytes)

		stdoutPipe.fileHandleForReading.readabilityHandler = { handle in
			stdoutBuffer.append(handle.availableData)
		}

		stderrPipe.fileHandleForReading.readabilityHandler = { handle in
			stderrBuffer.append(handle.availableData)
		}

		do {
			try process.run()
		} catch {
			stdoutPipe.fileHandleForReading.readabilityHandler = nil
			stderrPipe.fileHandleForReading.readabilityHandler = nil
			// 126 is "found it, could not run it", which is not the same answer as 127 ("this host does
			// not run commands"): a feed that treats 127 as "fall back to another path" must not be told
			// that when the real problem is a missing binary.
			return UsageResponse(status: 126, body: error.localizedDescription)
		}

		if let stdin = request.stdin, let data = stdin.data(using: .utf8) {
			try? stdinPipe.fileHandleForWriting.write(contentsOf: data)
			try? stdinPipe.fileHandleForWriting.close()
		}

		let deadline = Date().addingTimeInterval(timeoutSec)
		while process.isRunning && Date() < deadline {
			try? await Task.sleep(nanoseconds: 20_000_000)
		}

		if process.isRunning {
			process.terminate()
			var termWait = 0
			while process.isRunning && termWait < 10 {
				try? await Task.sleep(nanoseconds: 20_000_000)
				termWait += 1
			}
		}

		stdoutPipe.fileHandleForReading.readabilityHandler = nil
		stderrPipe.fileHandleForReading.readabilityHandler = nil

		stdoutBuffer.append(stdoutPipe.fileHandleForReading.availableData)
		stderrBuffer.append(stderrPipe.fileHandleForReading.availableData)

		var body = String(data: stdoutBuffer.snapshot(), encoding: .utf8) ?? ""
		let exitCode = Int(process.terminationStatus)

		// A command that wrote nothing to stdout but failed still has something to say, and saying it
		// on stderr is how a CLI reports a missing argument or a bad credential. Without this the feed
		// only sees an empty body and cannot tell the user why nothing came back.
		if exitCode != 0, body.isEmpty, let errorText = String(data: stderrBuffer.snapshot(), encoding: .utf8), !errorText.isEmpty {
			body = errorText
		}

		return UsageResponse(status: exitCode, body: body)
	}

	/// The core asks for one credential kind; a real app resolves it from its own store (or from
	/// prime-agent's auth.json). The shape is provider-defined and documented in usage-feeds.md.
	public static func credentialJSON(kind: String, value: String, accountId: String? = nil, email: String? = nil, extra: [String: String] = [:]) -> String {
		var object: [String: Any] = ["kind": kind]
		switch kind {
		case "oauth": object["access"] = value
		case "apiKey": object["apiKey"] = value
		case "session": object["session"] = (try? JSONSerialization.jsonObject(with: Data(value.utf8))) ?? [:]
		default: object[kind] = value
		}
		if let accountId { object["accountId"] = accountId }
		if let email { object["email"] = email }
		for (key, val) in extra { object[key] = val }
		let data = try? JSONSerialization.data(withJSONObject: object)
		return data.flatMap { String(data: $0, encoding: .utf8) } ?? "{}"
	}

	private static func send(_ request: UsageRequest, session: URLSession) async throws -> UsageResponse {
		guard let url = URL(string: request.url) else {
			return UsageResponse(status: 0, body: "invalid url: \(request.url)")
		}
		var urlRequest = URLRequest(url: url)
		urlRequest.httpMethod = request.method
		for (name, value) in request.headers {
			urlRequest.setValue(value, forHTTPHeaderField: name)
		}
		if let body = request.body {
			urlRequest.httpBody = Data(body.utf8)
		}
		do {
			let (data, response) = try await session.data(for: urlRequest)
			let status = (response as? HTTPURLResponse)?.statusCode ?? 0
			return UsageResponse(status: status, body: String(data: data, encoding: .utf8) ?? "")
		} catch {
			// A transport failure is data for the feed, not an exception: the report carries it.
			return UsageResponse(status: 0, body: "transport error: \(error.localizedDescription)")
		}
	}
}