import Foundation
import Darwin
import CSignalSupport

// Pure decision extracted for unit testing.
func isPeerAuthorized(peerUID: uid_t, currentUID: uid_t) -> Bool {
    return peerUID == currentUID
}

// MARK: - Path validation
// Trust boundary: see SocketPathValidation.swift for full documentation.
// Tiered model:
//  1) platform allowlist root-owned (uid 0, no g/o write except tmp)
//  2) user-owned intermediates: uid current, no g/o *write* (mode & 022 == 0) -> allows 0700/0750/0755, rejects 0770/0777
//  3) dedicated runtime (immediate socket parent): uid current, mode & 077 == 0 -> requires 0700 family only.
// Socket itself 0600.

// MARK: - Reused single-component validator (authoritative)
// NOTE: this is the ONLY place allowed to decide if an existing component is safe.
// It must stay in sync with validateParentChainPureResultProvider logic.
func validateExistingComponentLiveOrThrow(path: String, info: LStatInfo, currentUID: uid_t, isDedicatedRuntimeParent: Bool = false) throws {
    // Centralized call to shared validation in SocketPathValidation
    try validateSingleLStatInfoOrThrow(path: path, info: info, currentUID: currentUID, isDedicatedRuntimeParent: isDedicatedRuntimeParent)
}

func ensureParentDirectories(for socketPath: String) throws {
    // Authoritative validation reused; fail-closed on lstat errors
    try rejectIfDotComponentsPure(in: socketPath)
    try validateExistingParentChainLive(for: socketPath)

    let fm = FileManager.default
    let url = URL(fileURLWithPath: socketPath)
    let parent = url.deletingLastPathComponent()
    let parentPath = parent.path
    if parentPath.isEmpty { return }

    let comps = parent.pathComponents
    var cur = ""
    for comp in comps {
        if comp == "/" {
            cur = "/"
            continue
        }
        if cur == "/" {
            cur = "/" + comp
        } else if cur.isEmpty {
            cur = comp
        } else {
            cur = cur + "/" + comp
        }

        switch liveLStatResultProvider(path: cur) {
        case .absent:
            // Create missing component with 0700 – privacy preserving. Even intermediates now get 0700.
            do {
                try fm.createDirectory(atPath: cur, withIntermediateDirectories: false, attributes: [.posixPermissions: 0o700])
                chmod(cur, 0o700)
            } catch {
                // mkdir race: re-lstat and revalidate rather than assuming missing
                switch liveLStatResultProvider(path: cur) {
                case .absent:
                    throw error
                case .failed(let ec):
                    throw NSError(domain: "SocketServer", code: 22, userInfo: [NSLocalizedDescriptionKey: "lstat failed for \(cur): \(String(cString: strerror(ec)))"])
                case .present(let info):
                    do {
                        let isDedicated = (cur == parentPath)
                        try validateSingleLStatInfoOrThrow(path: cur, info: info, currentUID: getuid(), isDedicatedRuntimeParent: isDedicated)
                    } catch {
                        throw error
                    }
                    continue
                }
            }
        case .failed(let ec):
            throw NSError(domain: "SocketServer", code: 22, userInfo: [NSLocalizedDescriptionKey: "lstat failed for \(cur): \(String(cString: strerror(ec)))"])
        case .present(let info):
            // Reuse authoritative single-component validator (no duplicated policy)
            let isDedicated = (cur == parentPath)
            try validateSingleLStatInfoOrThrow(path: cur, info: info, currentUID: getuid(), isDedicatedRuntimeParent: isDedicated)
        }
    }
}

func safeUnlinkIfStaleSocket(at path: String) throws {
    var st = stat()
    let r = lstat(path, &st)
    if r != 0 {
        if errno == ENOENT { return }
        throw NSError(domain: "SocketServer", code: 11, userInfo: [NSLocalizedDescriptionKey: "lstat failed for \(path): \(String(cString: strerror(errno)))"])
    }
    let isSock = (st.st_mode & S_IFMT) == S_IFSOCK
    if !isSock {
        throw NSError(domain: "SocketServer", code: 12, userInfo: [NSLocalizedDescriptionKey: "Refusing to unlink: \(path) exists and is not a socket"])
    }
    if st.st_uid != getuid() {
        throw NSError(domain: "SocketServer", code: 13, userInfo: [NSLocalizedDescriptionKey: "Refusing to unlink: socket at \(path) not owned by current uid"])
    }
    if unlink(path) != 0 && errno != ENOENT {
        throw NSError(domain: "SocketServer", code: 14, userInfo: [NSLocalizedDescriptionKey: "Failed to unlink stale socket \(path): \(String(cString: strerror(errno)))"])
    }
}

private let kMaxRequestBytes = 64 * 1024
private let kClientRecvTimeoutSec = 5

private func makeErrorResponse(id: String, code: String, message: String) -> Data? {
    let err = ErrorPayload(code: code, message: message)
    let resp = Response(id: id, ok: false, result: nil, error: err)
    guard let json = try? JSONEncoder().encode(resp),
          let str = String(data: json, encoding: .utf8) else { return nil }
    return (str + "\n").data(using: .utf8)
}

private func processRequestData(_ data: Data) -> Data? {
    if data.isEmpty { return nil }
    if let s = String(data: data, encoding: .utf8),
       s.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
        return nil
    }
    if let req = try? JSONDecoder().decode(Request.self, from: data) {
        let resp = dispatch(request: req)
        guard let json = try? JSONEncoder().encode(resp),
              let str = String(data: json, encoding: .utf8) else { return nil }
        return (str + "\n").data(using: .utf8)
    } else {
        let recovered = extractRecoverableId(from: data) ?? ""
        return makeErrorResponse(id: recovered, code: "invalid_request", message: "Invalid request JSON")
    }
}

func runSocketServer(socketPath: String) -> Never {
    if !socketPath.hasPrefix("/") {
        fputs("error: --socket path must be absolute: \(socketPath)\n", stderr)
        Darwin.exit(2)
    }
    if socketPath.utf8.count >= 104 {
        fputs("error: --socket path too long\n", stderr)
        Darwin.exit(2)
    }

    do {
        try ensureParentDirectories(for: socketPath)
        try safeUnlinkIfStaleSocket(at: socketPath)
    } catch {
        fputs("error: \(error.localizedDescription)\n", stderr)
        Darwin.exit(3)
    }

    // Store for signal cleanup in C
    socketPath.withCString { cStr in
        reyna_store_socket_path(cStr)
    }
    reyna_install_signal_handlers()

    let fd = socket(AF_UNIX, SOCK_STREAM, 0)
    if fd < 0 {
        fputs("error: socket() failed: \(String(cString: strerror(errno)))\n", stderr)
        Darwin.exit(4)
    }
    _ = fcntl(fd, F_SETFD, FD_CLOEXEC)

    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) { dst in
            dst.withMemoryRebound(to: CChar.self, capacity: 104) { p in
                strncpy(p, cStr, 103)
            }
        }
    }

    let oldMask = umask(0o077)
    let bindRes = withUnsafePointer(to: &addr) { ptr in
        ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { saddr in
            bind(fd, saddr, socklen_t(MemoryLayout<sockaddr_un>.size))
        }
    }
    umask(oldMask)

    if bindRes != 0 {
        fputs("error: bind() \(socketPath): \(String(cString: strerror(errno)))\n", stderr)
        close(fd)
        Darwin.exit(5)
    }

    if chmod(socketPath, 0o600) != 0 {
        fputs("warning: chmod 0600 failed: \(String(cString: strerror(errno)))\n", stderr)
    }

    if listen(fd, 32) != 0 {
        fputs("error: listen() failed: \(String(cString: strerror(errno)))\n", stderr)
        close(fd)
        reyna_cleanup_socket_sync()
        Darwin.exit(6)
    }

    // Main loop: one client at a time, one request per connection
    while true {
        let cfd = accept(fd, nil, nil)
        if cfd < 0 {
            if errno == EINTR { continue }
            fputs("error: accept() failed: \(String(cString: strerror(errno)))\n", stderr)
            break
        }

        // --- Peer credential check (macOS getpeereid) ---
        var peerEuid: uid_t = 0
        var peerEgid: gid_t = 0
        if getpeereid(cfd, &peerEuid, &peerEgid) != 0 {
            // If we cannot obtain peer credentials, reject
            close(cfd)
            continue
        }
        if !isPeerAuthorized(peerUID: peerEuid, currentUID: getuid()) {
            if let d = makeErrorResponse(id: "", code: "unauthorized", message: "Peer UID not authorized") {
                _ = d.withUnsafeBytes { p in send(cfd, p.baseAddress!, p.count, 0) }
            }
            close(cfd)
            continue
        }

        // Set receive timeout to bound slow clients
        var tv = timeval()
        tv.tv_sec = kClientRecvTimeoutSec
        tv.tv_usec = 0
        setsockopt(cfd, SOL_SOCKET, SO_RCVTIMEO, &tv, socklen_t(MemoryLayout<timeval>.size))

        var buf = Data()
        buf.reserveCapacity(8192)
        var tmp = [UInt8](repeating: 0, count: 4096)
        var exceeded = false
        var gotAny = false
        var timedOut = false

        // Poll-based timeout additionally enforced
        while true {
            // Wait for data with timeout
            var pfd = pollfd(fd: cfd, events: Int16(POLLIN), revents: 0)
            let pollTimeoutMs: Int32 = Int32(kClientRecvTimeoutSec * 1000)
            let pr = poll(&pfd, 1, pollTimeoutMs)
            if pr < 0 {
                if errno == EINTR { continue }
                break
            }
            if pr == 0 {
                // timeout
                timedOut = true
                break
            }
            let n = recv(cfd, &tmp, tmp.count, 0)
            if n < 0 {
                if errno == EINTR { continue }
                if errno == EWOULDBLOCK || errno == EAGAIN {
                    timedOut = true
                    break
                }
                break
            }
            if n == 0 { break }
            gotAny = true

            // Oversized handling with newline-in-same-chunk fix
            if buf.count + n > kMaxRequestBytes {
                // Look for newline in the new chunk
                var newlineIdx: Int? = nil
                for i in 0..<n {
                    if tmp[i] == 0x0A {
                        newlineIdx = i
                        break
                    }
                }
                if let nl = newlineIdx {
                    // First line length = buf.count + nl (excluding newline char itself)
                    let firstLineLen = buf.count + nl
                    if firstLineLen <= kMaxRequestBytes {
                        // Accept up to newline and ignore rest
                        buf.append(contentsOf: tmp[0..<nl])
                        // Break to process – we have a complete line within limit
                        break
                    } else {
                        exceeded = true
                        break
                    }
                } else {
                    // No newline in this chunk and would exceed -> oversized
                    exceeded = true
                    break
                }
            }
            buf.append(contentsOf: tmp[0..<n])
            if buf.contains(0x0A) { break }
        }

        if timedOut {
            close(cfd)
            continue
        }

        if !gotAny && !exceeded {
            close(cfd)
            continue
        }

        if exceeded {
            if let d = makeErrorResponse(id: "", code: "payload_too_large", message: "Request exceeds 64KiB limit") {
                _ = d.withUnsafeBytes { p in send(cfd, p.baseAddress!, p.count, 0) }
            }
            close(cfd)
            continue
        }

        let lineData: Data
        if let idx = buf.firstIndex(of: 0x0A) {
            lineData = buf.prefix(upTo: idx)
        } else {
            lineData = buf
        }

        if lineData.count > kMaxRequestBytes {
            if let d = makeErrorResponse(id: "", code: "payload_too_large", message: "Request exceeds 64KiB limit") {
                _ = d.withUnsafeBytes { p in send(cfd, p.baseAddress!, p.count, 0) }
            }
            close(cfd)
            continue
        }

        if let resp = processRequestData(lineData) {
            _ = resp.withUnsafeBytes { p in
                var sent = 0
                while sent < resp.count {
                    let n = send(cfd, p.baseAddress!.advanced(by: sent), resp.count - sent, 0)
                    if n <= 0 { break }
                    sent += n
                }
            }
        }
        close(cfd)
    }

    close(fd)
    reyna_cleanup_socket_sync()
    Darwin.exit(0)
}