//! # Combined RFC 6186 SRV discovery coroutine
//!
//! [`DiscoverySrv`] runs the four RFC 6186 / RFC 8314 SRV queries
//! (`_imap._tcp.<domain>`, `_imaps._tcp.<domain>`,
//! `_submission._tcp.<domain>`, `_submissions._tcp.<domain>`) in
//! series, picks the best record per service (lowest priority, highest
//! weight on ties; already sorted by [`DiscoveryDnsSrv`]), and yields a
//! single [`DiscoverySrvReport`] when all four steps have completed.
//!
//! A per-service DNS failure (`InvalidQname`, `QueryTooLarge`,
//! `InvalidResponse`) terminates the whole coroutine with
//! [`DiscoverySrvError`]; empty SRV answers do not, the matching slot
//! is simply left as `None` in the report.

use core::mem;

use alloc::{
    format,
    string::{String, ToString},
};

use thiserror::Error;
use url::Url;

use crate::{
    coroutine::{DiscoveryCoroutine, DiscoveryCoroutineState, DiscoveryYield},
    rfc6186::{
        service::{DiscoverySrvReport, DiscoverySrvService},
        srv::{DiscoveryDnsSrv, DiscoveryDnsSrvError, SrvRecord},
    },
};

/// Errors emitted by [`DiscoverySrv`].
#[derive(Debug, Error)]
pub enum DiscoverySrvError {
    /// The `_imap._tcp` SRV lookup failed with a DNS-level error.
    #[error("DNS SRV lookup for `_imap._tcp` failed: {0}")]
    Imap(#[source] DiscoveryDnsSrvError),
    /// The `_imaps._tcp` SRV lookup failed with a DNS-level error.
    #[error("DNS SRV lookup for `_imaps._tcp` failed: {0}")]
    Imaps(#[source] DiscoveryDnsSrvError),
    /// The `_submission._tcp` SRV lookup failed with a DNS-level error.
    #[error("DNS SRV lookup for `_submission._tcp` failed: {0}")]
    Submission(#[source] DiscoveryDnsSrvError),
    /// The `_submissions._tcp` SRV lookup failed with a DNS-level error.
    #[error("DNS SRV lookup for `_submissions._tcp` failed: {0}")]
    Submissions(#[source] DiscoveryDnsSrvError),
}

#[derive(Default)]
enum State {
    Imap(DiscoveryDnsSrv),
    Imaps(DiscoveryDnsSrv),
    Submission(DiscoveryDnsSrv),
    Submissions(DiscoveryDnsSrv),
    #[default]
    Done,
}

/// I/O-free combined coroutine that runs the four RFC 6186 / RFC 8314
/// SRV queries and assembles their best records into a
/// [`DiscoverySrvReport`].
pub struct DiscoverySrv {
    state: State,
    domain: String,
    resolver: Url,
    report: DiscoverySrvReport,
}

impl DiscoverySrv {
    /// Builds the orchestrator. `resolver` must be a `tcp://host:port`
    /// URL pointing at a DNS-over-TCP resolver; it is yielded back on
    /// every `WantsRead` / `WantsWrite` so the runtime can route the
    /// bytes to the correct stream.
    pub fn new(domain: impl AsRef<str>, resolver: Url) -> Self {
        let domain = domain.as_ref().trim_matches('.').to_string();
        let imap = DiscoveryDnsSrv::new(format!("_imap._tcp.{domain}"), resolver.clone());

        Self {
            state: State::Imap(imap),
            domain,
            resolver,
            report: DiscoverySrvReport::default(),
        }
    }
}

impl DiscoveryCoroutine for DiscoverySrv {
    type Yield = DiscoveryYield;
    type Return = Result<DiscoverySrvReport, DiscoverySrvError>;

    fn resume(&mut self, arg: Option<&[u8]>) -> DiscoveryCoroutineState<Self::Yield, Self::Return> {
        match mem::take(&mut self.state) {
            State::Imap(mut srv) => match srv.resume(arg) {
                DiscoveryCoroutineState::Complete(Ok(records)) => {
                    self.report.imap = records.into_iter().next().map(into_service);
                    self.state = State::Imaps(DiscoveryDnsSrv::new(
                        format!("_imaps._tcp.{}", self.domain),
                        self.resolver.clone(),
                    ));
                    self.resume(None)
                }
                DiscoveryCoroutineState::Yielded(y) => {
                    self.state = State::Imap(srv);
                    DiscoveryCoroutineState::Yielded(y)
                }
                DiscoveryCoroutineState::Complete(Err(err)) => {
                    DiscoveryCoroutineState::Complete(Err(DiscoverySrvError::Imap(err)))
                }
            },
            State::Imaps(mut srv) => match srv.resume(arg) {
                DiscoveryCoroutineState::Complete(Ok(records)) => {
                    self.report.imaps = records.into_iter().next().map(into_service);
                    self.state = State::Submission(DiscoveryDnsSrv::new(
                        format!("_submission._tcp.{}", self.domain),
                        self.resolver.clone(),
                    ));
                    self.resume(None)
                }
                DiscoveryCoroutineState::Yielded(y) => {
                    self.state = State::Imaps(srv);
                    DiscoveryCoroutineState::Yielded(y)
                }
                DiscoveryCoroutineState::Complete(Err(err)) => {
                    DiscoveryCoroutineState::Complete(Err(DiscoverySrvError::Imaps(err)))
                }
            },
            State::Submission(mut srv) => match srv.resume(arg) {
                DiscoveryCoroutineState::Complete(Ok(records)) => {
                    self.report.submission = records.into_iter().next().map(into_service);
                    self.state = State::Submissions(DiscoveryDnsSrv::new(
                        format!("_submissions._tcp.{}", self.domain),
                        self.resolver.clone(),
                    ));
                    self.resume(None)
                }
                DiscoveryCoroutineState::Yielded(y) => {
                    self.state = State::Submission(srv);
                    DiscoveryCoroutineState::Yielded(y)
                }
                DiscoveryCoroutineState::Complete(Err(err)) => {
                    DiscoveryCoroutineState::Complete(Err(DiscoverySrvError::Submission(err)))
                }
            },
            State::Submissions(mut srv) => match srv.resume(arg) {
                DiscoveryCoroutineState::Complete(Ok(records)) => {
                    self.report.submissions = records.into_iter().next().map(into_service);
                    DiscoveryCoroutineState::Complete(Ok(mem::take(&mut self.report)))
                }
                DiscoveryCoroutineState::Yielded(y) => {
                    self.state = State::Submissions(srv);
                    DiscoveryCoroutineState::Yielded(y)
                }
                DiscoveryCoroutineState::Complete(Err(err)) => {
                    DiscoveryCoroutineState::Complete(Err(DiscoverySrvError::Submissions(err)))
                }
            },
            State::Done => panic!("DiscoverySrv::resume called after completion"),
        }
    }
}

fn into_service(record: SrvRecord) -> DiscoverySrvService {
    DiscoverySrvService {
        host: record
            .rdata
            .name
            .to_string()
            .trim_end_matches('.')
            .to_string(),
        port: record.rdata.port.get(),
        priority: record.rdata.priority.get(),
        weight: record.rdata.weight.get(),
    }
}
