use async_graphql::{
    extensions::{Extension, ExtensionContext, ExtensionFactory, NextPrepareRequest}, http::GraphiQLSource, EmptyMutation, EmptySubscription, Object, Request, Schema, ServerResult
};
use async_graphql_axum::{GraphQLRequest, GraphQLResponse};
use axum::{
    async_trait, response::{self, IntoResponse}, routing::get, Router
};
use tokio::net::TcpListener;

pub type MySchema = Schema<Query, EmptyMutation, EmptySubscription>;

pub struct Query;

#[Object]
impl Query {
    async fn hello_world(&self) -> String {
        "Hello world!".to_string()
    }
}

async fn graphiql() -> impl IntoResponse {
    response::Html(GraphiQLSource::build().endpoint("/").finish())
}

struct MyData(String);

async fn graphql_handler(
    schema: axum::Extension<MySchema>,
    req: GraphQLRequest,
) -> GraphQLResponse {
    let mut req = req.into_inner();

    req = req.data(MyData("foobar".to_string()));

    schema.execute(req).await.into()
}

#[tokio::main]
async fn main() {
    let schema = Schema::build(Query, EmptyMutation, EmptySubscription)
        .extension(MyExtension::new())
        .finish();

    let app = Router::new().route("/", get(graphiql).post(graphql_handler))
    .layer(axum::Extension(schema));

    println!("GraphiQL IDE: http://localhost:8000");

    axum::serve(TcpListener::bind("127.0.0.1:8000").await.unwrap(), app)
        .await
        .unwrap();
}

pub struct MyExtension;

impl MyExtension {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {}
    }
}

impl ExtensionFactory for MyExtension {
    fn create(&self) -> Arc<dyn Extension> {
        Arc::new(MyExtensionFactory {})
    }
}

struct MyExtensionFactory;

#[async_trait]
impl Extension for MyExtensionFactory {
    async fn prepare_request(
        &self,
        ctx: &ExtensionContext<'_>,
        request: Request,
        next: NextPrepareRequest<'_>,
    ) -> ServerResult<Request> {
        /*
        With async-graphql and async-graphql-axum "=7.0.2" this works.

        With async-graphql and async-graphql-axum "=7.0.3" I get
        ```
        thread 'tokio-runtime-worker' panicked at src/main.rs:90:44:
        called `Result::unwrap()` on an `Err` value: Error { message: "Data `async_graphql_context_bug::MyData` does not exist.", extensions: None }
        ```
        */
        let my_data = ctx.data::<MyData>().unwrap();
        println!("{}", my_data.0);

        next.run(ctx, request).await
    }
}

[package]
name = "async_graphql_context_bug"
version = "0.1.0"
edition = "2021"

[dependencies]
async-graphql = "=7.0.3"
async-graphql-axum = "=7.0.3"
axum = { version = "0.7.5" }
tokio = { version = "1.32.0", features = ["macros", "rt-multi-thread"] }

thread 'tokio-runtime-worker' panicked at src/main.rs:90:44:
called `Result::unwrap()` on an `Err` value: Error { message: "Data `async_graphql_context_bug::MyData` does not exist.", extensions: None }