import XCTest
import Foundation
import Darwin

// Mirrors the pure auth decision we expect to exist in SocketServer.swift after fix.
func referenceIsPeerAuthorized(peerUID: uid_t, currentUID: uid_t) -> Bool {
    return peerUID == currentUID
}

final class SecurityHardeningTests: XCTestCase {

    func hostExecutableURL() throws -> URL {
        let fm = FileManager.default
        let candidates = [
            ".build/debug/ReynaCLIHost",
            ".build/arm64-apple-macosx/debug/ReynaCLIHost",
            ".build/x86_64-apple-macosx/debug/ReynaCLIHost",
        ]
        let cwd = fm.currentDirectoryPath
        var tried: [String] = []
        for c in candidates {
            let url = URL(fileURLWithPath: cwd).appendingPathComponent(c)
            tried.append(url.path)
            if fm.isExecutableFile(atPath: url.path) { return url }
        }
        var parent = URL(fileURLWithPath: cwd)
        for _ in 0..<6 {
            let p1 = parent.appendingPathComponent(".build/debug/ReynaCLIHost")
            tried.append(p1.path)
            if fm.isExecutableFile(atPath: p1.path) { return p1 }
            parent = parent.deletingLastPathComponent()
        }
        throw NSError(domain: "SecurityHardeningTests", code: 1, userInfo: [NSLocalizedDescriptionKey: "binary not found Tried:\n"+tried.joined(separator: "\n")])
    }

    func makeShortUniqueDirChecked() throws -> URL {
        let fm = FileManager.default
        for _ in 0..<20 {
            let hex = String(format: "%08x", UInt32.random(in: 0...UInt32.max))
            let url = URL(fileURLWithPath: "/tmp/rh-\(hex)")
            if !fm.fileExists(atPath: url.path) { return url }
        }
        return URL(fileURLWithPath: "/tmp/rh-\(String(format: "%08x", UInt32.random(in: 0...UInt32.max)))")
    }

    final class HostProcess {
        let process: Process
        let socketPath: String
        let tempDir: URL
        init(process: Process, socketPath: String, tempDir: URL) {
            self.process = process; self.socketPath = socketPath; self.tempDir = tempDir
        }
        func terminate() {
            if process.isRunning { process.terminate() }
            let deadline = Date().addingTimeInterval(2)
            while process.isRunning && Date() < deadline { usleep(100_000) }
            if process.isRunning { process.interrupt() }
        }
        deinit { terminate() }
    }

    func startSocketHost(socketPath: String, tempDir: URL) throws -> HostProcess {
        let exe = try hostExecutableURL()
        let process = Process()
        process.executableURL = exe
        process.arguments = ["--socket", socketPath]
        let stderr = Pipe()
        process.standardError = stderr
        process.standardOutput = Pipe()
        process.standardInput = Pipe()
        try process.run()
        let fm = FileManager.default
        let deadline = Date().addingTimeInterval(5)
        while Date() < deadline {
            if fm.fileExists(atPath: socketPath) { break }
            if !process.isRunning {
                let data = stderr.fileHandleForReading.readDataToEndOfFile()
                let s = String(data: data, encoding: .utf8) ?? ""
                throw NSError(domain: "SecurityHardeningTests", code: 2, userInfo: [NSLocalizedDescriptionKey: "Host exited early. stderr: \(s)"])
            }
            usleep(100_000)
        }
        if !fm.fileExists(atPath: socketPath) {
            process.terminate()
            let data = stderr.fileHandleForReading.readDataToEndOfFile()
            let s = String(data: data, encoding: .utf8) ?? ""
            throw NSError(domain: "SecurityHardeningTests", code: 3, userInfo: [NSLocalizedDescriptionKey: "Socket not created at \(socketPath). stderr: \(s)"])
        }
        return HostProcess(process: process, socketPath: socketPath, tempDir: tempDir)
    }

    func socketRequestResponse(socketPath: String, requestLine: String, timeout: TimeInterval = 3) throws -> String {
        let fd = socket(AF_UNIX, SOCK_STREAM, 0)
        guard fd >= 0 else { throw NSError(domain: "SecurityHardeningTests", code: 10, userInfo: nil) }
        defer { close(fd) }
        var addr = sockaddr_un()
        addr.sun_family = sa_family_t(AF_UNIX)
        memset(&addr.sun_path, 0, MemoryLayout.size(ofValue: addr.sun_path))
        _ = socketPath.withCString { cStr in
            withUnsafeMutablePointer(to: &addr.sun_path) { dstPtr in
                dstPtr.withMemoryRebound(to: CChar.self, capacity: 104) { charPtr in
                    strncpy(charPtr, cStr, 103)
                }
            }
        }
        let addrLen = socklen_t(MemoryLayout<sockaddr_un>.size)
        let cr = withUnsafePointer(to: &addr) { ptr in
            ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { saddr in
                connect(fd, saddr, addrLen)
            }
        }
        guard cr == 0 else { throw NSError(domain: "SecurityHardeningTests", code: 11, userInfo: [NSLocalizedDescriptionKey: "connect failed: \(String(cString: strerror(errno)))"]) }
        let toSend = requestLine.hasSuffix("\n") ? requestLine : requestLine + "\n"
        guard let data = toSend.data(using: .utf8) else { throw NSError(domain: "SecurityHardeningTests", code: 12, userInfo: nil) }
        var sent = 0
        while sent < data.count {
            let n = data.withUnsafeBytes { raw in send(fd, raw.baseAddress!.advanced(by: sent), data.count - sent, 0) }
            if n <= 0 { throw NSError(domain: "SecurityHardeningTests", code: 13, userInfo: nil) }
            sent += n
        }
        var responseData = Data()
        var buffer = [UInt8](repeating: 0, count: 4096)
        let start = Date()
        while true {
            if Date().timeIntervalSince(start) > timeout {
                throw NSError(domain: "SecurityHardeningTests", code: 14, userInfo: [NSLocalizedDescriptionKey: "read timeout"])
            }
            var pfd = pollfd(fd: fd, events: Int16(POLLIN), revents: 0)
            let pr = poll(&pfd, 1, 200)
            if pr < 0 { if errno == EINTR { continue }; throw NSError(domain: "SecurityHardeningTests", code: 15, userInfo: nil) }
            if pr == 0 { continue }
            let r = recv(fd, &buffer, buffer.count, 0)
            if r < 0 { if errno == EINTR { continue }; throw NSError(domain: "SecurityHardeningTests", code: 16, userInfo: nil) }
            if r == 0 { break }
            responseData.append(contentsOf: buffer[0..<r])
            if let str = String(data: responseData, encoding: .utf8), str.contains("\n") { break }
        }
        guard let respString = String(data: responseData, encoding: .utf8) else { throw NSError(domain: "SecurityHardeningTests", code: 17, userInfo: nil) }
        let first = respString.split(separator: "\n", omittingEmptySubsequences: false).first.map { String($0) } ?? respString
        return first.trimmingCharacters(in: .whitespacesAndNewlines)
    }

    // MARK: - 1) peer UID pure decision

    func testPeerAuthorizationPureDecision() {
        let me = getuid()
        let foreign: uid_t = (me == 0) ? 1 : 0
        XCTAssertTrue(referenceIsPeerAuthorized(peerUID: me, currentUID: me), "own UID should be authorized")
        XCTAssertFalse(referenceIsPeerAuthorized(peerUID: foreign, currentUID: me), "foreign UID should be rejected")
    }

    func testHostPeerAuthorizationFunctionExists() throws {
        // If implementation exposes isPeerAuthorized, test it indirectly by exercising server.
        // We assert current process connecting is allowed (same UID) – existing health test proves this.
        // For this TDD RED, we also attempt to check source contains getpeereid.
        let fm = FileManager.default
        let srcURL = URL(fileURLWithPath: fm.currentDirectoryPath).appendingPathComponent("Sources/ReynaCLIHostCore/SocketServer.swift")
        // Walk up
        var found: URL? = nil
        var cur = URL(fileURLWithPath: fm.currentDirectoryPath)
        for _ in 0..<8 {
            let cand = cur.appendingPathComponent("Sources/ReynaCLIHostCore/SocketServer.swift")
            if fm.fileExists(atPath: cand.path) { found = cand; break }
            let candOld = cur.appendingPathComponent("Sources/ReynaCLIHost/SocketServer.swift")
            if fm.fileExists(atPath: candOld.path) { found = candOld; break }
            let cand2 = cur.appendingPathComponent("native/ReynaCLIHost/Sources/ReynaCLIHostCore/SocketServer.swift")
            if fm.fileExists(atPath: cand2.path) { found = cand2; break }
            let cand2Old = cur.appendingPathComponent("native/ReynaCLIHost/Sources/ReynaCLIHost/SocketServer.swift")
            if fm.fileExists(atPath: cand2Old.path) { found = cand2Old; break }
            cur = cur.deletingLastPathComponent()
        }
        let url = found ?? srcURL
        guard let content = try? String(contentsOf: url) else {
            XCTFail("Could not read SocketServer.swift at \(url.path)")
            return
        }
        XCTAssertTrue(content.contains("getpeereid") || content.contains("getpeerid"), "SocketServer.swift must call getpeereid for peer UID check")
    }

    // MARK: - 2) signal-handler safety

    func testSignalHandlerNoUnsafeGlobals() throws {
        let fm = FileManager.default
        var cur = URL(fileURLWithPath: fm.currentDirectoryPath)
        var srcPath: URL? = nil
        for _ in 0..<8 {
            let cand = cur.appendingPathComponent("Sources/ReynaCLIHostCore/SocketServer.swift")
            if fm.fileExists(atPath: cand.path) { srcPath = cand; break }
            let candOld = cur.appendingPathComponent("Sources/ReynaCLIHost/SocketServer.swift")
            if fm.fileExists(atPath: candOld.path) { srcPath = candOld; break }
            let cand2 = cur.appendingPathComponent("native/ReynaCLIHost/Sources/ReynaCLIHostCore/SocketServer.swift")
            if fm.fileExists(atPath: cand2.path) { srcPath = cand2; break }
            let cand2Old = cur.appendingPathComponent("native/ReynaCLIHost/Sources/ReynaCLIHost/SocketServer.swift")
            if fm.fileExists(atPath: cand2Old.path) { srcPath = cand2Old; break }
            cur = cur.deletingLastPathComponent()
        }
        guard let url = srcPath, let content = try? String(contentsOf: url) else {
            XCTFail("Cannot find SocketServer.swift for signal safety check")
            return
        }
        // No Swift mutable global storing path
        XCTAssertFalse(content.contains("gSocketPathCStr"), "Should not have Swift mutable global gSocketPathCStr")
        XCTAssertFalse(content.contains("nonisolated(unsafe)"), "Should not have nonisolated(unsafe) global for signal handling")
        // No unsafeBitCast to sig_t
        XCTAssertFalse(content.contains("unsafeBitCast") && content.contains("sig_t"), "Should not use unsafeBitCast to sig_t")
        // Signal handler itself should not be Swift using stat/lstat
        // Check that reynaSocketSignalHandler Swift func with lstat/stat is gone
        // Allow C file to handle signals; here check that Swift file doesn't define reynaSocketSignalHandler with lstat
        // This part will pass when we move handler to C target.
        // Also check Package.swift contains C target
        var pkgURL: URL? = nil
        cur = URL(fileURLWithPath: fm.currentDirectoryPath)
        for _ in 0..<8 {
            let cand = cur.appendingPathComponent("Package.swift")
            if fm.fileExists(atPath: cand.path) { pkgURL = cand; break }
            cur = cur.deletingLastPathComponent()
        }
        if let purl = pkgURL, (try? String(contentsOf: purl)) != nil {
            // Should contain C target for signal support OR no signal unsafe patterns above already covers
            // Not failing if C target missing yet, but signal safety tests still need to show RED via earlier checks
        }
    }

    // MARK: - 3) path validation

    func testSocketPathRejectsDotDotComponents() throws {
        // Host should refuse paths containing .. or . components
        let fm = FileManager.default
        let unique = try makeShortUniqueDirChecked()
        try fm.createDirectory(at: unique, withIntermediateDirectories: true, attributes: [.posixPermissions: 0o700])
        defer { try? fm.removeItem(at: unique) }
        let dotDotPath = unique.appendingPathComponent("../evil.sock").path
        // This contains .. – should be rejected, process exits non-zero quickly
        let exe = try hostExecutableURL()
        let proc = Process()
        proc.executableURL = exe
        proc.arguments = ["--socket", dotDotPath]
        proc.standardError = Pipe()
        proc.standardOutput = Pipe()
        try proc.run()
        let deadline = Date().addingTimeInterval(2)
        while proc.isRunning && Date() < deadline { usleep(100_000) }
        if proc.isRunning {
            proc.terminate()
            XCTFail("Host should reject path containing .. and exit, but kept running for \(dotDotPath)")
        } else {
            XCTAssertNotEqual(proc.terminationStatus, 0, "Should exit non-zero for .. path")
        }
        // Also test ./ component
        let dotPath = unique.appendingPathComponent("./evil.sock").path
        let proc2 = Process()
        proc2.executableURL = exe
        proc2.arguments = ["--socket", dotPath]
        proc2.standardError = Pipe()
        proc2.standardOutput = Pipe()
        try proc2.run()
        let deadline2 = Date().addingTimeInterval(2)
        while proc2.isRunning && Date() < deadline2 { usleep(100_000) }
        if proc2.isRunning {
            proc2.terminate()
            XCTFail("Host should reject path containing . and exit")
        } else {
            XCTAssertNotEqual(proc2.terminationStatus, 0, "Should exit non-zero for . path")
        }
    }

    func testSocketParentRejectsSymlink() throws {
        let fm = FileManager.default
        let unique = try makeShortUniqueDirChecked()
        try fm.createDirectory(at: unique, withIntermediateDirectories: true, attributes: [.posixPermissions: 0o700])
        defer { try? fm.removeItem(at: unique) }
        let real = unique.appendingPathComponent("real")
        try fm.createDirectory(at: real, withIntermediateDirectories: false, attributes: [.posixPermissions: 0o700])
        let link = unique.appendingPathComponent("linkdir")
        try fm.createSymbolicLink(at: link, withDestinationURL: real)
        let sockPath = link.appendingPathComponent("reyna.sock").path
        let exe = try hostExecutableURL()
        let proc = Process()
        proc.executableURL = exe
        proc.arguments = ["--socket", sockPath]
        proc.standardError = Pipe()
        proc.standardOutput = Pipe()
        try proc.run()
        let deadline = Date().addingTimeInterval(2)
        while proc.isRunning && Date() < deadline { usleep(100_000) }
        if proc.isRunning {
            proc.terminate()
            XCTFail("Host should reject symlink parent and exit")
        } else {
            XCTAssertNotEqual(proc.terminationStatus, 0, "Should exit non-zero when parent is symlink")
        }
    }

    func testSocketParentRejectsWorldWritable() throws {
        let fm = FileManager.default
        let unique = try makeShortUniqueDirChecked()
        try fm.createDirectory(at: unique, withIntermediateDirectories: true, attributes: [.posixPermissions: 0o700])
        defer { try? fm.removeItem(at: unique) }
        // Make parent world-writable 0777 but owned by us – should be rejected
        chmod(unique.path, 0o777)
        let sockPath = unique.appendingPathComponent("reyna.sock").path
        let exe = try hostExecutableURL()
        let proc = Process()
        proc.executableURL = exe
        proc.arguments = ["--socket", sockPath]
        proc.standardError = Pipe()
        proc.standardOutput = Pipe()
        try proc.run()
        let deadline = Date().addingTimeInterval(2)
        while proc.isRunning && Date() < deadline { usleep(100_000) }
        if proc.isRunning {
            proc.terminate()
            XCTFail("Host should reject world-writable dedicated parent")
        } else {
            XCTAssertNotEqual(proc.terminationStatus, 0, "Should exit non-zero for world-writable parent")
        }
    }

    // MARK: - 4) recv timeout / slow client

    func testSlowClientDoesNotBlockHealthClient() throws {
        let fm = FileManager.default
        let unique = try makeShortUniqueDirChecked()
        try fm.createDirectory(at: unique, withIntermediateDirectories: true, attributes: [.posixPermissions: 0o700])
        defer { try? fm.removeItem(at: unique) }
        let sockPath = unique.appendingPathComponent("reyna.sock").path
        let host = try startSocketHost(socketPath: sockPath, tempDir: unique)
        defer { host.terminate() }

        // Open slow client that connects and sends partial without newline and keeps open
        let slowFd = socket(AF_UNIX, SOCK_STREAM, 0)
        XCTAssertTrue(slowFd >= 0)
        var addr = sockaddr_un()
        addr.sun_family = sa_family_t(AF_UNIX)
        memset(&addr.sun_path, 0, MemoryLayout.size(ofValue: addr.sun_path))
        _ = sockPath.withCString { cStr in
            withUnsafeMutablePointer(to: &addr.sun_path) { dst in
                dst.withMemoryRebound(to: CChar.self, capacity: 104) { p in strncpy(p, cStr, 103) }
            }
        }
        let len = socklen_t(MemoryLayout<sockaddr_un>.size)
        let cr = withUnsafePointer(to: &addr) { ptr in
            ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { saddr in connect(slowFd, saddr, len) }
        }
        XCTAssertEqual(cr, 0, "slow client connect should succeed")
        // Send incomplete data (no newline)
        let partial = "{\"id\":\"slow\",\"operation\":\"service.health\""
        _ = partial.withCString { cStr in send(slowFd, cStr, strlen(cStr), 0) }

        // Give server a moment to be blocked in recv if vulnerable
        usleep(300_000)

        // Now try health client – should succeed within timeout + small margin, not blocked forever.
        // Server should have a recv timeout ~5s, so this health client should succeed in < (timeout+2)s
        let start = Date()
        let req = #"{"id":"fast","operation":"service.health","arguments":{}}"#
        var gotResponse = false
        var lastError: Error? = nil
        for _ in 0..<3 {
            do {
                let respLine = try socketRequestResponse(socketPath: sockPath, requestLine: req, timeout: 6)
                if let data = respLine.data(using: .utf8),
                   let obj = try JSONSerialization.jsonObject(with: data) as? [String: Any],
                   obj["id"] as? String == "fast",
                   obj["ok"] as? Bool == true {
                    gotResponse = true
                    break
                }
            } catch {
                lastError = error
                usleep(200_000)
            }
        }
        let elapsed = Date().timeIntervalSince(start)
        close(slowFd)

        XCTAssertTrue(gotResponse, "Fast client should succeed despite slow client; lastError: \(String(describing: lastError)) elapsed: \(elapsed)s")
        XCTAssertLessThan(elapsed, 8, "Slow client should not block health client beyond timeout; elapsed \(elapsed)s")
    }

    // MARK: - 5) oversized-line boundary

    func testOversizedLineWithNewlineInSameChunk() throws {
        // This tests the bug where buf+chunk > limit and newline in most recent chunk is ignored.
        // Build a valid JSON line exactly 500 bytes, then newline, then extra garbage in same TCP chunk.
        // The server must accept the first line (<=64KiB) even if same recv includes bytes after newline.
        let fm = FileManager.default
        let unique = try makeShortUniqueDirChecked()
        try fm.createDirectory(at: unique, withIntermediateDirectories: true, attributes: [.posixPermissions: 0o700])
        defer { try? fm.removeItem(at: unique) }
        let sockPath = unique.appendingPathComponent("reyna.sock").path
        let host = try startSocketHost(socketPath: sockPath, tempDir: unique)
        defer { host.terminate() }

        let fd = socket(AF_UNIX, SOCK_STREAM, 0)
        XCTAssertTrue(fd >= 0)
        defer { close(fd) }
        var addr = sockaddr_un()
        addr.sun_family = sa_family_t(AF_UNIX)
        memset(&addr.sun_path, 0, MemoryLayout.size(ofValue: addr.sun_path))
        _ = sockPath.withCString { cStr in
            withUnsafeMutablePointer(to: &addr.sun_path) { dst in
                dst.withMemoryRebound(to: CChar.self, capacity: 104) { p in strncpy(p, cStr, 103) }
            }
        }
        let addrLen = socklen_t(MemoryLayout<sockaddr_un>.size)
        let cr = withUnsafePointer(to: &addr) { ptr in
            ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { saddr in connect(fd, saddr, addrLen) }
        }
        XCTAssertEqual(cr, 0)

        // Build payload: first line is valid health request (<64KiB) + "\n" + extra bytes that would make total > limit if counted, but second line invalid.
        // Actually to trigger bug, we need first line <=64KiB, but the chunk that contains newline also contains extra bytes making total > limit? The bug checks buf.count+n > limit before looking at newline in new chunk.
        // So simulate by sending one large send that is exactly 64KiB + extra.
        // We'll send a health request (~50 bytes) + "\n" + 70KiB of 'X's in ONE send call. The server reads up to 4096 at a time, but could still get newline in first recv.
        // Better: send health request + "\n" + large extra, and ensure server still returns ok for first line, not payload_too_large.
        let healthReq = #"{"id":"line-ok","operation":"service.health","arguments":{}}"#
        let extra = String(repeating: "X", count: 70*1024)
        let combined = healthReq + "\n" + extra
        guard let data = combined.data(using: .utf8) else { XCTFail("encode fail"); return }
        // Ignore SIGPIPE in this process to avoid signal 13 when server closes early
        signal(SIGPIPE, SIG_IGN)
        var sent = 0
        while sent < data.count {
            let n = data.withUnsafeBytes { ptr in send(fd, ptr.baseAddress!.advanced(by: sent), data.count - sent, 0) }
            if n <= 0 {
                if errno == EPIPE || errno == ECONNRESET { break }
                break
            }
            sent += n
        }

        var responseData = Data()
        var buf = [UInt8](repeating: 0, count: 4096)
        let start = Date()
        while true {
            if Date().timeIntervalSince(start) > 3 { break }
            var pfd = pollfd(fd: fd, events: Int16(POLLIN), revents: 0)
            let pr = poll(&pfd, 1, 200)
            if pr <= 0 { continue }
            let r = recv(fd, &buf, buf.count, 0)
            if r <= 0 { break }
            responseData.append(contentsOf: buf[0..<r])
            if let s = String(data: responseData, encoding: .utf8), s.contains("\n") { break }
        }
        guard let respStr = String(data: responseData, encoding: .utf8) else {
            XCTFail("No utf8 response")
            return
        }
        let firstLine = respStr.split(separator: "\n").first.map { String($0) } ?? respStr
        guard let d = firstLine.data(using: .utf8),
              let obj = try? JSONSerialization.jsonObject(with: d) as? [String: Any] else {
            XCTFail("Response not JSON: \(firstLine)")
            return
        }
        XCTAssertEqual(obj["id"] as? String, "line-ok", "Should preserve id of first line")
        XCTAssertEqual(obj["ok"] as? Bool, true, "First line <=64KiB should be accepted even when same chunk has extra bytes after newline, got: \(obj)")
    }

    func testOversizedFirstLineStillRejected() throws {
        let fm = FileManager.default
        let unique = try makeShortUniqueDirChecked()
        try fm.createDirectory(at: unique, withIntermediateDirectories: true, attributes: [.posixPermissions: 0o700])
        defer { try? fm.removeItem(at: unique) }
        let sockPath = unique.appendingPathComponent("reyna.sock").path
        let host = try startSocketHost(socketPath: sockPath, tempDir: unique)
        defer { host.terminate() }

        let largeString = String(repeating: "A", count: 70*1024)
        let req = #"{"id":"big","operation":"service.health","arguments":{},"data":"\#(largeString)"}"#
        XCTAssertTrue(req.utf8.count > 65536)
        let respLine = try socketRequestResponse(socketPath: sockPath, requestLine: req)
        guard let data = respLine.data(using: .utf8),
              let obj = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            XCTFail("Not JSON: \(respLine)"); return
        }
        XCTAssertEqual(obj["ok"] as? Bool, false)
        let code = (obj["error"] as? [String: Any])?["code"] as? String ?? ""
        XCTAssertTrue(code.contains("too_large") || code.contains("payload") || code.contains("invalid"), "Expected too_large code, got \(code)")
    }
}