pub mod events;
pub mod handshake;
pub mod path;
pub mod space;
pub mod state;
pub mod termination;
pub mod tls;
pub mod tx;

pub mod prelude {
    pub use qbase::{
        cid::ConnectionId,
        frame::ConnectionCloseFrame,
        net::route::*,
        sid::{ControlStreamsConcurrency, ProductStreamsConcurrencyController, StreamId},
        varint::VarInt,
    };
    pub use qinterface::{QuicInterface, router::QuicProto};
    #[cfg(feature = "unreliable")]
    pub use qunreliable::{DatagramReader, DatagramWriter};

    #[allow(unused_imports)]
    pub mod handy {
        pub use qbase::{param::handy::*, sid::handy::*, token::handy::*};
        pub use qinterface::handy::*;
    }

    pub use crate::{
        Connection, StreamReader, StreamWriter,
        events::{EmitEvent, Event},
        path::idle::HeartbeatConfig,
    };
}

pub mod builder;

use std::{
    borrow::Cow,
    future::Future,
    io,
    sync::{Arc, RwLock},
};

use enum_dispatch::enum_dispatch;
use events::ArcEventBroker;
use path::ArcPathContexts;
use prelude::HeartbeatConfig;
use qbase::{
    cid, flow,
    frame::{ConnectionCloseFrame, CryptoFrame, ReliableFrame, StreamFrame},
    net::route::{Link, Pathway},
    param::{ArcParameters, ParameterId},
    sid::StreamId,
    token::ArcTokenRegistry,
    varint::VarInt,
};
use qevent::telemetry::Instrument;
use qinterface::{
    queue::RcvdPacketQueue,
    router::{QuicProto, RouterRegistry},
};
use qrecovery::{
    journal, recv, reliable, send,
    streams::{self, Ext},
};
#[cfg(feature = "unreliable")]
use qunreliable::{DatagramReader, DatagramWriter};
use space::Spaces;
use state::ConnState;
use termination::Termination;
use tls::ArcTlsSession;
use tracing::Instrument as _;

/// The kind of frame which guaratend to be received by peer.
///
/// The bundle of [`StreamFrame`], [`CryptoFrame`] and [`ReliableFrame`].
#[derive(Debug, Clone, Eq, PartialEq)]
#[enum_dispatch(EncodeFrame, FrameFeture)]
pub enum GuaranteedFrame {
    Stream(StreamFrame),
    Crypto(CryptoFrame),
    Reliable(ReliableFrame),
}

/// For initial space, only reliable transmission of crypto frames is required.
pub type InitialJournal = journal::Journal<CryptoFrame>;
/// For handshake space, only reliable transmission of crypto frames is required.
pub type HandshakeJournal = journal::Journal<CryptoFrame>;
/// For data space, reliable transmission of [`GuaranteedFrame`] (crypto frames, stream frames and reliable frames) is required.
pub type DataJournal = journal::Journal<GuaranteedFrame>;

pub type ArcReliableFrameDeque = reliable::ArcReliableFrameDeque<ReliableFrame>;
pub type ArcLocalCids = cid::ArcLocalCids<RouterRegistry<ArcReliableFrameDeque>>;
pub type ArcRemoteCids = cid::ArcRemoteCids<ArcReliableFrameDeque>;
pub type CidRegistry = cid::Registry<ArcLocalCids, ArcRemoteCids>;
pub type ArcDcidCell = cid::ArcCidCell<ArcReliableFrameDeque>;

pub type FlowController = flow::FlowController<ArcReliableFrameDeque>;
pub type Credit<'a> = flow::Credit<'a, ArcReliableFrameDeque>;

pub type Handshake = handshake::Handshake<ArcReliableFrameDeque>;
pub type RawHandshake = handshake::RawHandshake<ArcReliableFrameDeque>;

pub type DataStreams = streams::DataStreams<ArcReliableFrameDeque>;
pub type StreamReader = recv::Reader<Ext<ArcReliableFrameDeque>>;
pub type StreamWriter = send::Writer<Ext<ArcReliableFrameDeque>>;

#[derive(Clone)]
pub struct Components {
    parameters: ArcParameters,
    tls_session: ArcTlsSession,
    handshake: Handshake,
    token_registry: ArcTokenRegistry,
    cid_registry: CidRegistry,
    flow_ctrl: FlowController,
    spaces: Spaces,
    paths: ArcPathContexts,
    proto: Arc<QuicProto>,
    rcvd_pkt_q: Arc<RcvdPacketQueue>,
    defer_idle_timeout: HeartbeatConfig,
    event_broker: ArcEventBroker,
    conn_state: ConnState,
}

impl Components {
    pub fn open_bi_stream(
        &self,
    ) -> impl Future<Output = io::Result<Option<(StreamId, (StreamReader, StreamWriter))>>> + Send
    {
        let params = self.parameters.clone();
        let streams = self.spaces.data().streams().clone();
        async move {
            let snd_wnd_size = params
                .get_remote_as::<VarInt>(ParameterId::InitialMaxStreamDataBidiRemote)
                .await?;
            Ok(streams.open_bi(snd_wnd_size.into_inner()).await?)
        }
        .instrument_in_current()
        .in_current_span()
    }

    pub fn open_uni_stream(
        &self,
    ) -> impl Future<Output = io::Result<Option<(StreamId, StreamWriter)>>> + Send {
        let params = self.parameters.clone();
        let streams = self.spaces.data().streams().clone();
        async move {
            let snd_wnd_size = params
                .get_remote_as::<VarInt>(ParameterId::InitialMaxStreamDataUni)
                .await?;
            Ok(streams.open_uni(snd_wnd_size.into_inner()).await?)
        }
        .instrument_in_current()
        .in_current_span()
    }

    pub fn accept_bi_stream(
        &self,
    ) -> impl Future<Output = io::Result<Option<(StreamId, (StreamReader, StreamWriter))>>> + Send
    {
        let params = self.parameters.clone();
        let streams = self.spaces.data().streams().clone();
        async move {
            let snd_wnd_size = params
                .get_remote_as::<VarInt>(ParameterId::InitialMaxStreamDataBidiLocal)
                .await?;
            Ok(Some(streams.accept_bi(snd_wnd_size.into_inner()).await?))
        }
        .instrument_in_current()
        .in_current_span()
    }

    pub fn accept_uni_stream(
        &self,
    ) -> impl Future<Output = io::Result<Option<(StreamId, StreamReader)>>> + Send {
        let streams = self.spaces.data().streams().clone();
        async move { Ok(Some(streams.accept_uni().await?)) }
            .instrument_in_current()
            .in_current_span()
    }

    #[cfg(feature = "unreliable")]
    pub fn unreliable_reader(&self) -> io::Result<DatagramReader> {
        self.spaces.data().datagrams().reader()
    }

    #[cfg(feature = "unreliable")]
    pub fn unreliable_writer(&self) -> impl Future<Output = io::Result<DatagramWriter>> + Send {
        let params = self.parameters.clone();
        let datagrams = self.spaces.data().datagrams().clone();
        async move {
            let max_datagram_frame_size = params
                .get_remote_as::<VarInt>(ParameterId::MaxDatagramFrameSize)
                .await?;
            datagrams.writer(max_datagram_frame_size.into_inner())
        }
        .instrument_in_current()
        .in_current_span()
    }

    pub fn add_path(&self, link: Link, pathway: Pathway) -> io::Result<()> {
        self.get_or_try_create_path(link, pathway, false)
            .map(|_| ())
    }

    pub fn del_path(&self, pathway: &Pathway) {
        self.paths.remove(pathway, "application removed");
    }
}

type ConnectionState = RwLock<Result<Components, Termination>>;

pub struct Connection {
    state: ConnectionState,
    qlog_span: qevent::telemetry::Span,
    tracing_span: tracing::Span,
}

impl Connection {
    pub fn enter_closing(&self, ccf: ConnectionCloseFrame) {
        let _span = (self.qlog_span.enter(), self.tracing_span.enter());
        let mut conn = self.state.write().unwrap();
        if let Ok(components) = conn.as_mut() {
            *conn = Err(components.clone().enter_closing(ccf));
        }
    }

    pub fn enter_draining(&self, ccf: ConnectionCloseFrame) {
        let _span = (self.qlog_span.enter(), self.tracing_span.enter());
        let mut conn = self.state.write().unwrap();
        match conn.as_mut() {
            Ok(core_conn) => *conn = Err(core_conn.clone().enter_draining(ccf)),
            Err(termination) => termination.enter_draining(),
        }
    }

    pub fn close(&self, reason: Cow<'static, str>, code: u64) {
        // let _span = (self.qlog_span.enter(), self.tracing_span.enter());
        let error_code = code.try_into().unwrap();
        self.enter_closing(ConnectionCloseFrame::new_app(error_code, reason));
    }

    fn try_map_components<T>(&self, op: impl Fn(&Components) -> T) -> io::Result<T> {
        let _span = (self.qlog_span.enter(), self.tracing_span.enter());
        self.state
            .read()
            .unwrap()
            .as_ref()
            .map(op)
            .map_err(|termination| termination.error().into())
    }

    pub async fn open_bi_stream(
        &self,
    ) -> io::Result<Option<(StreamId, (StreamReader, StreamWriter))>> {
        self.try_map_components(|core_conn| core_conn.open_bi_stream())?
            .await
    }

    pub async fn open_uni_stream(&self) -> io::Result<Option<(StreamId, StreamWriter)>> {
        self.try_map_components(|core_conn| core_conn.open_uni_stream())?
            .await
    }

    pub async fn accept_bi_stream(
        &self,
    ) -> io::Result<Option<(StreamId, (StreamReader, StreamWriter))>> {
        self.try_map_components(|core_conn| core_conn.accept_bi_stream())?
            .await
    }

    pub async fn accept_uni_stream(&self) -> io::Result<Option<(StreamId, StreamReader)>> {
        self.try_map_components(|core_conn| core_conn.accept_uni_stream())?
            .await
    }

    #[cfg(feature = "unreliable")]
    pub fn unreliable_reader(&self) -> io::Result<DatagramReader> {
        self.try_map_components(|core_conn| core_conn.unreliable_reader())?
    }

    #[cfg(feature = "unreliable")]
    pub async fn unreliable_writer(&self) -> io::Result<DatagramWriter> {
        self.try_map_components(|core_conn| core_conn.unreliable_writer())?
            .await
    }

    pub fn add_path(&self, link: Link, pathway: Pathway) -> io::Result<()> {
        self.try_map_components(|core_conn| core_conn.add_path(link, pathway))?
    }

    pub fn del_path(&self, pathway: &Pathway) -> io::Result<()> {
        self.try_map_components(|core_conn| core_conn.del_path(pathway))
    }

    pub fn is_active(&self) -> bool {
        self.try_map_components(|_| true).unwrap_or_default()
    }

    pub fn origin_dcid(&self) -> io::Result<cid::ConnectionId> {
        self.try_map_components(|core_conn| Ok(core_conn.parameters.get_origin_dcid()?))?
    }

    pub async fn handshaked(&self) {
        if let Ok(f) = self.try_map_components(|core_conn| core_conn.conn_state.handshaked()) {
            f.await
        }
    }

    pub async fn terminated(&self) {
        if let Ok(f) = self.try_map_components(|core_conn| core_conn.conn_state.terminated()) {
            f.await
        }
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        if let Ok(origin_dcid) = self.origin_dcid() {
            tracing::warn!("Connection {origin_dcid:x} is still active when dropped",);
        }
    }
}
