import CryptoKit
import Darwin
import Foundation

private let maximumCapturedBytes = 64 * 1024

private struct ExecutionRecord: Encodable {
    let schemaVersion: String
    let binarySha256: String
    let binaryIdentityMatched: Bool
    let argvFingerprintSha256: String
    let startedAt: String
    let endedAt: String
    let exitCode: Int32?
    let terminationSignal: Int32?
    let stdout: String
    let stderr: String
    let outputTruncated: Bool
    let sessionIdentity: String
    let cellIdentity: String
    let launchStatus: String
    let snapshotCleanupStatus: String

    enum CodingKeys: String, CodingKey {
        case schemaVersion, binarySha256, binaryIdentityMatched, argvFingerprintSha256
        case startedAt, endedAt, exitCode, terminationSignal, stdout, stderr
        case outputTruncated, sessionIdentity, cellIdentity, launchStatus, snapshotCleanupStatus
    }

    func encode(to encoder: Encoder) throws {
        var container = encoder.container(keyedBy: CodingKeys.self)
        try container.encode(schemaVersion, forKey: .schemaVersion)
        try container.encode(binarySha256, forKey: .binarySha256)
        try container.encode(binaryIdentityMatched, forKey: .binaryIdentityMatched)
        try container.encode(argvFingerprintSha256, forKey: .argvFingerprintSha256)
        try container.encode(startedAt, forKey: .startedAt)
        try container.encode(endedAt, forKey: .endedAt)
        if let exitCode { try container.encode(exitCode, forKey: .exitCode) }
        else { try container.encodeNil(forKey: .exitCode) }
        if let terminationSignal { try container.encode(terminationSignal, forKey: .terminationSignal) }
        else { try container.encodeNil(forKey: .terminationSignal) }
        try container.encode(stdout, forKey: .stdout)
        try container.encode(stderr, forKey: .stderr)
        try container.encode(outputTruncated, forKey: .outputTruncated)
        try container.encode(sessionIdentity, forKey: .sessionIdentity)
        try container.encode(cellIdentity, forKey: .cellIdentity)
        try container.encode(launchStatus, forKey: .launchStatus)
        try container.encode(snapshotCleanupStatus, forKey: .snapshotCleanupStatus)
    }
}

private func sha256(_ data: Data) -> String {
    SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
}

private func fileSha256(_ url: URL) throws -> String {
    let handle = try FileHandle(forReadingFrom: url)
    defer { try? handle.close() }
    var digest = SHA256()
    while true {
        let data = try handle.read(upToCount: 64 * 1024) ?? Data()
        if data.isEmpty { break }
        digest.update(data: data)
    }
    return digest.finalize().map { String(format: "%02x", $0) }.joined()
}

private struct ExecutableSnapshot {
    let directoryURL: URL
    let executableURL: URL
    let sha256: String
    let handle: FileHandle
}

private struct SnapshotFailure: Error {
    let stage: String
    let cleanupFailure: String?
}

private func sha256(_ handle: FileHandle) throws -> String {
    try handle.seek(toOffset: 0)
    var digest = SHA256()
    while true {
        let data = try handle.read(upToCount: 64 * 1024) ?? Data()
        if data.isEmpty { break }
        digest.update(data: data)
    }
    return digest.finalize().map { String(format: "%02x", $0) }.joined()
}

private func sameObject(_ left: stat, _ right: stat) -> Bool {
    left.st_dev == right.st_dev && left.st_ino == right.st_ino
}

private func cleanupFailedSnapshot(
    handle: FileHandle?, executableURL: URL?, directoryURL: URL?
) -> String? {
    var failures: [String] = []
    if let handle {
        if fchflags(handle.fileDescriptor, 0) != 0 { failures.append("bound-immutable-flag") }
        try? handle.close()
    }
    if let executableURL {
        _ = chflags(executableURL.path, 0)
        do { try FileManager.default.removeItem(at: executableURL) }
        catch where (error as NSError).code != NSFileNoSuchFileError { failures.append("executable") }
        catch {}
    }
    if let directoryURL {
        do { try FileManager.default.removeItem(at: directoryURL) }
        catch where (error as NSError).code != NSFileNoSuchFileError { failures.append("directory") }
        catch {}
    }
    return failures.isEmpty ? nil : "SNAPSHOT_CLEANUP_FAILED:\(failures.joined(separator: ","))"
}

private func createExecutableSnapshot(from sourceURL: URL) throws -> ExecutableSnapshot {
    #if EVIDENCE_RUNNER_TESTING
    if ProcessInfo.processInfo.environment["MINDRAW_EVIDENCE_TEST_FAIL_SNAPSHOT_CREATE"] == "1" {
        throw SnapshotFailure(stage: "SNAPSHOT_CREATE_FAILED", cleanupFailure: nil)
    }
    #endif
    let manager = FileManager.default
    let directoryURL = manager.temporaryDirectory
        .appendingPathComponent("mindraw-evidence-exec-\(UUID().uuidString)", isDirectory: true)
    do {
        try manager.createDirectory(
            at: directoryURL,
            withIntermediateDirectories: false,
            attributes: [.posixPermissions: 0o700]
        )
    } catch {
        throw SnapshotFailure(stage: "SNAPSHOT_CREATE_FAILED", cleanupFailure: nil)
    }
    var directoryStatus = stat()
    guard lstat(directoryURL.path, &directoryStatus) == 0,
          (directoryStatus.st_mode & S_IFMT) == S_IFDIR,
          directoryStatus.st_mode & 0o777 == 0o700 else {
        try? manager.removeItem(at: directoryURL)
        throw SnapshotFailure(stage: "SNAPSHOT_CREATE_FAILED", cleanupFailure: nil)
    }

    let executableURL = directoryURL.appendingPathComponent("verified-tool", isDirectory: false)
    let descriptor = open(executableURL.path, O_RDWR | O_CREAT | O_EXCL | O_NOFOLLOW, 0o600)
    guard descriptor >= 0 else {
        try? manager.removeItem(at: directoryURL)
        throw SnapshotFailure(stage: "SNAPSHOT_CREATE_FAILED", cleanupFailure: nil)
    }
    let output = FileHandle(fileDescriptor: descriptor, closeOnDealloc: true)
    do {
        let input = try FileHandle(forReadingFrom: sourceURL)
        defer { try? input.close() }
        var digest = SHA256()
        while true {
            let data = try input.read(upToCount: 64 * 1024) ?? Data()
            if data.isEmpty { break }
            digest.update(data: data)
            try output.write(contentsOf: data)
        }
        try output.synchronize()
        let copiedSha = digest.finalize().map { String(format: "%02x", $0) }.joined()
        var createdStatus = stat()
        guard fstat(descriptor, &createdStatus) == 0,
              (createdStatus.st_mode & S_IFMT) == S_IFREG,
              fchmod(descriptor, 0o500) == 0 else {
            throw SnapshotFailure(stage: "SNAPSHOT_VERIFY_FAILED", cleanupFailure: nil)
        }
        #if EVIDENCE_RUNNER_TESTING
        if ProcessInfo.processInfo.environment["MINDRAW_EVIDENCE_TEST_FAIL_SNAPSHOT_VERIFY"] == "1" {
            throw SnapshotFailure(stage: "SNAPSHOT_VERIFY_FAILED", cleanupFailure: nil)
        }
        guard try sha256(output) == copiedSha else {
            throw SnapshotFailure(stage: "SNAPSHOT_VERIFY_FAILED", cleanupFailure: nil)
        }
        if let replacement = ProcessInfo.processInfo.environment["MINDRAW_EVIDENCE_TEST_REPLACE_SNAPSHOT_WITH"] {
            try manager.removeItem(at: executableURL)
            try manager.copyItem(at: URL(fileURLWithPath: replacement), to: executableURL)
            try manager.setAttributes([.posixPermissions: 0o500], ofItemAtPath: executableURL.path)
        }
        #endif
        guard fchflags(descriptor, UInt32(UF_IMMUTABLE)) == 0 else {
            throw SnapshotFailure(stage: "SNAPSHOT_HARDENING_FAILED", cleanupFailure: nil)
        }
        var boundStatus = stat()
        var pathStatus = stat()
        guard fstat(descriptor, &boundStatus) == 0,
              lstat(executableURL.path, &pathStatus) == 0,
              sameObject(createdStatus, boundStatus), sameObject(boundStatus, pathStatus) else {
            throw SnapshotFailure(stage: "SNAPSHOT_OBJECT_REPLACED", cleanupFailure: nil)
        }
        guard (boundStatus.st_mode & S_IFMT) == S_IFREG,
              boundStatus.st_mode & 0o777 == 0o500,
              boundStatus.st_flags & UInt32(UF_IMMUTABLE) != 0,
              try sha256(output) == copiedSha else {
            throw SnapshotFailure(stage: "SNAPSHOT_VERIFY_FAILED", cleanupFailure: nil)
        }
        return ExecutableSnapshot(
            directoryURL: directoryURL,
            executableURL: executableURL,
            sha256: copiedSha,
            handle: output
        )
    } catch {
        let cleanupFailure = cleanupFailedSnapshot(
            handle: output, executableURL: executableURL, directoryURL: directoryURL
        )
        if let failure = error as? SnapshotFailure {
            throw SnapshotFailure(stage: failure.stage, cleanupFailure: cleanupFailure)
        }
        throw SnapshotFailure(stage: "SNAPSHOT_VERIFY_FAILED", cleanupFailure: cleanupFailure)
    }
}

private func cleanupExecutableSnapshot(_ snapshot: ExecutableSnapshot) -> String? {
    let manager = FileManager.default
    var failures: [String] = []
    if fchflags(snapshot.handle.fileDescriptor, 0) != 0 {
        failures.append("immutable-flag")
    }
    try? snapshot.handle.close()
    if failures.isEmpty {
        do { try manager.removeItem(at: snapshot.executableURL) }
        catch { failures.append("executable") }
    }
    do { try manager.removeItem(at: snapshot.directoryURL) }
    catch { failures.append("directory") }
    #if EVIDENCE_RUNNER_TESTING
    if ProcessInfo.processInfo.environment["MINDRAW_EVIDENCE_TEST_FAIL_SNAPSHOT_CLEANUP"] == "1" {
        failures.append("injected")
    }
    #endif
    return failures.isEmpty ? nil : "SNAPSHOT_CLEANUP_FAILED:\(failures.joined(separator: ","))"
}

private func validOpaqueIdentity(_ value: String) -> Bool {
    value.range(of: #"^[A-Za-z0-9._-]{1,128}$"#, options: .regularExpression) != nil
}

private func replaceMatches(_ pattern: String, in value: String, with replacement: String) -> String {
    guard let expression = try? NSRegularExpression(pattern: pattern) else { return value }
    let range = NSRange(value.startIndex..<value.endIndex, in: value)
    return expression.stringByReplacingMatches(in: value, range: range, withTemplate: replacement)
}

private func sanitizeOutput(_ data: Data, statePath: String, sourceIdentifier: String) -> (String, Bool) {
    let truncated = data.count > maximumCapturedBytes
    let bounded = data.prefix(maximumCapturedBytes)
    var value = String(decoding: bounded, as: UTF8.self)
    value = value.replacingOccurrences(of: statePath, with: "<STATE_PATH>")
    value = value.replacingOccurrences(of: sourceIdentifier, with: "<SOURCE_IDENTIFIER>")
    value = replaceMatches(#"/(?:Users|private|var|tmp)/[^\s]+"#, in: value, with: "<PRIVATE_PATH>")
    value = replaceMatches(#"(?i)\bpid=\d+"#, in: value, with: "pid=<PID>")
    value = replaceMatches(#"(?i)\bwindow(?:_?id)?=\d+"#, in: value, with: "window=<WINDOW_ID>")
    if truncated { value += "<OUTPUT_TRUNCATED>\n" }
    return (value, truncated)
}

private func argvFingerprint(binarySha: String, mode: String) -> String {
    sha256(Data("\(binarySha)|\(mode)|<STATE_PATH>|<SOURCE_IDENTIFIER>".utf8))
}

private func writeRecord(_ record: ExecutionRecord, to url: URL) throws {
    let encoder = JSONEncoder()
    encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
    let data = try encoder.encode(record)
    try data.write(to: url, options: .atomic)
    let recordSha = Data("\(sha256(data))\n".utf8)
    try recordSha.write(to: URL(fileURLWithPath: url.path + ".sha256"), options: .atomic)
}

private final class BoundedCapture {
    private let lock = NSLock()
    private var data = Data()
    private(set) var truncated = false

    func append(_ incoming: Data) {
        lock.lock()
        defer { lock.unlock() }
        let remaining = max(0, maximumCapturedBytes + 1 - data.count)
        if remaining > 0 { data.append(incoming.prefix(remaining)) }
        if incoming.count > remaining || data.count > maximumCapturedBytes { truncated = true }
    }

    func snapshot() -> Data {
        lock.lock()
        defer { lock.unlock() }
        return data
    }
}

private func run() -> Int32 {
    let arguments = CommandLine.arguments
    guard arguments.count == 10, arguments[1] == "run" else {
        fputs("usage: EvidenceToolRunner run record.json expectedBinarySha sessionIdentity cellIdentity binary isolate|restore state.json sourceIdentifier\n", stderr)
        return 64
    }
    let recordURL = URL(fileURLWithPath: arguments[2])
    let expectedSha = arguments[3]
    let sessionIdentity = arguments[4]
    let cellIdentity = arguments[5]
    let binaryURL = URL(fileURLWithPath: arguments[6])
    let mode = arguments[7]
    let statePath = arguments[8]
    let sourceIdentifier = arguments[9]
    guard expectedSha.range(of: #"^[a-f0-9]{64}$"#, options: .regularExpression) != nil,
          validOpaqueIdentity(sessionIdentity), validOpaqueIdentity(cellIdentity),
          ["isolate", "restore"].contains(mode), !statePath.isEmpty, !sourceIdentifier.isEmpty else {
        fputs("invalid evidence runner arguments\n", stderr)
        return 64
    }

    let formatter = ISO8601DateFormatter()
    formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
    let startedAt = formatter.string(from: Date())
    let snapshot: ExecutableSnapshot
    do {
        snapshot = try createExecutableSnapshot(from: binaryURL)
    } catch {
        let failure = error as? SnapshotFailure
        let record = ExecutionRecord(
            schemaVersion: "mindraw.qa-evidence-exec.v1",
            binarySha256: expectedSha,
            binaryIdentityMatched: false,
            argvFingerprintSha256: argvFingerprint(binarySha: expectedSha, mode: mode),
            startedAt: startedAt,
            endedAt: formatter.string(from: Date()),
            exitCode: nil,
            terminationSignal: nil,
            stdout: "",
            stderr: failure?.cleanupFailure.map { $0 + "\n" } ?? "",
            outputTruncated: false,
            sessionIdentity: sessionIdentity,
            cellIdentity: cellIdentity,
            launchStatus: failure?.stage ?? "SNAPSHOT_VERIFY_FAILED",
            snapshotCleanupStatus: failure?.cleanupFailure == nil ? "CLEANUP_COMPLETE" : "CLEANUP_FAILED"
        )
        do { try writeRecord(record, to: recordURL) } catch { return 77 }
        fputs("executable snapshot unavailable\n", stderr)
        return 77
    }
    let actualSha = snapshot.sha256
    let fingerprint = argvFingerprint(binarySha: actualSha, mode: mode)
    if actualSha != expectedSha {
        let cleanupFailure = cleanupExecutableSnapshot(snapshot)
        let endedAt = formatter.string(from: Date())
        let record = ExecutionRecord(
            schemaVersion: "mindraw.qa-evidence-exec.v1",
            binarySha256: actualSha,
            binaryIdentityMatched: false,
            argvFingerprintSha256: fingerprint,
            startedAt: startedAt,
            endedAt: endedAt,
            exitCode: nil,
            terminationSignal: nil,
            stdout: "",
            stderr: cleanupFailure ?? "",
            outputTruncated: false,
            sessionIdentity: sessionIdentity,
            cellIdentity: cellIdentity,
            launchStatus: cleanupFailure == nil ? "BINARY_IDENTITY_MISMATCH" : "SNAPSHOT_CLEANUP_FAILED",
            snapshotCleanupStatus: cleanupFailure == nil ? "CLEANUP_COMPLETE" : "CLEANUP_FAILED"
        )
        do { try writeRecord(record, to: recordURL) } catch { return 77 }
        return cleanupFailure == nil ? 78 : 77
    }

    #if EVIDENCE_RUNNER_TESTING
    if let replacement = ProcessInfo.processInfo.environment["MINDRAW_EVIDENCE_TEST_REPLACE_SOURCE_WITH"] {
        do {
            try FileManager.default.removeItem(at: binaryURL)
            try FileManager.default.copyItem(at: URL(fileURLWithPath: replacement), to: binaryURL)
            try FileManager.default.setAttributes([.posixPermissions: 0o700], ofItemAtPath: binaryURL.path)
        } catch {
            _ = cleanupExecutableSnapshot(snapshot)
            return 77
        }
    }
    #endif

    let process = Process()
    let stdoutPipe = Pipe()
    let stderrPipe = Pipe()
    let stdoutCapture = BoundedCapture()
    let stderrCapture = BoundedCapture()
    process.executableURL = snapshot.executableURL
    process.arguments = [mode, statePath, sourceIdentifier]
    process.standardOutput = stdoutPipe
    process.standardError = stderrPipe
    stdoutPipe.fileHandleForReading.readabilityHandler = { handle in
        let data = handle.availableData
        if !data.isEmpty { stdoutCapture.append(data) }
    }
    stderrPipe.fileHandleForReading.readabilityHandler = { handle in
        let data = handle.availableData
        if !data.isEmpty { stderrCapture.append(data) }
    }
    do {
        try process.run()
    } catch {
        let cleanupFailure = cleanupExecutableSnapshot(snapshot)
        stdoutPipe.fileHandleForReading.readabilityHandler = nil
        stderrPipe.fileHandleForReading.readabilityHandler = nil
        let record = ExecutionRecord(
            schemaVersion: "mindraw.qa-evidence-exec.v1",
            binarySha256: actualSha,
            binaryIdentityMatched: cleanupFailure == nil,
            argvFingerprintSha256: fingerprint,
            startedAt: startedAt,
            endedAt: formatter.string(from: Date()),
            exitCode: nil,
            terminationSignal: nil,
            stdout: "",
            stderr: cleanupFailure.map { $0 + "\n" } ?? "",
            outputTruncated: false,
            sessionIdentity: sessionIdentity,
            cellIdentity: cellIdentity,
            launchStatus: cleanupFailure == nil ? "LAUNCH_FAILED" : "SNAPSHOT_CLEANUP_FAILED",
            snapshotCleanupStatus: cleanupFailure == nil ? "CLEANUP_COMPLETE" : "CLEANUP_FAILED"
        )
        do { try writeRecord(record, to: recordURL) } catch { return 77 }
        return 77
    }
    process.waitUntilExit()
    let cleanupFailure = cleanupExecutableSnapshot(snapshot)
    stdoutPipe.fileHandleForReading.readabilityHandler = nil
    stderrPipe.fileHandleForReading.readabilityHandler = nil
    stdoutCapture.append(stdoutPipe.fileHandleForReading.readDataToEndOfFile())
    stderrCapture.append(stderrPipe.fileHandleForReading.readDataToEndOfFile())

    let (safeStdout, stdoutTruncated) = sanitizeOutput(
        stdoutCapture.snapshot(), statePath: statePath, sourceIdentifier: sourceIdentifier
    )
    let (safeStderr, stderrTruncated) = sanitizeOutput(
        stderrCapture.snapshot(), statePath: statePath, sourceIdentifier: sourceIdentifier
    )
    let signal = process.terminationReason == .uncaughtSignal ? process.terminationStatus : nil
    let exitCode = process.terminationReason == .exit ? process.terminationStatus : nil
    let record = ExecutionRecord(
        schemaVersion: "mindraw.qa-evidence-exec.v1",
        binarySha256: actualSha,
        binaryIdentityMatched: cleanupFailure == nil,
        argvFingerprintSha256: fingerprint,
        startedAt: startedAt,
        endedAt: formatter.string(from: Date()),
        exitCode: exitCode,
        terminationSignal: signal,
        stdout: safeStdout,
        stderr: cleanupFailure.map { safeStderr + $0 + "\n" } ?? safeStderr,
        outputTruncated: stdoutTruncated || stderrTruncated || stdoutCapture.truncated || stderrCapture.truncated,
        sessionIdentity: sessionIdentity,
        cellIdentity: cellIdentity,
        launchStatus: cleanupFailure == nil ? "EXECUTED" : "SNAPSHOT_CLEANUP_FAILED",
        snapshotCleanupStatus: cleanupFailure == nil ? "CLEANUP_COMPLETE" : "CLEANUP_FAILED"
    )
    do { try writeRecord(record, to: recordURL) } catch { return 77 }
    return cleanupFailure == nil ? (exitCode ?? 76) : 77
}

private func selfTest() -> Int32 {
    let state = "/private/tmp/private-state.json"
    let source = "private-source"
    let sanitized = sanitizeOutput(
        Data("/Users/private/name \(state) \(source) pid=123 window_id=456\n".utf8),
        statePath: state,
        sourceIdentifier: source
    ).0
    guard sanitized == "<PRIVATE_PATH> <STATE_PATH> <SOURCE_IDENTIFIER> pid=<PID> window=<WINDOW_ID>\n",
          validOpaqueIdentity("session-a"), !validOpaqueIdentity("session/path"),
          argvFingerprint(binarySha: String(repeating: "a", count: 64), mode: "isolate").count == 64 else {
        return 1
    }
    print("self_test=PASS bounded_output=true private_values_redacted=true identity_mismatch_fail_closed=true verified_snapshot_launch=true")
    return 0
}

if CommandLine.arguments == [CommandLine.arguments[0], "--self-test"] {
    exit(selfTest())
}
exit(run())