// Copyright (c) 2023 Huawei Device Co., Ltd.
// 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.

//! `Resolver` trait and `DefaultDnsResolver` implementation.

use std::future::Future;
use std::io;
use std::io::Error;
use std::net::{SocketAddr, ToSocketAddrs};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::vec::IntoIter;

use crate::runtime::JoinHandle;

/// `SocketAddr` resolved by `Resolver`.
pub type Addrs = Box<dyn Iterator<Item = SocketAddr> + Sync + Send>;
/// Possible errors that this resolver may generate when attempting to
/// resolve.
pub type StdError = Box<dyn std::error::Error + Send + Sync>;
/// Futures generated by this resolve when attempting to resolve an address.
pub type SocketFuture = Pin<Box<dyn Future<Output = Result<Addrs, StdError>> + Sync + Send>>;

/// `Resolver` trait used by `async_impl::connector::HttpConnector`. `Resolver`
/// provides asynchronous dns resolve interfaces.
pub trait Resolver: Send + Sync + 'static {
    /// resolve domain to a `SocketAddr` `Future`.
    fn resolve(&self, domain: &str) -> SocketFuture;
}

/// `SocketAddr` resolved by `DefaultDnsResolver`.
pub struct ResolvedAddrs {
    iter: IntoIter<SocketAddr>,
}

impl ResolvedAddrs {
    pub(super) fn new(iter: IntoIter<SocketAddr>) -> Self {
        Self { iter }
    }

    // The first ip in the dns record is the preferred addrs type.
    pub(super) fn split_preferred_addrs(self) -> (Vec<SocketAddr>, Vec<SocketAddr>) {
        // get preferred address family type.
        let is_ipv6 = self
            .iter
            .as_slice()
            .first()
            .map(SocketAddr::is_ipv6)
            .unwrap_or(false);
        self.iter
            .partition::<Vec<_>, _>(|addr| addr.is_ipv6() == is_ipv6)
    }
}

impl Iterator for ResolvedAddrs {
    type Item = SocketAddr;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next()
    }
}

/// Futures generated by `DefaultDnsResolver`.
pub struct DefaultDnsFuture {
    inner: JoinHandle<Result<ResolvedAddrs, Error>>,
}

impl Future for DefaultDnsFuture {
    type Output = Result<Addrs, StdError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.inner).poll(cx).map(|res| match res {
            Ok(Ok(addrs)) => Ok(Box::new(addrs) as Addrs),
            Ok(Err(err)) => Err(Box::new(err) as StdError),
            Err(err) => Err(Box::new(Error::new(io::ErrorKind::Interrupted, err)) as StdError),
        })
    }
}

/// Default dns resolver used by the `Client`。
#[derive(Default)]
pub struct DefaultDnsResolver {}

impl Resolver for DefaultDnsResolver {
    fn resolve(&self, domain: &str) -> SocketFuture {
        let domain: Box<str> = domain.into();
        let blocking = crate::runtime::spawn_blocking(move || {
            (*domain)
                .to_socket_addrs()
                .map(|iter| ResolvedAddrs { iter })
        });
        Box::pin(DefaultDnsFuture { inner: blocking })
    }
}