firestore 0.61.1

use crate::db::endpoint::FirestoreEmulatorCredentials;
use crate::{FirestoreDb, FirestoreDbOptions};
#[cfg(feature = "admin")]
use gcloud_sdk::google::firestore::admin::v1::{
    BulkDeleteDocumentsMetadata, Field as ProtoField, Index as ProtoIndex,
};
use gcloud_sdk::google::firestore::v1::{
    get_document_request, BeginTransactionResponse, CommitResponse, GetDocumentRequest,
};
#[cfg(feature = "admin")]
use gcloud_sdk::google::longrunning::{operation, Operation};
#[cfg(feature = "admin")]
use gcloud_sdk::google::rpc::Status as RpcStatus;
use gcloud_sdk::prost::Message as _;
use gcloud_sdk::tonic::Code;
use h2::server::SendResponse;
use h2::RecvStream;
use hyper::body::Bytes;
use hyper::header::{HeaderValue, AUTHORIZATION};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{watch, Notify};

/// A handler's answer to one RPC.
pub(super) enum FakeResponse {
    /// A framed protobuf message, sent with a `grpc-status: 0` trailer.
    Message(Vec<u8>),
    /// A stream of framed messages, the first at once and each next one the given interval after
    /// the one before, then a `grpc-status: 0` trailer: a stream that is slow but healthy.
    Paced(Vec<Vec<u8>>, Duration),
    /// A Trailers-Only error response with no body - the only way a streaming call's error is
    /// visible at its initial `await` rather than only once polled. Never `Code::Ok`: a unary
    /// call answered that way fails for want of a response message; use [`FakeResponse::empty`].
    Status(Code),
    /// Like [`FakeResponse::Status`], with a `grpc-message` trailer a test can assert the client
    /// surfaced, rather than the empty message every bare `Status(code)` carries.
    #[cfg(feature = "admin")]
    StatusWithMessage(Code, &'static str),
    /// The server closes the whole connection without answering, simulating a response lost in
    /// transit: the client sees a transport error, not a `grpc-status` it can read off the wire.
    Drop,
    /// The server never answers and keeps the stream open, simulating a call stuck in flight
    /// until the client gives up on it.
    Hang,
}

impl FakeResponse {
    /// A successful answer carrying an empty message, such as `Rollback`'s `Empty`.
    pub fn empty() -> Self {
        Self::Message(Vec::new())
    }

    /// A successful `Commit` with no write results.
    pub fn committed() -> Self {
        Self::Message(CommitResponse::default().encode_to_vec())
    }
}

/// Answers `BeginTransaction` with the next sequential transaction ID (1, 2, 3, ...), logged as
/// `Begin→<id>` - the one RPC every transaction fixture in this crate answers the same way.
pub(super) fn begin_response(next_id: &AtomicU8) -> (String, FakeResponse) {
    let id = next_id.fetch_add(1, Ordering::SeqCst) + 1;
    let transaction = BeginTransactionResponse {
        transaction: vec![id],
    };
    (
        format!("Begin→{id}"),
        FakeResponse::Message(transaction.encode_to_vec()),
    )
}

/// Answers `BeginTransaction` with `code` instead, logged as `Begin failed: <code>`. No
/// transaction ID is used up, so the next [`begin_response`] opens the one this would have.
pub(super) fn begin_failure(code: Code) -> (String, FakeResponse) {
    (
        format!("Begin failed: {code:?}"),
        FakeResponse::Status(code),
    )
}

/// The transaction ID, as [`begin_response`] numbers them, that a `GetDocument` request reads in.
///
/// Panics if the read does not belong to a transaction.
pub(super) fn read_transaction_id(get_document_request: &[u8]) -> u8 {
    let request = GetDocumentRequest::decode(get_document_request).unwrap();
    let Some(get_document_request::ConsistencySelector::Transaction(id)) =
        request.consistency_selector
    else {
        panic!("read must belong to a transaction");
    };
    id[0]
}

/// A `ListIndexes`/`ListFields` response, or a `CreateIndex`/`UpdateField` operation, or a
/// `DeleteIndex`/`Empty` acknowledgement - the admin and long-running operations RPCs an index
/// sync test drives, alongside the data RPC fixtures above.
#[cfg(feature = "admin")]
pub(super) fn list_indexes_response(indexes: Vec<ProtoIndex>) -> FakeResponse {
    FakeResponse::Message(
        gcloud_sdk::google::firestore::admin::v1::ListIndexesResponse {
            indexes,
            next_page_token: String::new(),
        }
        .encode_to_vec(),
    )
}

#[cfg(feature = "admin")]
pub(super) fn list_fields_response(fields: Vec<ProtoField>) -> FakeResponse {
    FakeResponse::Message(
        gcloud_sdk::google::firestore::admin::v1::ListFieldsResponse {
            fields,
            next_page_token: String::new(),
        }
        .encode_to_vec(),
    )
}

/// A long-running operation still in progress.
#[cfg(feature = "admin")]
pub(super) fn pending_operation_response(name: &str) -> FakeResponse {
    FakeResponse::Message(
        Operation {
            name: name.to_string(),
            metadata: None,
            done: false,
            result: None,
        }
        .encode_to_vec(),
    )
}

/// A long-running operation that completed successfully.
#[cfg(feature = "admin")]
pub(super) fn done_operation_response(name: &str) -> FakeResponse {
    FakeResponse::Message(
        Operation {
            name: name.to_string(),
            metadata: None,
            done: true,
            result: Some(operation::Result::Response(
                gcloud_sdk::prost_types::Any::default(),
            )),
        }
        .encode_to_vec(),
    )
}

/// A `BulkDeleteDocuments` long-running operation carrying `metadata`, pending or done - the
/// fixture a bulk-delete wait test needs that [`pending_operation_response`]/
/// [`done_operation_response`] cannot provide, since those never attach metadata.
#[cfg(feature = "admin")]
pub(super) fn bulk_delete_operation_response(
    name: &str,
    done: bool,
    metadata: BulkDeleteDocumentsMetadata,
) -> FakeResponse {
    FakeResponse::Message(
        Operation {
            name: name.to_string(),
            metadata: Some(gcloud_sdk::prost_types::Any {
                type_url:
                    "type.googleapis.com/google.firestore.admin.v1.BulkDeleteDocumentsMetadata"
                        .to_string(),
                value: metadata.encode_to_vec(),
            }),
            done,
            result: done
                .then(|| operation::Result::Response(gcloud_sdk::prost_types::Any::default())),
        }
        .encode_to_vec(),
    )
}

/// A long-running operation that completed with an error.
#[cfg(feature = "admin")]
pub(super) fn failed_operation_response(name: &str, code: i32, message: &str) -> FakeResponse {
    FakeResponse::Message(
        Operation {
            name: name.to_string(),
            metadata: None,
            done: true,
            result: Some(operation::Result::Error(RpcStatus {
                code,
                message: message.to_string(),
                details: vec![],
            })),
        }
        .encode_to_vec(),
    )
}

type Handler = dyn Fn(&str, &[u8]) -> (String, FakeResponse) + Send + Sync;

/// A minimal fake Firestore gRPC server for driving a real `FirestoreDb` against a controlled
/// backend. `handler` runs once per RPC, after its request body has ended, and returns one line
/// for `calls()` plus the response - so a test can assert a whole RPC sequence in one comparison.
///
/// `Listen` is the exception: it is answered at once with an empty stream, or with an open and
/// quiet one by [`start_with_quiet_listen`](Self::start_with_quiet_listen), and its handler runs
/// only when the client closes the request, so a logged `Listen` call is itself the proof that
/// the request was closed. The handler's response to it is ignored.
pub(super) struct FakeFirestore {
    pub db: FirestoreDb,
    address: SocketAddr,
    log: RequestLog,
}

impl FakeFirestore {
    pub async fn start<F>(handler: F) -> Self
    where
        F: Fn(&str, &[u8]) -> (String, FakeResponse) + Send + Sync + 'static,
    {
        Self::serve(3, false, handler).await
    }

    /// Like [`start`](Self::start), but with a `max_retries` other than the client default -
    /// enough for a test to bound a persistent failure to a small, fast number of attempts.
    pub async fn start_with_max_retries<F>(max_retries: usize, handler: F) -> Self
    where
        F: Fn(&str, &[u8]) -> (String, FakeResponse) + Send + Sync + 'static,
    {
        Self::serve(max_retries, false, handler).await
    }

    /// Like [`start`](Self::start), but `Listen` is answered with the response headers and then
    /// nothing until the client closes the request: a listener on a database where nothing
    /// changes.
    pub async fn start_with_quiet_listen<F>(handler: F) -> Self
    where
        F: Fn(&str, &[u8]) -> (String, FakeResponse) + Send + Sync + 'static,
    {
        Self::serve(3, true, handler).await
    }

    async fn serve<F>(max_retries: usize, quiet_listen: bool, handler: F) -> Self
    where
        F: Fn(&str, &[u8]) -> (String, FakeResponse) + Send + Sync + 'static,
    {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();
        let handler: Arc<Handler> = Arc::new(handler);
        let log = RequestLog {
            calls: Arc::new(watch::Sender::new(Vec::new())),
            authorizations: Arc::new(Mutex::new(Vec::new())),
        };
        let served = log.clone();
        // Detached: #[tokio::test] drops every spawned task, this one included, at test end.
        tokio::spawn(async move {
            while let Ok((socket, _)) = listener.accept().await {
                tokio::spawn(serve_connection(
                    socket,
                    handler.clone(),
                    served.clone(),
                    quiet_listen,
                ));
            }
        });
        let mut options = FirestoreDbOptions::new("fake-firestore".into());
        options.firebase_api_url = Some(format!("http://{address}"));
        options.max_retries = max_retries;
        let db = FirestoreDb::with_options_auth(options, FirestoreEmulatorCredentials::new())
            .await
            .unwrap();
        Self { db, address, log }
    }

    /// The address the server listens on, as `FIRESTORE_EMULATOR_HOST` would name it.
    pub fn address(&self) -> SocketAddr {
        self.address
    }

    /// Every RPC handled so far, in the order the server received it.
    pub fn calls(&self) -> Vec<String> {
        self.log.calls.borrow().clone()
    }

    /// The `authorization` header of every RPC received so far, in the order the server
    /// received them, `None` for an RPC that carried none.
    pub fn authorizations(&self) -> Vec<Option<HeaderValue>> {
        self.log.authorizations.lock().unwrap().clone()
    }

    /// Resolves once at least `count` RPCs have been handled.
    pub async fn wait_for_calls(&self, count: usize) {
        self.log
            .calls
            .subscribe()
            .wait_for(|calls| calls.len() >= count)
            .await
            .expect("the call log lives as long as the server that owns it");
    }
}

/// What the server records of every RPC, shared by all its connections.
#[derive(Clone)]
struct RequestLog {
    calls: Arc<watch::Sender<Vec<String>>>,
    authorizations: Arc<Mutex<Vec<Option<HeaderValue>>>>,
}

/// Serves one client connection until the client closes it or an RPC answers
/// [`FakeResponse::Drop`], which drops the connection and its socket.
async fn serve_connection(
    socket: TcpStream,
    handler: Arc<Handler>,
    log: RequestLog,
    quiet_listen: bool,
) {
    let Ok(mut connection) = h2::server::handshake(socket).await else {
        return;
    };
    let close = Arc::new(Notify::new());
    let serve = async {
        while let Some(Ok((request, respond))) = connection.accept().await {
            tokio::spawn(answer(
                request,
                respond,
                handler.clone(),
                log.clone(),
                close.clone(),
                quiet_listen,
            ));
        }
    };
    tokio::select! {
        () = serve => {}
        () = close.notified() => {}
    }
}

async fn answer(
    request: hyper::Request<RecvStream>,
    mut respond: SendResponse<Bytes>,
    handler: Arc<Handler>,
    log: RequestLog,
    close: Arc<Notify>,
    quiet_listen: bool,
) {
    let method = request.uri().path().to_owned();
    log.authorizations
        .lock()
        .unwrap()
        .push(request.headers().get(AUTHORIZATION).cloned());
    let listening = method.ends_with("/Listen");
    // Held until the client closes the request, since dropping it would end a quiet `Listen`.
    let _listen_response = if listening {
        let headers = if quiet_listen {
            ok_headers()
        } else {
            ok_headers().header("grpc-status", "0")
        };
        let Ok(send) = respond.send_response(headers.body(()).unwrap(), !quiet_listen) else {
            return;
        };
        Some(send)
    } else {
        None
    };
    let mut body = request.into_body();
    let mut bytes = Vec::new();
    while let Some(Ok(chunk)) = body.data().await {
        body.flow_control().release_capacity(chunk.len()).unwrap();
        bytes.extend_from_slice(&chunk);
    }
    let (call, response) = handler(&method, bytes.get(5..).unwrap_or_default());
    log.calls.send_modify(|calls| calls.push(call));
    if listening {
        return;
    }
    match response {
        FakeResponse::Message(message) => {
            send_messages(respond, vec![message], Duration::ZERO).await
        }
        FakeResponse::Paced(messages, interval) => send_messages(respond, messages, interval).await,
        // A single Trailers-Only frame, with no body ever opened, is what makes even a streaming
        // call's error visible at its initial `await` rather than only once its caller polls the
        // response.
        FakeResponse::Status(code) => {
            assert_ne!(code, Code::Ok, "answer success with FakeResponse::empty");
            let headers = ok_headers().header("grpc-status", (code as i32).to_string());
            let _ = respond.send_response(headers.body(()).unwrap(), true);
        }
        #[cfg(feature = "admin")]
        FakeResponse::StatusWithMessage(code, message) => {
            assert_ne!(code, Code::Ok, "answer success with FakeResponse::empty");
            let headers = ok_headers()
                .header("grpc-status", (code as i32).to_string())
                .header("grpc-message", message);
            let _ = respond.send_response(headers.body(()).unwrap(), true);
        }
        FakeResponse::Hang => std::future::pending::<()>().await,
        FakeResponse::Drop => {
            close.notify_one();
            // Dropping `respond` while the connection is still up would reset just this stream,
            // which the client reads as CANCELLED rather than as a lost connection. The task ends
            // with the test's runtime.
            std::future::pending::<()>().await;
        }
    }
}

/// Answers with each of `messages`, `interval` apart and the first at once, then a
/// `grpc-status: 0` trailer.
async fn send_messages(
    mut respond: SendResponse<Bytes>,
    messages: Vec<Vec<u8>>,
    interval: Duration,
) {
    let Ok(mut send) = respond.send_response(ok_headers().body(()).unwrap(), false) else {
        return;
    };
    for (index, message) in messages.iter().enumerate() {
        if index > 0 {
            tokio::time::sleep(interval).await;
        }
        if send.send_data(grpc_frame(message), false).is_err() {
            return;
        }
    }
    let mut trailers = hyper::HeaderMap::new();
    trailers.insert("grpc-status", HeaderValue::from_static("0"));
    let _ = send.send_trailers(trailers);
}

/// `message` with the gRPC length prefix, uncompressed.
fn grpc_frame(message: &[u8]) -> Bytes {
    let mut frame = vec![0u8; 5];
    frame[1..5].copy_from_slice(&u32::try_from(message.len()).unwrap().to_be_bytes());
    frame.extend_from_slice(message);
    frame.into()
}

fn ok_headers() -> hyper::http::response::Builder {
    hyper::Response::builder()
        .status(200)
        .header("content-type", "application/grpc")
}