use std::{
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use futures::StreamExt;
use rs_zero::{
    discovery::{InstanceEndpoint, MemoryRegistry, Registry, ServiceInstance},
    rpc::{
        LoadBalancePolicy, RpcClientConfig, RpcRetryConfig, RpcStreamingConfig,
        balancer::WeightedRoundRobinBalancer,
        deadline::{DeadlineBudget, insert_grpc_timeout, remaining_timeout_from_metadata},
        deadline_interceptor,
        discovery::resolve_service_endpoint,
        endpoint_from_rpc_endpoint,
        retry::{is_retryable_status, next_backoff, run_with_retry, run_with_retry_budget},
        streaming::{
            ObservedRecvStream, RpcStreamingObserver, record_stream_send, run_observed_stream,
        },
    },
};
use tonic::{Code, Request, Status, service::Interceptor};

#[test]
fn rpc_go_zero_defaults_enable_production_chain() {
    let config = RpcClientConfig::go_zero_defaults("http://127.0.0.1:50051");

    assert!(config.retry.enabled);
    assert!(config.deadline.propagate);
    assert!(config.streaming.observe);
    assert_eq!(
        config.load_balance.policy,
        LoadBalancePolicy::WeightedRoundRobin
    );
}

#[tokio::test]
async fn retry_runs_retryable_status_until_success() {
    let attempts = Arc::new(AtomicUsize::new(0));
    let seen = attempts.clone();
    let config = RpcRetryConfig {
        enabled: true,
        max_attempts: 3,
        initial_backoff: Duration::from_millis(1),
        max_backoff: Duration::from_millis(2),
        ..RpcRetryConfig::default()
    };

    let result = run_with_retry(&config, || {
        let seen = seen.clone();
        async move {
            let attempt = seen.fetch_add(1, Ordering::SeqCst);
            if attempt == 0 {
                Err(Status::unavailable("try again"))
            } else {
                Ok("ok")
            }
        }
    })
    .await;

    assert_eq!(result.expect("retry succeeds"), "ok");
    assert_eq!(attempts.load(Ordering::SeqCst), 2);
    assert!(is_retryable_status(
        &Status::deadline_exceeded("late"),
        &config
    ));
    assert_eq!(next_backoff(&config, 8), Duration::from_millis(2));
}

#[test]
fn deadline_inserts_and_reads_grpc_timeout_metadata() {
    let budget = DeadlineBudget::new(Duration::from_millis(150));
    let mut request = Request::new(());

    insert_grpc_timeout(&mut request, budget.remaining()).expect("metadata");

    let remaining = remaining_timeout_from_metadata(request.metadata()).expect("timeout");
    assert!(remaining <= Duration::from_millis(150));
    assert!(remaining >= Duration::from_millis(1));

    let mut interceptor = deadline_interceptor(Duration::from_millis(25));
    let intercepted = interceptor.call(Request::new(())).expect("intercepted");
    assert_eq!(
        remaining_timeout_from_metadata(intercepted.metadata()),
        Some(Duration::from_millis(25))
    );
}

#[tokio::test]
async fn retry_budget_passes_remaining_timeout_to_attempts() {
    let config = RpcRetryConfig {
        enabled: true,
        max_attempts: 2,
        initial_backoff: Duration::from_millis(1),
        max_backoff: Duration::from_millis(1),
        ..RpcRetryConfig::default()
    };
    let budget = DeadlineBudget::new(Duration::from_millis(50));
    let attempts = Arc::new(AtomicUsize::new(0));
    let seen = attempts.clone();

    let result = run_with_retry_budget(&config, &budget, |remaining| {
        let seen = seen.clone();
        async move {
            assert!(remaining <= Duration::from_millis(50));
            if seen.fetch_add(1, Ordering::SeqCst) == 0 {
                Err(Status::unavailable("try again"))
            } else {
                Ok(())
            }
        }
    })
    .await;

    result.expect("budget retry succeeds");
    assert_eq!(attempts.load(Ordering::SeqCst), 2);
}

#[tokio::test]
async fn discovery_resolves_weighted_rpc_endpoint() {
    let registry = MemoryRegistry::new();
    registry
        .register(
            ServiceInstance::new(
                "greeter",
                "a",
                InstanceEndpoint::new("127.0.0.1", 50051).expect("endpoint"),
            )
            .with_weight(2),
        )
        .await
        .expect("register a");
    registry
        .register(ServiceInstance::new(
            "greeter",
            "b",
            InstanceEndpoint::new("127.0.0.1", 50052).expect("endpoint"),
        ))
        .await
        .expect("register b");

    let selector = WeightedRoundRobinBalancer::new();
    let first = resolve_service_endpoint(&registry, &selector, "greeter")
        .await
        .expect("first");
    let second = resolve_service_endpoint(&registry, &selector, "greeter")
        .await
        .expect("second");

    assert!(first.uri.starts_with("http://127.0.0.1:"));
    assert_eq!(first.service, "greeter");
    assert_eq!(first.instance_id, "a");
    assert_eq!(second.instance_id, "a");
    assert!(endpoint_from_rpc_endpoint(&first, &RpcClientConfig::new(&first.uri)).is_ok());
}

#[tokio::test]
async fn streaming_observer_tracks_messages_and_status() {
    let observer =
        RpcStreamingObserver::new("chat.Chat", "Talk", RpcStreamingConfig::go_zero_defaults());

    observer.record_send().await;
    observer.record_recv().await;
    observer
        .finish::<()>(Err(Status::new(Code::Unavailable, "closed")))
        .await;

    let snapshot = observer.snapshot().await;
    assert_eq!(snapshot.sent_messages, 1);
    assert_eq!(snapshot.received_messages, 1);
    assert_eq!(snapshot.code, Some(Code::Unavailable));
}

#[tokio::test]
async fn streaming_wrapper_records_receive_and_finish() {
    let observer =
        RpcStreamingObserver::new("chat.Chat", "Watch", RpcStreamingConfig::go_zero_defaults());
    let stream = futures::stream::iter(vec![Ok::<_, Status>(1), Ok(2)]);
    let mut observed = ObservedRecvStream::new(Box::pin(stream), observer.clone());

    assert_eq!(observed.next().await.expect("first").expect("ok"), 1);
    assert_eq!(observed.next().await.expect("second").expect("ok"), 2);
    assert!(observed.next().await.is_none());
    tokio::task::yield_now().await;

    let snapshot = observer.snapshot().await;
    assert_eq!(snapshot.received_messages, 2);
    assert!(snapshot.completed);
    assert_eq!(snapshot.code, Some(Code::Ok));
}

#[tokio::test]
async fn streaming_send_wrapper_records_errors() {
    let observer =
        RpcStreamingObserver::new("chat.Chat", "Talk", RpcStreamingConfig::go_zero_defaults());

    let error = record_stream_send(&observer, async {
        Err::<(), _>(Status::unavailable("down"))
    })
    .await
    .expect_err("send error");

    assert_eq!(error.code(), Code::Unavailable);
    let snapshot = observer.snapshot().await;
    assert!(snapshot.completed);
    assert_eq!(snapshot.code, Some(Code::Unavailable));
}

#[tokio::test]
async fn observed_stream_runner_maps_timeout() {
    let config = RpcStreamingConfig {
        timeout: Some(Duration::from_millis(1)),
        ..RpcStreamingConfig::go_zero_defaults()
    };
    let observer = RpcStreamingObserver::new("chat.Chat", "Talk", config.clone());

    let error = run_observed_stream(&observer, &config, async {
        tokio::time::sleep(Duration::from_millis(20)).await;
        Ok::<_, Status>(())
    })
    .await
    .expect_err("timeout");

    assert_eq!(error.code(), Code::DeadlineExceeded);
    let snapshot = observer.snapshot().await;
    assert!(snapshot.completed);
    assert_eq!(snapshot.code, Some(Code::DeadlineExceeded));
}
