//
// Copyright 2023 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::SocketAddr;

use clap::{Parser, ValueEnum};
use oak_attestation_gcp::attestation::request_attestation_token;
use oak_functions_service::wasm::wasmtime::WasmtimeHandler;
use oak_functions_standalone::{serve, AttestationArgs, OakFunctionsSessionArgs};
use oak_proto_rust::oak::functions::{
    config::ApplicationConfig, InitializeRequest, LookupDataChunk,
};
use oak_session::attestation::AttestationType;
use p256::ecdsa::{signature::rand_core::OsRng, SigningKey};
use prost::Message;
use sha2::Digest;
use tokio::net::TcpListener;
use tokio_stream::wrappers::TcpListenerStream;

#[global_allocator]
static ALLOCATOR: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

// Supported AttestationTypes for Google Cloud Platform, derived from
// oak/oak_session/src/attestation.rs.
#[derive(Debug, Clone, Eq, PartialEq, ValueEnum)]
enum AttestationTypeParam {
    Unattested,
    SelfUnidirectional,
}

// For both the Wasm and lookup data arguments, either their paths or URLs must
// be provided. If both are provided in either case, the program will panic.
// Note that all path arguements must be built into the container to ensure that
// they are accessible at runtime.
#[derive(Parser, Debug)]
struct Args {
    // TODO: b/424407998 - Have wasm_path point to content addressable storage
    #[arg(long, help = "The URI for fetching the wasm logic")]
    wasm_uri: String,

    #[arg(
        long,
        help = "The URI for fetching the serialized LookupDataChunk data",
        hide_short_help = true,
        default_value = ""
    )]
    lookup_data_uri: String,

    // Peer attestation verification is currently unsuported.
    #[arg(
        short,
        long,
        value_enum,
        default_value_t = AttestationTypeParam::Unattested,
        help = "The attestation scheme for the server. By default, no attestation is generated",
        hide_short_help = true,
    )]
    attestation_type: AttestationTypeParam,

    // Address to listen on
    #[arg(long, help = "The port to listen on", default_value = "0.0.0.0:8080")]
    listen_address: SocketAddr,
}

fn fetch_data_from_uri(uri: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    println!("fetching data from uri: {uri}");
    let response = ureq::get(uri).call()?;
    let mut buffer = Vec::new();
    response.into_reader().read_to_end(&mut buffer)?;
    Ok(buffer)
}

// Constructs the attestation arguement for the Oak Functions server.
// Attestations generated by the server assume the workload is running on Google
// Cloud's Confidential Space environment. If the server is not running in such
// an environment, only `AttestationTypeParam::Unattested` should be used.
// For details on Google Cloud's attestion scheme, read:
// https://cloud.google.com/confidential-computing/confidential-space/docs/reference/token-claims
fn create_attestation_args_for_gcp(
    attestation_type_param: AttestationTypeParam,
) -> AttestationArgs {
    let attestation_type = match attestation_type_param {
        AttestationTypeParam::SelfUnidirectional => AttestationType::SelfUnidirectional,
        AttestationTypeParam::Unattested => AttestationType::Unattested,
    };

    if attestation_type_param == AttestationTypeParam::Unattested {
        return AttestationArgs { attestation_type, binding_key: None, endorsement: None };
    }

    // Create the binding key and get the attestation token endorsement.
    let binding_key = SigningKey::random(&mut OsRng);
    let public_key_hash = sha2::Sha256::digest(binding_key.verifying_key().to_sec1_bytes());
    let public_key_hash = hex::encode(public_key_hash);

    println!("Requesting attestation token for {public_key_hash}...");
    let endorsement =
        request_attestation_token("oak://session/attestation", public_key_hash.as_str()).unwrap();
    AttestationArgs {
        attestation_type,
        binding_key: Some(binding_key),
        endorsement: Some(endorsement),
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    let wasm_module_bytes: Vec<u8> = if !args.wasm_uri.is_empty() {
        fetch_data_from_uri(&args.wasm_uri).expect("unable to fetch Wasm data")
    } else {
        panic!("--wasm_uri must be specified")
    };

    let lookup_data_option: Option<LookupDataChunk> = if !args.lookup_data_uri.is_empty() {
        let uri = &args.lookup_data_uri;
        // Issues an HTTP GET request to fetch the lookup data
        println!("reading LookupDataChunk from: {}", uri);
        let lookup_data_bytes = fetch_data_from_uri(uri).expect("unable to fetch lookup data");
        Some(LookupDataChunk::decode(lookup_data_bytes.as_slice()).unwrap())
    } else {
        None
    };

    let attestation_args = create_attestation_args_for_gcp(args.attestation_type);

    // This is a hack to get _some_ logging out of the binary, and should be
    // replaced with proper OTLP logging (or logging to journald, or something) in
    // the not too distant future. Debug logging is also only enabled for the
    // `oak_functions_service` module as Tonic tends to be rather chatty if
    // you enable debug logs everywhere; also, this could end up in a feedback
    // loop as if we create a RPC do do the debug logging, it'll mean the RPC
    // itself will generate more debug logs, which in turn will be sent via a
    // RPC, and the cycle continues.
    env_logger::builder()
        .filter_module("oak_functions_service", log::LevelFilter::Debug)
        .try_init()?;

    let application_config = ApplicationConfig::default();
    let wasmtime_config = application_config.wasmtime_config.unwrap_or_default();

    let addr = args.listen_address;

    let oak_functions_session_args = OakFunctionsSessionArgs {
        wasm_initialization: InitializeRequest {
            constant_response_size: 100, // This value is ultimately ignored.
            wasm_module: wasm_module_bytes,
        },
        attestation_args,
        lookup_data: lookup_data_option,
    };

    let server_handle = {
        let listener = TcpListener::bind(addr).await?;
        serve::<WasmtimeHandler>(
            Box::new(TcpListenerStream::new(listener)),
            wasmtime_config,
            oak_functions_session_args,
        )
    };

    println!("Listening on Address: {}", addr);

    Ok(server_handle.await?)
}
