import AppKit
import Darwin

// Chrome's pipe transport belongs to one process. All hands processes on this
// account share that owner through a socket in the account's private temp dir.
private let chromeSocketDir = NSTemporaryDirectory() + "byre-chrome"
private let chromeSocketPath = chromeSocketDir + "/control"

private func chromeAddress() throws -> sockaddr_un {
    var address = sockaddr_un()
    address.sun_family = sa_family_t(AF_UNIX)
    let bytes = Array(chromeSocketPath.utf8) + [0]
    guard bytes.count <= MemoryLayout.size(ofValue: address.sun_path) else {
        throw ToolError("Chrome's private socket path is too long")
    }
    withUnsafeMutableBytes(of: &address.sun_path) { $0.copyBytes(from: bytes) }
    address.sun_len = UInt8(MemoryLayout<sockaddr_un>.size)
    return address
}

private func connectChromeBroker() throws -> Int32 {
    var address = try chromeAddress()
    let fd = socket(AF_UNIX, SOCK_STREAM, 0)
    guard fd >= 0 else { throw ToolError("create Chrome socket: errno \(errno)") }
    let result = withUnsafePointer(to: &address) {
        $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
            Darwin.connect(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
        }
    }
    guard result == 0 else {
        close(fd)
        throw ToolError("Chrome's private helper is not running")
    }
    var noSignal: Int32 = 1
    setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &noSignal, socklen_t(MemoryLayout<Int32>.size))
    return fd
}

// Both the broker and Chrome are launched in their own process groups, so ending
// one agent's hands process cannot take down the shared browser.
func spawnChromeProcess(_ executable: String, _ args: [String], pipes: (Int32, Int32)? = nil) throws -> pid_t {
    var actions: posix_spawn_file_actions_t?
    var attrs: posix_spawnattr_t?
    posix_spawn_file_actions_init(&actions)
    posix_spawnattr_init(&attrs)
    defer {
        posix_spawn_file_actions_destroy(&actions)
        posix_spawnattr_destroy(&attrs)
    }
    posix_spawnattr_setflags(&attrs, Int16(POSIX_SPAWN_SETPGROUP | POSIX_SPAWN_CLOEXEC_DEFAULT))
    posix_spawnattr_setpgroup(&attrs, 0)
    for fd: Int32 in [0, 1, 2] {
        posix_spawn_file_actions_addopen(&actions, fd, "/dev/null", O_RDWR, 0)
    }
    if let (input, output) = pipes {
        posix_spawn_file_actions_adddup2(&actions, input, 3)
        posix_spawn_file_actions_adddup2(&actions, output, 4)
    }
    let argv = ([executable] + args).map { strdup($0) } + [nil]
    defer { for p in argv { free(p) } }
    var pid: pid_t = 0
    let result = argv.withUnsafeBufferPointer {
        posix_spawn(&pid, executable, &actions, &attrs, $0.baseAddress!, environ)
    }
    guard result == 0 else { throw ToolError("launch \(executable): errno \(result)") }
    return pid
}

// Foundation's read(upToCount:) can wait to fill its buffer on a socket or
// pipe. A single POSIX read returns the available bytes without waiting for EOF.
func readChromeChunk(_ handle: FileHandle) throws -> Data? {
    var bytes = [UInt8](repeating: 0, count: 64 << 10)
    while true {
        let count = Darwin.read(handle.fileDescriptor, &bytes, bytes.count)
        if count > 0 { return Data(bytes.prefix(count)) }
        if count == 0 { return nil }
        if errno != EINTR { throw ToolError("read Chrome connection: errno \(errno)") }
    }
}

private func chromeFrame(_ handle: FileHandle) throws -> Data {
    var data = Data()
    // One request per connection. Chunked reads avoid per-byte system calls.
    while true {
        guard let chunk = try readChromeChunk(handle), !chunk.isEmpty else {
            throw ToolError("Chrome's private connection closed")
        }
        if let end = chunk.firstIndex(of: 0) {
            data.append(chunk[..<end])
            guard data.count <= 256 << 20 else { throw ToolError("Chrome's reply is too large") }
            return data
        }
        data.append(chunk)
        guard data.count <= 256 << 20 else { throw ToolError("Chrome's reply is too large") }
    }
}

private let chromeClientID = UUID().uuidString

private actor ChromeContexts {
    var clients: [String: ChromeContext] = [:]

    func get(_ id: String, pid: pid_t) async -> ChromeContext {
        clients = clients.filter { _, context in
            kill(context.pid, 0) == 0 || errno != ESRCH
        }
        let context = clients[id] ?? ChromeContext(pid: pid)
        clients[id] = context
        // A departed helper's target sessions should stop collecting events.
        if let connection = await chromeLink.conn, let (chrome, _) = try? await connection.value {
            let detached = chrome.lock.withLock {
                let ids = chrome.sessions.filter { $0.value.context == nil }.map { $0.key }
                for id in ids { chrome.sessions.removeValue(forKey: id) }
                return ids
            }
            for id in detached {
                _ = try? await chrome.send("Target.detachFromTarget", ["sessionId": id])
            }
        }
        return context
    }
}

private let chromeContexts = ChromeContexts()
let chromeBrokerCalls = DispatchGroup()

func sharedChromeTool(_ args: [String: Any]) async throws -> [Content] {
    try await Task.detached {
        var fd = try? connectChromeBroker()
        if fd == nil {
            let pid = try spawnChromeProcess(Bundle.main.executablePath!, ["chrome-broker"])
            Thread.detachNewThread {
                var status: Int32 = 0
                while waitpid(pid, &status, 0) < 0, errno == EINTR {}
            }
            let start = Date()
            while fd == nil, Date().timeIntervalSince(start) < 15 {
                await pause(100)
                fd = try? connectChromeBroker()
            }
        }
        guard let fd else { throw ToolError("Chrome's private helper did not start within 15s") }
        let handle = FileHandle(fileDescriptor: fd, closeOnDealloc: true)
        defer { try? handle.close() }
        var request: [String: Any] = ["arguments": args, "client": chromeClientID]
        request["media_dir"] = mediaDir
        var data = try JSONSerialization.data(withJSONObject: request)
        data.append(0)
        try handle.write(contentsOf: data)
        let response = try JSONSerialization.jsonObject(with: chromeFrame(handle)) as? [String: Any] ?? [:]
        if let error = response.string("error") { throw ToolError(error) }
        return (response["content"] as? [[String: Any]] ?? []).compactMap { item in
            if item.string("type") == "image", let data = Data(base64Encoded: item.string("data") ?? "") {
                return Content.png(data)
            }
            return Content.text(item.string("text") ?? "")
        }
    }.value
}

func serveChromeBroker() throws {
    let fm = FileManager.default
    if !fm.fileExists(atPath: chromeSocketDir) {
        do {
            try fm.createDirectory(atPath: chromeSocketDir, withIntermediateDirectories: false, attributes: [.posixPermissions: 0o700])
        } catch {
            if !fm.fileExists(atPath: chromeSocketDir) { throw error }
        }
    }
    let attributes = try fm.attributesOfItem(atPath: chromeSocketDir)
    guard attributes[.type] as? FileAttributeType == .typeDirectory,
          (attributes[.ownerAccountID] as? NSNumber)?.uint32Value == getuid() else {
        throw ToolError("Chrome's socket directory is not owned by this account")
    }
    try fm.setAttributes([.posixPermissions: 0o700], ofItemAtPath: chromeSocketDir)
    let lock = Darwin.open(chromeSocketDir + "/lock", O_CREAT | O_RDWR | O_NOFOLLOW | O_CLOEXEC, 0o600)
    guard lock >= 0 else { throw ToolError("open Chrome helper lock: errno \(errno)") }
    defer { close(lock) }
    guard flock(lock, LOCK_EX | LOCK_NB) == 0 else { return }
    unlink(chromeSocketPath)
    let fd = socket(AF_UNIX, SOCK_STREAM, 0)
    guard fd >= 0 else { throw ToolError("create Chrome listener: errno \(errno)") }
    defer {
        close(fd)
        unlink(chromeSocketPath)
    }
    var address = try chromeAddress()
    let bound = withUnsafePointer(to: &address) {
        $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
            bind(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
        }
    }
    guard bound == 0, listen(fd, 16) == 0 else { throw ToolError("listen on Chrome socket: errno \(errno)") }
    while true {
        let client = accept(fd, nil, nil)
        if client < 0 {
            if errno == EINTR { continue }
            throw ToolError("accept Chrome connection: errno \(errno)")
        }
        // Check the peer too: filesystem permissions alone are not the boundary.
        var uid: uid_t = 0
        var gid: gid_t = 0
        guard getpeereid(client, &uid, &gid) == 0, uid == getuid() else {
            close(client)
            continue
        }
        var pid: pid_t = 0
        var pidSize = socklen_t(MemoryLayout<pid_t>.size)
        guard getsockopt(client, SOL_LOCAL, LOCAL_PEERPID, &pid, &pidSize) == 0 else {
            close(client)
            continue
        }
        var noSignal: Int32 = 1
        setsockopt(client, SOL_SOCKET, SO_NOSIGPIPE, &noSignal, socklen_t(MemoryLayout<Int32>.size))
        chromeBrokerCalls.enter()
        Task.detached {
            defer { chromeBrokerCalls.leave() }
            let handle = FileHandle(fileDescriptor: client, closeOnDealloc: true)
            defer { try? handle.close() }
            var response: [String: Any]
            do {
                let request = try JSONSerialization.jsonObject(with: chromeFrame(handle)) as? [String: Any] ?? [:]
                let context = await chromeContexts.get(request.string("client") ?? "", pid: pid)
                let content = try await chromeTool(request.dict("arguments") ?? [:], saveDir: request.string("media_dir"), context: context)
                response = ["content": content.map { $0.json }]
            } catch let error as ToolError {
                response = ["error": error.message]
            } catch {
                response = ["error": "\(error)"]
            }
            if var data = try? JSONSerialization.data(withJSONObject: response) {
                data.append(0)
                try? handle.write(contentsOf: data)
            }
        }
    }
}