// Copyright (c) 2024 Jonas Schäfer <jonas@zombofant.net>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::borrow::Cow;
use std::io;

use futures::{ready, Sink, SinkExt, Stream, StreamExt};

use bytes::{Buf, BytesMut};

use tokio::io::{AsyncBufRead, AsyncWrite};

use xso::{
    exports::rxml::{self, writer::TrackNamespace, xml_ncname, Event, Namespace},
    AsXml, FromEventsBuilder, FromXml, Item,
};

use super::capture::{log_enabled, log_recv, log_send, CaptureBufRead};

use xmpp_parsers::ns::STREAM as XML_STREAM_NS;

pin_project_lite::pin_project! {
    // NOTE: due to limitations of pin_project_lite, the field comments are
    // no doc comments. Luckily, this struct is only `pub(super)` anyway.
    #[project = RawXmlStreamProj]
    pub(super) struct RawXmlStream<Io> {
        // The parser used for deserialising data.
        #[pin]
        parser: rxml::AsyncReader<CaptureBufRead<Io>>,

        // The writer used for serialising data.
        writer: rxml::writer::Encoder<rxml::writer::SimpleNamespaces>,

        // The default namespace to declare on the stream header.
        stream_ns: &'static str,

        // Buffer containing serialised data which will then be sent through
        // the inner `Io`. Sending that serialised data happens in
        // `poll_ready` and `poll_flush`, while appending serialised data
        // happens in `start_send`.
        tx_buffer: BytesMut,

        // Position inside tx_buffer up to which to-be-sent data has already
        // been logged.
        tx_buffer_logged: usize,

        // This signifies the limit at the point of which the Sink will
        // refuse to accept more data: if the `tx_buffer`'s size grows beyond
        // that high water mark, poll_ready will return Poll::Pending until
        // it has managed to flush enough data down the inner writer.
        //
        // Note that poll_ready will always attempt to progress the writes,
        // which further reduces the chance of hitting this limit unless
        // either the underlying writer gets stuck (e.g. TCP connection
        // breaking in a timeouty way) or a lot of data is written in bulk.
        // In both cases, the backpressure created by poll_ready returning
        // Pending is desirable.
        //
        // However, there is a catch: We don't assert this condition
        // in `start_send` at all. The reason is that we cannot suspend
        // serialisation of an XSO in the middle of writing it: it has to be
        // written in one batch or you have to start over later (this has to
        // do with the iterator state borrowing the data and futures getting
        // cancelled e.g. in tokio::select!). In order to facilitate
        // implementing a `Sink<T: AsXml>` on top of `RawXmlStream`, we
        // cannot be strict about what is going on in `start_send`:
        // `poll_ready` does not know what kind of data will be written (so
        // it could not make a size estimate, even if that was at all
        // possible with AsXml) and `start_send` is not a coroutine. So if
        // `Sink<T: AsXml>` wants to use `RawXmlStream`, it must be able to
        // submit an entire XSO's items in one batch to `RawXmlStream` after
        // it has reported to be ready once. That may easily make the buffer
        // reach its high water mark.
        //
        // So if we checked that condition in `start_send` (as opposed to
        // `poll_ready`), we would cause situations where submitting XSOs
        // failed randomly (with a panic or other errors) and would have to
        // be retried later.
        //
        // While failing with e.g. io::ErrorKind::WouldBlock is something
        // that could be investigated later, it would still require being
        // able to make an accurate estimate of the number of bytes needed to
        // serialise any given `AsXml`, because as pointed out earlier, once
        // we have started, there is no going back.
        //
        // Finally, none of that hurts much because `RawXmlStream` is only an
        // internal API. The high-level APIs will always call `poll_ready`
        // before sending an XSO, which means that we won't *grossly* go over
        // the TX buffer high water mark---unless you send a really large
        // XSO at once.
        tx_buffer_high_water_mark: usize,
    }
}

impl<Io: AsyncBufRead + AsyncWrite> RawXmlStream<Io> {
    fn new_writer(
        stream_ns: &'static str,
    ) -> rxml::writer::Encoder<rxml::writer::SimpleNamespaces> {
        let mut writer = rxml::writer::Encoder::new();
        writer
            .ns_tracker_mut()
            .declare_fixed(Some(xml_ncname!("stream")), XML_STREAM_NS.into());
        writer
            .ns_tracker_mut()
            .declare_fixed(None, stream_ns.into());
        writer
    }

    pub(super) fn new(io: Io, stream_ns: &'static str) -> Self {
        let parser = rxml::Parser::default();
        let mut io = CaptureBufRead::wrap(io);
        if log_enabled() {
            io.enable_capture();
        }
        Self {
            parser: rxml::AsyncReader::wrap(io, parser),
            writer: Self::new_writer(stream_ns),
            tx_buffer_logged: 0,
            stream_ns,
            tx_buffer: BytesMut::new(),

            // This basically means: "if we already have 2 kiB in our send
            // buffer, do not accept more data".
            // Please see the extensive words at
            //`Self::tx_buffer_high_water_mark` for details.
            tx_buffer_high_water_mark: 2048,
        }
    }

    pub(super) fn reset_state(self: Pin<&mut Self>) {
        let this = self.project();
        *this.parser.parser_pinned() = rxml::Parser::default();
        *this.writer = Self::new_writer(this.stream_ns);
    }

    pub(super) fn into_inner(self) -> Io {
        self.parser.into_inner().0.into_inner()
    }
}

impl<Io: AsyncWrite> RawXmlStream<Io> {
    /// Start sending an entire XSO.
    ///
    /// Unlike the `Sink` implementation, this provides nice syntax
    /// highlighting for the serialised data in log outputs (if enabled) *and*
    /// is error safe: if the XSO fails to serialise completely, it will be as
    /// if it hadn't been attempted to serialise it at all.
    ///
    /// Note that, like with `start_send`, the caller is responsible for
    /// ensuring that the stream is ready by polling
    /// [`<Self as Sink>::poll_ready`] as needed.
    pub(super) fn start_send_xso<T: AsXml>(self: Pin<&mut Self>, xso: &T) -> io::Result<()> {
        let mut this = self.project();
        let prev_len = this.tx_buffer.len();
        match this.try_send_xso(xso) {
            Ok(()) => Ok(()),
            Err(e) => {
                let curr_len = this.tx_buffer.len();
                this.tx_buffer.truncate(prev_len);
                log::trace!(
                    "SEND failed: {}. Rewinding buffer by {} bytes.",
                    e,
                    curr_len - prev_len
                );
                Err(e)
            }
        }
    }
}

impl<Io> RawXmlStream<Io> {
    fn parser_pinned(self: Pin<&mut Self>) -> &mut rxml::Parser {
        self.project().parser.parser_pinned()
    }

    fn stream_pinned(self: Pin<&mut Self>) -> Pin<&mut CaptureBufRead<Io>> {
        self.project().parser.inner_pinned()
    }

    pub(super) fn get_stream(&self) -> &Io {
        self.parser.inner().inner()
    }
}

impl<Io: AsyncBufRead> Stream for RawXmlStream<Io> {
    type Item = Result<rxml::Event, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();
        loop {
            return Poll::Ready(
                match ready!(this.parser.as_mut().poll_read(cx)).transpose() {
                    // Skip the XML declaration, nobody wants to hear about that.
                    Some(Ok(rxml::Event::XmlDeclaration(_, _))) => continue,
                    other => other,
                },
            );
        }
    }
}

impl<'x, Io: AsyncWrite> RawXmlStreamProj<'x, Io> {
    fn flush_tx_log(&mut self) {
        let range = &self.tx_buffer[*self.tx_buffer_logged..];
        if range.len() == 0 {
            return;
        }
        log_send(range);
        *self.tx_buffer_logged = self.tx_buffer.len();
    }

    fn start_send(&mut self, item: &xso::Item<'_>) -> io::Result<()> {
        self.writer
            .encode_into_bytes(item.as_rxml_item(), self.tx_buffer)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
    }

    fn try_send_xso<T: AsXml>(&mut self, xso: &T) -> io::Result<()> {
        let iter = match xso.as_xml_iter() {
            Ok(v) => v,
            Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
        };
        for item in iter {
            let item = match item {
                Ok(v) => v,
                Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
            };
            self.start_send(&item)?;
        }
        Ok(())
    }

    fn progress_write(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        self.flush_tx_log();
        while self.tx_buffer.len() > 0 {
            let written = match ready!(self
                .parser
                .as_mut()
                .inner_pinned()
                .poll_write(cx, &self.tx_buffer))
            {
                Ok(v) => v,
                Err(e) => return Poll::Ready(Err(e)),
            };
            self.tx_buffer.advance(written);
            *self.tx_buffer_logged = self
                .tx_buffer_logged
                .checked_sub(written)
                .expect("Buffer arithmetic error");
        }
        Poll::Ready(Ok(()))
    }
}

impl<'x, Io: AsyncWrite> Sink<xso::Item<'x>> for RawXmlStream<Io> {
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let mut this = self.project();
        match this.progress_write(cx) {
            // No progress on write, but if we have enough space in the buffer
            // it's ok nonetheless.
            Poll::Pending => (),
            // Some progress and it went fine, move on.
            Poll::Ready(Ok(())) => (),
            // Something went wrong -> return the error.
            Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())),
        }
        if this.tx_buffer.len() < *this.tx_buffer_high_water_mark {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let mut this = self.project();
        ready!(this.progress_write(cx))?;
        this.parser.as_mut().inner_pinned().poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let mut this = self.project();
        ready!(this.progress_write(cx))?;
        this.parser.as_mut().inner_pinned().poll_shutdown(cx)
    }

    fn start_send(self: Pin<&mut Self>, item: xso::Item<'x>) -> Result<(), Self::Error> {
        let mut this = self.project();
        this.start_send(&item)
    }
}

/// Error returned by the [`ReadXso`] future and the [`ReadXsoState`] helper.
pub(super) enum ReadXsoError {
    /// The outer element was closed before a child element could be read.
    ///
    /// This is typically the stream footer in XML stream applications.
    Footer,

    /// A hard error occurred.
    ///
    /// This is either a real I/O error or an error from the XML parser.
    /// Neither are recoverable, because the nesting state is lost and
    /// in addition, XML errors are not recoverable because they indicate a
    /// not well-formed document.
    Hard(io::Error),

    /// A parse error occurred.
    ///
    /// The XML structure was well-formed, but the data contained did not
    /// match the XSO which was attempted to be parsed. This error is
    /// recoverable: when this error is emitted, the XML stream is at the same
    /// nesting level as it was before the XSO was attempted to be read; all
    /// XML structure which belonged to the XSO which failed to parse has
    /// been consumed. This allows to read more XSOs even if one fails to
    /// parse.
    Parse(xso::error::Error),
}

impl From<ReadXsoError> for io::Error {
    fn from(other: ReadXsoError) -> Self {
        match other {
            ReadXsoError::Hard(v) => v,
            ReadXsoError::Parse(e) => io::Error::new(io::ErrorKind::InvalidData, e),
            ReadXsoError::Footer => io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "element footer while waiting for XSO element start",
            ),
        }
    }
}

impl From<io::Error> for ReadXsoError {
    fn from(other: io::Error) -> Self {
        Self::Hard(other)
    }
}

impl From<xso::error::Error> for ReadXsoError {
    fn from(other: xso::error::Error) -> Self {
        Self::Parse(other)
    }
}

/// State for reading an XSO from a `Stream<Item = Result<rxml::Event, ...>>`.
///
/// Due to pinning, it is simpler to implement the statemachine in a dedicated
/// enum and let the actual (pinned) future pass the stream toward this enum's
/// function.
///
/// This is used by both [`ReadXso`] and the [`super::XmlStream`] itself.
#[derive(Default)]
pub(super) enum ReadXsoState<T: FromXml> {
    /// The [`rxml::Event::StartElement`] event was not seen yet.
    ///
    /// In this state, XML whitespace is ignored (as per RFC 6120 § 11.7), but
    /// other text data is rejected.
    #[default]
    PreData,

    /// The [`rxml::Event::StartElement`] event was received.
    ///
    /// The inner value is the builder for the "return type" of this enum and
    /// the implementation in the [`xso`] crate does all the heavy lifting:
    /// we'll only send events in its general direction.
    // We use the fallible parsing here so that we don't have to do the depth
    // accounting ourselves.
    Parsing(<Result<T, xso::error::Error> as FromXml>::Builder),

    /// The parsing has completed (successful or not).
    ///
    /// This is a final state and attempting to advance the state will panic.
    /// This is in accordance with [`core::future::Future::poll`]'s contract,
    /// for which this enum is primarily used.
    Done,
}

impl<T: FromXml> ReadXsoState<T> {
    /// Progress reading the XSO from the given source.
    ///
    /// This attempts to parse a single XSO from the underlying stream,
    /// while discarding any XML whitespace before the beginning of the XSO.
    ///
    /// If the XSO is parsed successfully, the method returns Ready with the
    /// parsed value. If parsing fails or an I/O error occurs, an appropriate
    /// error is returned.
    ///
    /// If parsing fails, the entire XML subtree belonging to the XSO is
    /// nonetheless processed. That makes parse errors recoverable: After
    /// `poll_advance` has returned Ready with either  an Ok result or a
    /// [`ReadXsoError::Parse`] error variant, another XSO can be read and the
    /// XML parsing will be at the same nesting depth as it was before the
    /// first call to `poll_advance`.
    ///
    /// Note that this guarantee does not hold for non-parse errors (i.e. for
    /// the other variants of [`ReadXsoError`]): I/O errors as well as
    /// occurrence of the outer closing element are fatal.
    ///
    /// The `source` passed to `poll_advance` should be the same on every
    /// call.
    pub(super) fn poll_advance<Io: AsyncBufRead>(
        &mut self,
        mut source: Pin<&mut RawXmlStream<Io>>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<T, ReadXsoError>> {
        loop {
            // Disable text buffering before the start event. That way, we
            // don't accumulate infinite amounts of XML whitespace caused by
            // whitespace keepalives.
            // (And also, we'll know faster when the remote side sends
            // non-whitespace garbage.)
            let text_buffering = match self {
                ReadXsoState::PreData => false,
                _ => true,
            };
            source
                .as_mut()
                .parser_pinned()
                .set_text_buffering(text_buffering);

            let ev = ready!(source.as_mut().poll_next(cx)).transpose()?;
            match self {
                ReadXsoState::PreData => {
                    log::trace!("ReadXsoState::PreData ev = {:?}", ev);
                    match ev {
                        Some(rxml::Event::XmlDeclaration(_, _)) => (),
                        Some(rxml::Event::Text(_, data)) => {
                            if xso::is_xml_whitespace(data.as_bytes()) {
                                log::trace!("Received {} bytes of whitespace", data.len());
                                source.as_mut().stream_pinned().discard_capture();
                                continue;
                            } else {
                                *self = ReadXsoState::Done;
                                return Poll::Ready(Err(io::Error::new(
                                    io::ErrorKind::InvalidData,
                                    "non-whitespace text content before XSO",
                                )
                                .into()));
                            }
                        }
                        Some(rxml::Event::StartElement(_, name, attrs)) => {
                            *self = ReadXsoState::Parsing(
                                <Result<T, xso::error::Error> as FromXml>::from_events(name, attrs)
                                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?,
                            );
                        }
                        // Amounts to EOF, as we expect to start on the stream level.
                        Some(rxml::Event::EndElement(_)) => {
                            *self = ReadXsoState::Done;
                            return Poll::Ready(Err(ReadXsoError::Footer));
                        }
                        None => {
                            *self = ReadXsoState::Done;
                            return Poll::Ready(Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "end of parent element before XSO started",
                            )
                            .into()));
                        }
                    }
                }
                ReadXsoState::Parsing(builder) => {
                    let Some(ev) = ev else {
                        *self = ReadXsoState::Done;
                        return Poll::Ready(Err(io::Error::new(
                            io::ErrorKind::UnexpectedEof,
                            "eof during XSO parsing",
                        )
                        .into()));
                    };

                    match builder.feed(ev) {
                        Err(err) => {
                            *self = ReadXsoState::Done;
                            source.as_mut().stream_pinned().discard_capture();
                            return Poll::Ready(Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                err,
                            )
                            .into()));
                        }
                        Ok(Some(Err(err))) => {
                            *self = ReadXsoState::Done;
                            log_recv(Some(&err), source.as_mut().stream_pinned().take_capture());
                            return Poll::Ready(Err(ReadXsoError::Parse(err)));
                        }
                        Ok(Some(Ok(value))) => {
                            *self = ReadXsoState::Done;
                            log_recv(None, source.as_mut().stream_pinned().take_capture());
                            return Poll::Ready(Ok(value));
                        }
                        Ok(None) => (),
                    }
                }

                // The error talks about "future", simply because that is
                // where `Self` is used (inside `core::future::Future::poll`).
                ReadXsoState::Done => panic!("future polled after completion"),
            }
        }
    }
}

/// Future to read a single XSO from a stream.
pub(super) struct ReadXso<'x, Io, T: FromXml> {
    /// Stream to read the future from.
    inner: Pin<&'x mut RawXmlStream<Io>>,

    /// Current state of parsing.
    state: ReadXsoState<T>,
}

impl<'x, Io: AsyncBufRead, T: FromXml> ReadXso<'x, Io, T> {
    /// Start reading a single XSO from a stream.
    pub(super) fn read_from(stream: Pin<&'x mut RawXmlStream<Io>>) -> Self {
        Self {
            inner: stream,
            state: ReadXsoState::PreData,
        }
    }
}

impl<'x, Io: AsyncBufRead, T: FromXml> Future for ReadXso<'x, Io, T>
where
    T::Builder: Unpin,
{
    type Output = Result<T, ReadXsoError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        this.state.poll_advance(this.inner.as_mut(), cx)
    }
}

/// Contains metadata from an XML stream header
#[derive(Default)]
pub struct StreamHeader<'x> {
    /// The optional `from` attribute.
    pub from: Option<Cow<'x, str>>,

    /// The optional `to` attribute.
    pub to: Option<Cow<'x, str>>,

    /// The optional `id` attribute.
    pub id: Option<Cow<'x, str>>,
}

impl<'x> StreamHeader<'x> {
    /// Take the contents and return them as new object.
    ///
    /// `self` will be left with all its parts set to `None`.
    pub fn take(&mut self) -> Self {
        Self {
            from: self.from.take(),
            to: self.to.take(),
            id: self.id.take(),
        }
    }

    pub(super) async fn send<Io: AsyncWrite>(
        self,
        mut stream: Pin<&mut RawXmlStream<Io>>,
    ) -> io::Result<()> {
        stream
            .send(Item::XmlDeclaration(rxml::XmlVersion::V1_0))
            .await?;
        stream
            .send(Item::ElementHeadStart(
                Namespace::from(XML_STREAM_NS),
                Cow::Borrowed(xml_ncname!("stream")),
            ))
            .await?;
        if let Some(from) = self.from {
            stream
                .send(Item::Attribute(
                    Namespace::NONE,
                    Cow::Borrowed(xml_ncname!("from")),
                    from,
                ))
                .await?;
        }
        if let Some(to) = self.to {
            stream
                .send(Item::Attribute(
                    Namespace::NONE,
                    Cow::Borrowed(xml_ncname!("to")),
                    to,
                ))
                .await?;
        }
        if let Some(id) = self.id {
            stream
                .send(Item::Attribute(
                    Namespace::NONE,
                    Cow::Borrowed(xml_ncname!("id")),
                    id,
                ))
                .await?;
        }
        stream
            .send(Item::Attribute(
                Namespace::NONE,
                Cow::Borrowed(xml_ncname!("version")),
                Cow::Borrowed("1.0"),
            ))
            .await?;
        stream.send(Item::ElementHeadEnd).await?;
        Ok(())
    }
}

impl StreamHeader<'static> {
    pub(super) async fn recv<Io: AsyncBufRead>(
        mut stream: Pin<&mut RawXmlStream<Io>>,
    ) -> io::Result<Self> {
        loop {
            match stream.as_mut().next().await.transpose()? {
                Some(Event::StartElement(_, (ns, name), mut attrs)) => {
                    if ns != XML_STREAM_NS || name != "stream" {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "unknown stream header",
                        ));
                    }

                    match attrs.remove(Namespace::none(), "version") {
                        Some(v) => {
                            if v != "1.0" {
                                return Err(io::Error::new(
                                    io::ErrorKind::InvalidData,
                                    format!("unsuppored stream version: {}", v),
                                ));
                            }
                        }
                        None => {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "required `version` attribute missing",
                            ))
                        }
                    }

                    let from = attrs.remove(Namespace::none(), "from");
                    let to = attrs.remove(Namespace::none(), "to");
                    let id = attrs.remove(Namespace::none(), "id");
                    let _ = attrs.remove(Namespace::xml(), "lang");

                    if let Some(((ns, name), _)) = attrs.into_iter().next() {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            format!("unexpected stream header attribute: {{{}}}{}", ns, name),
                        ));
                    }

                    return Ok(StreamHeader {
                        from: from.map(Cow::Owned),
                        to: to.map(Cow::Owned),
                        id: id.map(Cow::Owned),
                    });
                }
                Some(Event::Text(_, _)) | Some(Event::EndElement(_)) => {
                    return Err(io::Error::new(
                        io::ErrorKind::UnexpectedEof,
                        "unexpected content before stream header",
                    ))
                }
                // We cannot loop infinitely here because the XML parser will
                // prevent more than one XML declaration from being parsed.
                Some(Event::XmlDeclaration(_, _)) => (),
                None => {
                    return Err(io::Error::new(
                        io::ErrorKind::UnexpectedEof,
                        "eof before stream header",
                    ))
                }
            }
        }
    }
}