import Foundation
import NucleicCore

/// Tier-2 contract-test stub (OBSERVABILITY B.4), the Grok analog of `fake-claude`: speaks Grok's
/// real **ACP** wire protocol (`grok agent stdio`) — standards JSON-RPC 2.0 over stdio — off a
/// recorded fixture. It answers the client's `initialize`/`session/new`/`session/load` handshake,
/// then on `session/prompt` plays a scripted turn: emits `session/update` notifications and issues
/// `session/request_permission` *requests*, asserting the client's reply matches the fixture's
/// expectation, before answering the prompt with a `stopReason`. Runs through the production
/// `ProcessHost` + `GrokACPBackend` with no `grok` installed and no cost.
///
/// Fixture (FAKE_GROK_FIXTURE → one JSON object):
///   {
///     "sessionId": "grok-sess-1", "model": "grok-4",
///     "agentCapabilities": {"loadSession": true},   // optional; default {loadSession:true}
///     "authMethods": [],                            // optional; default [] (no auth round-trip)
///     "stopReason": "end_turn",                     // optional; default "end_turn"
///     "turn": [
///        {"notify": {"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"hi"}}},
///        {"permission": {"toolCall": {...}, "expect": "allow"|"deny"|"cancel"}},
///        {"sleep": 30}
///     ]
///   }
///
/// Exit codes: 0 ok · 4 approval-reply mismatch · 5 setup/handshake failure.
@main
struct FakeGrok {
    static let peer = ACPPeer()

    static func main() async {
        guard let fixturePath = ProcessInfo.processInfo.environment["FAKE_GROK_FIXTURE"] else {
            die(5, "FAKE_GROK_FIXTURE not set")
        }
        guard let data = FileManager.default.contents(atPath: fixturePath),
            let fixture = try? JSONValue(parsing: data), fixture.objectValue != nil
        else {
            die(5, "cannot read/parse fixture \(fixturePath)")
        }

        // Drive the JSON-RPC channel: read client messages on stdin, dispatch requests, resolve
        // the replies to our own outbound requests.
        do {
            for try await line in FileHandle.standardInput.bytes.lines {
                guard let msg = try? JSONValue(parsing: Data(line.utf8)) else { continue }
                let method = msg["method"]?.stringValue
                let id = msg["id"]
                if let method, let id, id != .null {
                    await handleRequest(method: method, id: id, params: msg["params"] ?? .object([:]),
                        fixture: fixture)
                } else if method == nil, let intID = msg["id"]?.intValue {
                    await peer.resolve(id: intID, result: msg["result"] ?? .null)
                }
                // method without id → a notification (e.g. session/cancel): ignored by the stub.
            }
        } catch {
            die(5, "stdin read failed: \(error)")
        }
        exit(0)  // client closed stdin (teardown) → clean exit
    }

    // MARK: - Request handling

    static func handleRequest(method: String, id: JSONValue, params: JSONValue, fixture: JSONValue) async {
        switch method {
        case "initialize":
            let caps = fixture["agentCapabilities"] ?? .object(["loadSession": .bool(true)])
            let authMethods = fixture["authMethods"] ?? .array([])
            reply(id: id, result: .object([
                "protocolVersion": .number(1),
                "agentCapabilities": caps,
                "authMethods": authMethods,
                "agentInfo": .object(["name": .string("fake-grok"), "version": .string("0.1")]),
            ]))
        case "authenticate":
            reply(id: id, result: .object([:]))
        case "session/new":
            let sid = fixture["sessionId"]?.stringValue ?? "grok-sess-fake"
            reply(id: id, result: .object(["sessionId": .string(sid)]))
        case "session/load":
            reply(id: id, result: .object([:]))
        case "session/prompt":
            // Play the turn concurrently so the read loop keeps draining (to receive our
            // request_permission replies).
            let sid = params["sessionId"]?.stringValue
                ?? fixture["sessionId"]?.stringValue ?? "grok-sess-fake"
            Task { await playTurn(promptID: id, sessionID: sid, fixture: fixture) }
        default:
            replyError(id: id, code: -32601, message: "unknown method \(method)")
        }
    }

    // MARK: - Turn playback

    static func playTurn(promptID: JSONValue, sessionID: String, fixture: JSONValue) async {
        for step in fixture["turn"]?.arrayValue ?? [] {
            if let update = step["notify"] {
                emitNotification(method: "session/update", params: .object([
                    "sessionId": .string(sessionID), "update": update,
                ]))
            } else if let permission = step["permission"] {
                await performPermission(permission, sessionID: sessionID)
            } else if let ms = step["sleep"]?.intValue {
                try? await Task.sleep(for: .milliseconds(ms))
            }
        }
        let stopReason = fixture["stopReason"]?.stringValue ?? "end_turn"
        reply(id: promptID, result: .object(["stopReason": .string(stopReason)]))
    }

    /// Issue an ACP `session/request_permission` request and assert the client's reply matches the
    /// fixture's `expect` (`allow`→an allow_* option selected, `deny`→a reject_* option selected,
    /// `cancel`→the cancelled outcome).
    static func performPermission(_ permission: JSONValue, sessionID: String) async {
        let toolCall = permission["toolCall"] ?? .object([:])
        let options = JSONValue.array([
            .object(["optionId": .string("allow-once"), "name": .string("Allow once"),
                "kind": .string("allow_once")]),
            .object(["optionId": .string("allow-always"), "name": .string("Always allow"),
                "kind": .string("allow_always")]),
            .object(["optionId": .string("reject-once"), "name": .string("Reject"),
                "kind": .string("reject_once")]),
        ])
        let id = await peer.nextOutboundID()
        emitRequest(id: id, method: "session/request_permission", params: .object([
            "sessionId": .string(sessionID), "toolCall": toolCall, "options": options,
        ]))
        let result = await peer.awaitReply(id: id)

        let outcome = result["outcome"] ?? .object([:])
        let kind = outcome["outcome"]?.stringValue
        let optionID = outcome["optionId"]?.stringValue ?? ""
        switch permission["expect"]?.stringValue {
        case "allow":
            guard kind == "selected", optionID.hasPrefix("allow") else {
                die(4, "expected an allow selection, got \(result.canonicalString())")
            }
        case "deny":
            guard kind == "selected", optionID.hasPrefix("reject") else {
                die(4, "expected a reject selection, got \(result.canonicalString())")
            }
        case "cancel":
            guard kind == "cancelled" else {
                die(4, "expected a cancelled outcome, got \(result.canonicalString())")
            }
        default:
            break  // no expectation — accept whatever the client decided
        }
    }

    // MARK: - Wire I/O

    static func reply(id: JSONValue, result: JSONValue) {
        write(.object(["jsonrpc": .string("2.0"), "id": id, "result": result]))
    }

    static func replyError(id: JSONValue, code: Int, message: String) {
        write(.object([
            "jsonrpc": .string("2.0"), "id": id,
            "error": .object(["code": .number(Double(code)), "message": .string(message)]),
        ]))
    }

    static func emitNotification(method: String, params: JSONValue) {
        write(.object(["jsonrpc": .string("2.0"), "method": .string(method), "params": params]))
    }

    static func emitRequest(id: Int, method: String, params: JSONValue) {
        write(.object([
            "jsonrpc": .string("2.0"), "id": .number(Double(id)),
            "method": .string(method), "params": params,
        ]))
    }

    private static let writeLock = NSLock()
    static func write(_ value: JSONValue) {
        guard let data = try? value.encodedData() else { return }
        writeLock.lock()
        defer { writeLock.unlock() }
        FileHandle.standardOutput.write(data)
        FileHandle.standardOutput.write(Data([0x0A]))
    }

    static func die(_ code: Int32, _ message: String) -> Never {
        FileHandle.standardError.write(Data(("fake-grok: " + message + "\n").utf8))
        exit(code)
    }
}

/// Serializes the stub's outbound request ids and the continuations awaiting their replies.
actor ACPPeer {
    private var nextID = 0
    private var pending: [Int: CheckedContinuation<JSONValue, Never>] = [:]

    func nextOutboundID() -> Int { nextID += 1; return nextID }

    func awaitReply(id: Int) async -> JSONValue {
        await withCheckedContinuation { continuation in pending[id] = continuation }
    }

    func resolve(id: Int, result: JSONValue) {
        pending.removeValue(forKey: id)?.resume(returning: result)
    }
}
