/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::io::Cursor;

use anyhow::anyhow;
use anyhow::Result;
use bufsize::SizeCounter;
use bytes::Bytes;
use bytes::BytesMut;
use ghost::phantom;

use crate::binary_type::CopyFromBuf;
use crate::bufext::BufExt;
use crate::bufext::BufMutExt;
use crate::bufext::DeserializeSource;
use crate::deserialize::Deserialize;
use crate::errors::ProtocolError;
use crate::framing::Framing;
use crate::protocol::Field;
use crate::protocol::Protocol;
use crate::protocol::ProtocolReader;
use crate::protocol::ProtocolWriter;
use crate::serialize::Serialize;
use crate::thrift_protocol::MessageType;
use crate::thrift_protocol::ProtocolID;
use crate::ttype::TType;
use crate::varint;

const COMPACT_PROTOCOL_VERSION: u8 = 0x02;
const PROTOCOL_ID: u8 = 0x82;
const TYPE_MASK: u8 = 0xE0;
const TYPE_SHIFT_AMOUNT: usize = 5;
const VERSION_MASK: u8 = 0x1f;

#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
#[repr(u8)]
pub enum CType {
    Stop = 0,
    BoolTrue = 1,
    BoolFalse = 2,
    Byte = 3,
    I16 = 4,
    I32 = 5,
    I64 = 6,
    Double = 7,
    Binary = 8,
    List = 9,
    Set = 10,
    Map = 11,
    Struct = 12,
    Float = 13,
}

impl TryFrom<i8> for CType {
    type Error = anyhow::Error;

    fn try_from(v: i8) -> Result<Self> {
        let ret = match v {
            0 => CType::Stop,
            1 => CType::BoolTrue,
            2 => CType::BoolFalse,
            3 => CType::Byte,
            4 => CType::I16,
            5 => CType::I32,
            6 => CType::I64,
            7 => CType::Double,
            8 => CType::Binary,
            9 => CType::List,
            10 => CType::Set,
            11 => CType::Map,
            12 => CType::Struct,
            13 => CType::Float,
            _ => bail_err!(ProtocolError::InvalidTypeTag),
        };
        Ok(ret)
    }
}

impl From<CType> for TType {
    fn from(cty: CType) -> TType {
        match cty {
            CType::Stop => TType::Stop,
            CType::BoolFalse | CType::BoolTrue => TType::Bool,
            CType::Byte => TType::Byte,
            CType::Double => TType::Double,
            CType::I16 => TType::I16,
            CType::I32 => TType::I32,
            CType::I64 => TType::I64,
            CType::Binary => TType::String,
            CType::Struct => TType::Struct,
            CType::Map => TType::Map,
            CType::Set => TType::Set,
            CType::List => TType::List,
            CType::Float => TType::Float,
        }
    }
}

impl From<TType> for CType {
    fn from(tty: TType) -> CType {
        match tty {
            TType::Stop => CType::Stop,
            TType::Bool => CType::BoolTrue,
            TType::Byte => CType::Byte,
            TType::Double => CType::Double,
            TType::I16 => CType::I16,
            TType::I32 => CType::I32,
            TType::I64 => CType::I64,
            TType::String => CType::Binary,
            TType::Struct => CType::Struct,
            TType::Map => CType::Map,
            TType::Set => CType::Set,
            TType::List => CType::List,
            TType::Float => CType::Float,
            bad => panic!("Don't know how to convert TType {:?} to CType", bad),
        }
    }
}

impl From<bool> for CType {
    fn from(b: bool) -> CType {
        if b { CType::BoolTrue } else { CType::BoolFalse }
    }
}

/// An efficient, dense encoding that uses variable-length integers and squeezes
/// data into unused bits wherever possible.
///
/// ```ignore
/// let protocol = CompactProtocol;
/// let transport = HttpClient::new(ENDPOINT)?;
/// let client = <dyn BuckGraphService>::new(protocol, transport);
/// ```
///
/// The type parameter is the Framing expected by the transport on which this
/// protocol is operating. Usually by convention the transport itself serves as
/// the Framing impl, so for example in the case of HttpClient above, the
/// compiler has inferred `F = HttpClient`.
///
/// Where the compiler reports that a Framing can't be inferred, one can be
/// specified explicitly:
///
/// ```ignore
/// let protocol = CompactProtocol::<SRHeaderTransport>;
/// ```
#[phantom]
#[derive(Copy, Clone)]
pub struct CompactProtocol<F = Bytes>;

#[derive(Debug, Clone)]
struct EncState {
    idxstack: Vec<i16>,
    lastidx: i16,
    field: Option<(TType, i16)>,
}

impl EncState {
    fn new() -> Self {
        EncState {
            idxstack: Vec::new(),
            lastidx: 0,
            field: None,
        }
    }

    fn struct_begin(&mut self) {
        self.idxstack.push(self.lastidx);
        self.lastidx = 0;
    }

    fn struct_end(&mut self) {
        self.lastidx = self.idxstack.pop().expect("struct stack underrun");
    }

    fn field_begin(&mut self, tty: TType, idx: i16) {
        self.field = Some((tty, idx));
    }

    fn field_get(&mut self) -> Option<(TType, i16)> {
        self.field.take()
    }

    fn field_end(&mut self) {
        debug_assert!(
            // the field was consumed during the write, or
            self.field.is_none() ||
            // the field contains a void which wasn't written
            matches!(self.field, Some((TType::Void, _)))
        )
    }

    fn in_field(&self) -> bool {
        self.field.is_some()
    }
}

pub struct CompactProtocolSerializer<B> {
    state: EncState,
    buffer: B,
    string_limit: Option<usize>,
    container_limit: Option<usize>,
}

pub struct CompactProtocolDeserializer<B> {
    state: EncState,
    boolfield: Option<CType>,
    buffer: B,
    string_limit: Option<usize>,
    container_limit: Option<usize>,
}

impl<F> Protocol for CompactProtocol<F>
where
    F: Framing + 'static,
{
    type Frame = F;
    type Sizer = CompactProtocolSerializer<SizeCounter>;
    type Serializer = CompactProtocolSerializer<F::EncBuf>;
    type Deserializer = CompactProtocolDeserializer<F::DecBuf>;

    const PROTOCOL_ID: ProtocolID = ProtocolID::CompactProtocol;

    fn serializer<SZ, SER>(size: SZ, ser: SER) -> <Self::Serializer as ProtocolWriter>::Final
    where
        SZ: FnOnce(&mut Self::Sizer),
        SER: FnOnce(&mut Self::Serializer),
    {
        // Instantiate a serializer with SizeCounter to count up the output buffer size
        let mut sizer = CompactProtocolSerializer {
            state: EncState::new(),
            buffer: SizeCounter::new(),
            container_limit: None,
            string_limit: None,
        };
        size(&mut sizer);
        let sz = sizer.finish();

        // Now that we have the size, allocate an output buffer and serialize into it
        let mut buf = CompactProtocolSerializer {
            state: EncState::new(),
            buffer: F::enc_with_capacity(sz),
            container_limit: None,
            string_limit: None,
        };
        ser(&mut buf);

        // Done
        buf.finish()
    }

    fn deserializer(buf: F::DecBuf) -> Self::Deserializer {
        CompactProtocolDeserializer::new(buf)
    }

    fn into_buffer(deser: Self::Deserializer) -> F::DecBuf {
        deser.into_inner()
    }
}

impl<B> CompactProtocolSerializer<B>
where
    B: BufMutExt,
{
    pub fn with_buffer(buffer: B) -> Self {
        Self {
            state: EncState::new(),
            buffer,
            string_limit: None,
            container_limit: None,
        }
    }

    pub fn into_inner(self) -> B {
        self.buffer
    }

    #[inline]
    fn write_varint_i64(&mut self, v: i64) {
        self.buffer.put_varint_i64(v)
    }

    #[inline]
    fn write_varint_i32(&mut self, v: i32) {
        self.write_varint_i64(v as i64);
    }

    #[inline]
    fn write_varint_i16(&mut self, v: i16) {
        self.write_varint_i64(v as i64);
    }

    /// If we've just got a write_field_begin, then this will emit an appropriate field begin
    /// record to the wire. If we haven't then it's a no-op. This is used to defer emitting
    /// the field header until we get the field value, mostly to handle the strangeness of bool.
    fn write_field_id(&mut self, cty: CType) {
        match self.state.field_get() {
            None => {}
            Some((tty, idx)) => {
                debug_assert_eq!(tty, TType::from(cty));

                let delta = idx - self.state.lastidx;
                self.state.lastidx = idx;

                if delta <= 0 || delta > 15 {
                    self.buffer.put_u8(cty as u8);
                    self.write_varint_i16(idx);
                } else {
                    self.buffer.put_u8(((delta as u8) << 4) | (cty as u8));
                }
            }
        }
    }

    /// Common code for writing the start of a sequence of elements, used for lists and sets
    fn write_sequence(&mut self, elem_type: TType, size: usize) {
        assert!(
            self.container_limit.map_or(true, |lim| size < lim),
            "container too large {}, lim {:?}",
            size,
            self.container_limit
        );

        let cty = CType::from(elem_type) as u8;

        // Size = 0x0f is used as the marker for a separate varint size
        if size < 0x0f {
            self.buffer.put_u8((size as u8) << 4 | cty);
        } else {
            self.buffer.put_u8((0x0f << 4) | cty);
            self.buffer.put_varint_u64(size as u64);
        }
    }
}

impl<B: BufMutExt> ProtocolWriter for CompactProtocolSerializer<B> {
    type Final = B::Final; // Our final form is whatever the buffer produces

    fn write_message_begin(&mut self, name: &str, msgtype: MessageType, seqid: u32) {
        let msgtype = msgtype as u8;
        self.buffer.put_u8(PROTOCOL_ID);
        self.buffer
            .put_u8(COMPACT_PROTOCOL_VERSION | ((msgtype << TYPE_SHIFT_AMOUNT) & TYPE_MASK));
        self.buffer.put_varint_u64(seqid as u64);
        self.write_string(name);
    }

    #[inline]
    fn write_message_end(&mut self) {}

    #[inline]
    fn write_struct_begin(&mut self, _name: &str) {
        self.write_field_id(CType::Struct);
        self.state.struct_begin();
    }

    #[inline]
    fn write_struct_end(&mut self) {
        self.state.struct_end();
    }

    #[inline]
    fn write_field_begin(&mut self, _name: &str, type_id: TType, id: i16) {
        // Note the field, but we'll emit it when the value comes
        self.state.field_begin(type_id, id);
    }

    #[inline]
    fn write_field_end(&mut self) {
        self.state.field_end()
    }

    #[inline]
    fn write_field_stop(&mut self) {
        self.buffer.put_u8(CType::Stop as u8)
    }

    fn write_map_begin(&mut self, key_type: TType, value_type: TType, size: usize) {
        assert!(
            self.container_limit.map_or(true, |lim| size < lim),
            "map too large {}, lim {:?}",
            size,
            self.container_limit
        );

        self.write_field_id(CType::Map);
        self.buffer.put_varint_u64(size as u64);
        if size > 0 {
            let ckty = CType::from(key_type);
            let cvty = CType::from(value_type);

            self.buffer.put_u8((ckty as u8) << 4 | (cvty as u8));
        }
    }

    #[inline]
    fn write_map_key_begin(&mut self) {}

    #[inline]
    fn write_map_value_begin(&mut self) {}

    #[inline]
    fn write_map_end(&mut self) {}

    fn write_list_begin(&mut self, elem_type: TType, size: usize) {
        assert!(self.container_limit.map_or(true, |lim| size < lim));
        self.write_field_id(CType::List);
        self.write_sequence(elem_type, size);
    }

    #[inline]
    fn write_list_value_begin(&mut self) {}

    #[inline]
    fn write_list_end(&mut self) {}

    fn write_set_begin(&mut self, elem_type: TType, size: usize) {
        self.write_field_id(CType::Set);
        self.write_sequence(elem_type, size);
    }

    #[inline]
    fn write_set_value_begin(&mut self) {}

    fn write_set_end(&mut self) {}

    fn write_bool(&mut self, value: bool) {
        // If we're in a field then we need to encode the bool value as the field's type.
        // Otherwise we just emit it as a byte, but still using the field type values as
        // the encoding.
        if self.state.in_field() {
            self.write_field_id(CType::from(value))
        } else if value {
            self.write_byte(CType::BoolTrue as i8)
        } else {
            self.write_byte(CType::BoolFalse as i8)
        }
    }

    fn write_byte(&mut self, value: i8) {
        self.write_field_id(CType::Byte);
        self.buffer.put_i8(value)
    }

    fn write_i16(&mut self, value: i16) {
        self.write_field_id(CType::I16);
        self.write_varint_i16(value)
    }

    fn write_i32(&mut self, value: i32) {
        self.write_field_id(CType::I32);
        self.write_varint_i32(value)
    }

    fn write_i64(&mut self, value: i64) {
        self.write_field_id(CType::I64);
        self.write_varint_i64(value)
    }

    fn write_double(&mut self, value: f64) {
        self.write_field_id(CType::Double);
        self.buffer.put_f64(value)
    }

    fn write_float(&mut self, value: f32) {
        self.write_field_id(CType::Float);
        self.buffer.put_f32(value)
    }

    #[inline]
    fn write_string(&mut self, value: &str) {
        self.write_binary(value.as_bytes());
    }

    fn write_binary(&mut self, value: &[u8]) {
        let size = value.len();
        assert!(
            self.string_limit.map_or(true, |lim| size < lim),
            "string too large {}, lim {:?}",
            size,
            self.string_limit
        );

        self.write_field_id(CType::Binary);
        self.buffer.put_varint_u64(size as u64);
        self.buffer.put_slice(value)
    }

    fn finish(self) -> B::Final {
        self.buffer.finalize()
    }
}

impl<B: BufExt> CompactProtocolDeserializer<B> {
    pub fn new(buffer: B) -> Self {
        CompactProtocolDeserializer {
            state: EncState::new(),
            boolfield: None,
            buffer,
            string_limit: None,
            container_limit: None,
        }
    }

    pub fn into_inner(self) -> B {
        self.buffer
    }

    fn peek_bytes(&self, len: usize) -> Option<&[u8]> {
        if self.buffer.chunk().len() >= len {
            Some(&self.buffer.chunk()[..len])
        } else {
            None
        }
    }

    fn read_varint_u64(&mut self) -> Result<u64> {
        varint::read_u64(&mut self.buffer)
    }

    fn read_varint_i64(&mut self) -> Result<i64> {
        self.read_varint_u64().map(varint::unzigzag)
    }

    fn read_varint_i32(&mut self) -> Result<i32> {
        self.read_varint_i64()
            .and_then(|v| i32::try_from(v).map_err(|_| ProtocolError::InvalidValue.into()))
    }

    fn read_varint_i16(&mut self) -> Result<i16> {
        self.read_varint_i64()
            .and_then(|v| i16::try_from(v).map_err(|_| (ProtocolError::InvalidValue).into()))
    }
}

impl<B: BufExt> ProtocolReader for CompactProtocolDeserializer<B> {
    fn read_message_begin<F, T>(&mut self, msgfn: F) -> Result<(T, MessageType, u32)>
    where
        F: FnOnce(&[u8]) -> T,
    {
        let protocolid = self.read_byte()? as u8;
        ensure_err!(protocolid == PROTOCOL_ID, ProtocolError::BadVersion);

        let vandty = self.read_byte()? as u8;
        ensure_err!(
            (vandty & VERSION_MASK) == COMPACT_PROTOCOL_VERSION,
            ProtocolError::BadVersion
        );

        let msgty = (vandty & TYPE_MASK) >> TYPE_SHIFT_AMOUNT;
        let msgty = MessageType::try_from(msgty as u32)?;

        let seqid = self.read_varint_u64()? as u32;

        let name = {
            let len = self.read_varint_u64()? as usize;
            let (len, name) = {
                if self.peek_bytes(len).is_some() {
                    let namebuf = self.peek_bytes(len).unwrap();
                    (namebuf.len(), msgfn(namebuf))
                } else {
                    ensure_err!(
                        self.buffer.remaining() >= len,
                        ProtocolError::InvalidDataLength
                    );
                    let namebuf: Vec<u8> = Vec::copy_from_buf(&mut self.buffer, len);
                    (0, msgfn(namebuf.as_slice()))
                }
            };
            self.buffer.advance(len);
            name
        };

        Ok((name, msgty, seqid))
    }

    fn read_message_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_struct_begin<F, T>(&mut self, namefn: F) -> Result<T>
    where
        F: FnOnce(&[u8]) -> T,
    {
        self.state.struct_begin();
        Ok(namefn(&[]))
    }

    fn read_struct_end(&mut self) -> Result<()> {
        self.state.struct_end();
        Ok(())
    }

    fn read_field_begin<F, T>(&mut self, fieldfn: F, _fields: &[Field]) -> Result<(T, TType, i16)>
    where
        F: FnOnce(&[u8]) -> T,
    {
        let tyid = self.read_byte()?;
        let cty = CType::try_from(tyid & 0x0f)?;
        let didx = (tyid >> 4) & 0x0f;

        let tty = TType::from(cty);

        // Check if we're stopping, or get the next field id
        let idx = match (tty, didx) {
            (TType::Stop, _) => 0,
            (_, 0) => self.read_varint_i16()?,
            (_, didx) => self.state.lastidx + (didx as i16),
        };

        self.state.lastidx = idx;

        if tty == TType::Bool {
            self.boolfield = Some(cty);
        }

        let f = fieldfn(&[]);
        Ok((f, tty, idx))
    }

    fn read_field_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_map_begin(&mut self) -> Result<(TType, TType, Option<usize>)> {
        let size = self.read_varint_u64()? as usize;

        ensure_err!(
            self.container_limit.map_or(true, |lim| size < lim),
            ProtocolError::InvalidDataLength
        );

        // If the size is 0 we treat the kv type byte as 0, which ends up being (Stop, Stop)
        let kvtype = if size > 0 { self.read_byte()? } else { 0 };

        let kcty = CType::try_from((kvtype >> 4) & 0x0f)?;
        let vcty = CType::try_from((kvtype) & 0x0f)?;

        Ok((TType::from(kcty), TType::from(vcty), Some(size)))
    }

    #[inline]
    fn read_map_key_begin(&mut self) -> Result<bool> {
        Ok(true)
    }

    #[inline]
    fn read_map_value_begin(&mut self) -> Result<()> {
        Ok(())
    }

    #[inline]
    fn read_map_value_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_map_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_list_begin(&mut self) -> Result<(TType, Option<usize>)> {
        let szty = self.read_byte()?;
        let cty = CType::try_from(szty & 0x0f)?;
        let elem_type = TType::from(cty);

        let size = match (szty >> 4) & 0x0f {
            0x0f => self.read_varint_u64()? as usize,
            sz => sz as usize,
        };

        ensure_err!(
            self.container_limit.map_or(true, |lim| size < lim),
            ProtocolError::InvalidDataLength
        );

        Ok((elem_type, Some(size)))
    }

    #[inline]
    fn read_list_value_begin(&mut self) -> Result<bool> {
        Ok(true)
    }

    #[inline]
    fn read_list_value_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_list_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_set_begin(&mut self) -> Result<(TType, Option<usize>)> {
        self.read_list_begin()
    }

    #[inline]
    fn read_set_value_begin(&mut self) -> Result<bool> {
        Ok(true)
    }

    #[inline]
    fn read_set_value_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_set_end(&mut self) -> Result<()> {
        Ok(())
    }

    fn read_bool(&mut self) -> Result<bool> {
        // If we're in a bool field then take the value from that, otherwise
        // read it as a byte from the wire.
        let cty = match self.boolfield.take() {
            None => CType::try_from(self.read_byte()?)?,
            Some(cty) => cty,
        };

        // Match cpp2 - not strict on value of false. This is supposed to be encoded as a
        // CType::BoolTrue/BoolFalse, which is 1/2 - but we just check for true == 1. I suspect
        // this is to allow false to be encoded as 0 (or anything other than 1).
        Ok(cty == CType::BoolTrue)
    }

    fn read_byte(&mut self) -> Result<i8> {
        ensure_err!(self.buffer.remaining() >= 1, ProtocolError::EOF);

        Ok(self.buffer.get_i8())
    }

    fn read_i16(&mut self) -> Result<i16> {
        self.read_varint_i16()
    }

    fn read_i32(&mut self) -> Result<i32> {
        self.read_varint_i32()
    }

    fn read_i64(&mut self) -> Result<i64> {
        self.read_varint_i64()
    }

    fn read_double(&mut self) -> Result<f64> {
        ensure_err!(self.buffer.remaining() >= 8, ProtocolError::EOF);

        Ok(self.buffer.get_f64())
    }

    fn read_float(&mut self) -> Result<f32> {
        ensure_err!(self.buffer.remaining() >= 4, ProtocolError::EOF);

        Ok(self.buffer.get_f32())
    }

    fn read_string(&mut self) -> Result<String> {
        let vec = self.read_binary::<Vec<u8>>()?;

        String::from_utf8(vec)
            .map_err(|utf8_error| anyhow!("deserializing `string` from Thrift compact protocol got invalid utf-8, you need to use `binary` instead: {utf8_error}"))
    }

    fn read_binary<V: CopyFromBuf>(&mut self) -> Result<V> {
        let received_len = self.read_varint_u64()? as usize;
        ensure_err!(
            self.string_limit.map_or(true, |lim| received_len < lim),
            ProtocolError::InvalidDataLength
        );
        ensure_err!(self.buffer.remaining() >= received_len, ProtocolError::EOF);

        Ok(V::copy_from_buf(&mut self.buffer, received_len))
    }
}

/// How large an item will be when `serialize()` is called
pub fn serialize_size<T>(v: &T) -> usize
where
    T: Serialize<CompactProtocolSerializer<SizeCounter>>,
{
    let mut sizer = CompactProtocolSerializer {
        state: EncState::new(),
        buffer: SizeCounter::new(),
        container_limit: None,
        string_limit: None,
    };
    v.write(&mut sizer);
    sizer.finish()
}

/// Serialize a Thrift value using the compact protocol to a pre-allocated buffer.
/// This will panic if the buffer is not large enough. A buffer at least as
/// large as the return value of `serialize_size` will not panic.
pub fn serialize_to_buffer<T>(v: T, buffer: BytesMut) -> CompactProtocolSerializer<BytesMut>
where
    T: Serialize<CompactProtocolSerializer<BytesMut>>,
{
    // Now that we have the size, allocate an output buffer and serialize into it
    let mut buf = CompactProtocolSerializer {
        state: EncState::new(),
        buffer,
        container_limit: None,
        string_limit: None,
    };
    v.write(&mut buf);
    buf
}

pub trait SerializeRef:
    Serialize<CompactProtocolSerializer<SizeCounter>> + Serialize<CompactProtocolSerializer<BytesMut>>
where
    for<'a> &'a Self: Serialize<CompactProtocolSerializer<SizeCounter>>,
    for<'a> &'a Self: Serialize<CompactProtocolSerializer<BytesMut>>,
{
}

impl<T> SerializeRef for T
where
    T: Serialize<CompactProtocolSerializer<BytesMut>>,
    T: Serialize<CompactProtocolSerializer<SizeCounter>>,
    for<'a> &'a T: Serialize<CompactProtocolSerializer<BytesMut>>,
    for<'a> &'a T: Serialize<CompactProtocolSerializer<SizeCounter>>,
{
}

/// Serialize a Thrift value using the compact protocol.
pub fn serialize<T>(v: T) -> Bytes
where
    T: Serialize<CompactProtocolSerializer<SizeCounter>>
        + Serialize<CompactProtocolSerializer<BytesMut>>,
{
    let sz = serialize_size(&v);
    let buf = serialize_to_buffer(v, BytesMut::with_capacity(sz));
    // Done
    buf.finish()
}

pub trait DeserializeSlice:
    for<'a> Deserialize<CompactProtocolDeserializer<Cursor<&'a [u8]>>>
{
}

impl<T> DeserializeSlice for T where
    T: for<'a> Deserialize<CompactProtocolDeserializer<Cursor<&'a [u8]>>>
{
}

/// Deserialize a Thrift blob using the compact protocol.
pub fn deserialize<T, B, C>(b: B) -> Result<T>
where
    B: Into<DeserializeSource<C>>,
    C: BufExt,
    T: Deserialize<CompactProtocolDeserializer<C>>,
{
    let source: DeserializeSource<C> = b.into();
    let mut deser = CompactProtocolDeserializer::new(source.0);
    T::read(&mut deser)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::marker::PhantomData;

use anyhow::Error;

use crate::bufext::BufExt;
use crate::framing::Framing;
use crate::framing::FramingDecoded;
use crate::framing::FramingEncodedFinal;
use crate::thrift_protocol::ProtocolID;

pub struct SerializedMessage<'a, Name: ?Sized, Buffer> {
    pub protocol: ProtocolID,
    pub buffer: Buffer,
    pub method_name: &'a Name,
}

pub trait ContextStack {
    /// Type for method names
    type Name: ?Sized;
    /// Type for serialized read and write buffers
    type Frame: Framing;

    /// Called before the request is read.
    fn pre_read(&mut self) -> Result<(), Error>;

    /// Called before post_read after reading arguments (server) / after reading
    /// reply (client), with the actual (unparsed, serialized) data.
    fn on_read_data(
        &mut self,
        msg: SerializedMessage<Self::Name, FramingDecoded<Self::Frame>>,
    ) -> Result<(), Error>;

    /// Called after the request is read.
    fn post_read(&mut self, bytes: u32) -> Result<(), Error>;

    /// Called before a response is written.
    fn pre_write(&mut self) -> Result<(), Error>;

    /// Called before post_write, after serializing response (server) / after
    /// serializing request (client), with the actual (serialized) data.
    fn on_write_data(
        &mut self,
        msg: SerializedMessage<Self::Name, FramingEncodedFinal<Self::Frame>>,
    ) -> Result<(), Error>;

    /// Called after a response a written.
    fn post_write(&mut self, bytes: u32) -> Result<(), Error>;
}

pub struct DummyContextStack<Name: ?Sized, Frame> {
    _phantom: PhantomData<(Frame, Name)>,
}

impl<Name: ?Sized, Frame> DummyContextStack<Name, Frame> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData,
        }
    }
}

impl<Name: ?Sized, Frame: Framing> ContextStack for DummyContextStack<Name, Frame>
where
    FramingEncodedFinal<Frame>: BufExt,
{
    type Name = Name;
    type Frame = Frame;

    fn pre_read(&mut self) -> Result<(), Error> {
        Ok(())
    }

    fn on_read_data(
        &mut self,
        _msg: SerializedMessage<Self::Name, FramingDecoded<Self::Frame>>,
    ) -> Result<(), Error> {
        Ok(())
    }

    fn post_read(&mut self, _bytes: u32) -> Result<(), Error> {
        Ok(())
    }

    fn pre_write(&mut self) -> Result<(), Error> {
        Ok(())
    }

    fn on_write_data(
        &mut self,
        _msg: SerializedMessage<Self::Name, FramingEncodedFinal<Self::Frame>>,
    ) -> Result<(), Error> {
        Ok(())
    }

    fn post_write(&mut self, _bytes: u32) -> Result<(), Error> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use bytes::Bytes;

    use super::*;

    fn assert_context_stack(_: &impl ContextStack) {}

    #[test]
    fn check_unsized() {
        assert_context_stack(&DummyContextStack::<std::ffi::CStr, Bytes>::new());
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::hash::Hash;
use std::sync::Arc;

use bytes::Bytes;
use ordered_float::OrderedFloat;

use crate::protocol::should_break;
use crate::protocol::ProtocolReader;
use crate::Result;

// Read trait. Every type that needs to be deserialized will implement this trait.
pub trait Deserialize<P>
where
    P: ProtocolReader,
{
    fn read(p: &mut P) -> Result<Self>
    where
        Self: Sized;
}

impl<P, T> Deserialize<P> for Box<T>
where
    P: ProtocolReader,
    T: Deserialize<P>,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        T::read(p).map(Box::new)
    }
}

impl<P, T> Deserialize<P> for Arc<T>
where
    P: ProtocolReader,
    T: Deserialize<P>,
{
    fn read(p: &mut P) -> Result<Self> {
        T::read(p).map(Arc::new)
    }
}

impl<P> Deserialize<P> for ()
where
    P: ProtocolReader,
{
    #[inline]
    fn read(_p: &mut P) -> Result<Self> {
        Ok(())
    }
}

impl<P> Deserialize<P> for bool
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_bool()
    }
}

impl<P> Deserialize<P> for i8
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_byte()
    }
}

impl<P> Deserialize<P> for i16
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_i16()
    }
}

impl<P> Deserialize<P> for i32
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_i32()
    }
}

impl<P> Deserialize<P> for i64
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_i64()
    }
}

impl<P> Deserialize<P> for f64
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_double()
    }
}

impl<P> Deserialize<P> for f32
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_float()
    }
}

impl<P> Deserialize<P> for OrderedFloat<f64>
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_double().map(OrderedFloat)
    }
}

impl<P> Deserialize<P> for OrderedFloat<f32>
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_float().map(OrderedFloat)
    }
}

impl<P> Deserialize<P> for String
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_string()
    }
}

impl<P> Deserialize<P> for Bytes
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_binary()
    }
}

impl<P> Deserialize<P> for Vec<u8>
where
    P: ProtocolReader,
{
    #[inline]
    fn read(p: &mut P) -> Result<Self> {
        p.read_binary()
    }
}

impl<P, T> Deserialize<P> for BTreeSet<T>
where
    P: ProtocolReader,
    T: Deserialize<P> + Ord,
{
    fn read(p: &mut P) -> Result<Self> {
        let (_elem_ty, len) = p.read_set_begin()?;
        let mut bset = BTreeSet::new();

        if let Some(0) = len {
            return Ok(bset);
        }

        let mut idx = 0;
        loop {
            let more = p.read_set_value_begin()?;
            if !more {
                break;
            }
            let item = Deserialize::read(p)?;
            p.read_set_value_end()?;
            bset.insert(item);

            idx += 1;
            if should_break(len, more, idx) {
                break;
            }
        }
        p.read_set_end()?;
        Ok(bset)
    }
}

impl<P, T, S> Deserialize<P> for HashSet<T, S>
where
    P: ProtocolReader,
    T: Deserialize<P> + Hash + Eq,
    S: std::hash::BuildHasher + Default,
{
    fn read(p: &mut P) -> Result<Self> {
        let (_elem_ty, len) = p.read_set_begin()?;
        let mut hset =
            HashSet::with_capacity_and_hasher(len.unwrap_or_default(), Default::default());

        if let Some(0) = len {
            return Ok(hset);
        }

        let mut idx = 0;
        loop {
            let more = p.read_set_value_begin()?;
            if !more {
                break;
            }
            let item = Deserialize::read(p)?;
            p.read_set_value_end()?;
            hset.insert(item);

            idx += 1;
            if should_break(len, more, idx) {
                break;
            }
        }
        p.read_set_end()?;
        Ok(hset)
    }
}

impl<P, K, V> Deserialize<P> for BTreeMap<K, V>
where
    P: ProtocolReader,
    K: Deserialize<P> + Ord,
    V: Deserialize<P>,
{
    fn read(p: &mut P) -> Result<Self> {
        let (_key_ty, _val_ty, len) = p.read_map_begin()?;
        let mut btree = BTreeMap::new();

        if let Some(0) = len {
            return Ok(btree);
        }

        let mut idx = 0;
        loop {
            let more = p.read_map_key_begin()?;
            if !more {
                break;
            }
            let key = Deserialize::read(p)?;
            p.read_map_value_begin()?;
            let val = Deserialize::read(p)?;
            p.read_map_value_end()?;
            btree.insert(key, val);

            idx += 1;
            if should_break(len, more, idx) {
                break;
            }
        }
        p.read_map_end()?;
        Ok(btree)
    }
}

impl<P, K, V, S> Deserialize<P> for HashMap<K, V, S>
where
    P: ProtocolReader,
    K: Deserialize<P> + Hash + Eq,
    V: Deserialize<P>,
    S: std::hash::BuildHasher + Default,
{
    fn read(p: &mut P) -> Result<Self> {
        let (_key_ty, _val_ty, len) = p.read_map_begin()?;
        let mut hmap =
            HashMap::with_capacity_and_hasher(len.unwrap_or_default(), Default::default());

        if let Some(0) = len {
            return Ok(hmap);
        }

        let mut idx = 0;
        loop {
            let more = p.read_map_key_begin()?;
            if !more {
                break;
            }
            let key = Deserialize::read(p)?;
            p.read_map_value_begin()?;
            let val = Deserialize::read(p)?;
            p.read_map_value_end()?;
            hmap.insert(key, val);

            idx += 1;
            if should_break(len, more, idx) {
                break;
            }
        }
        p.read_map_end()?;
        Ok(hmap)
    }
}

impl<P, T> Deserialize<P> for Vec<T>
where
    P: ProtocolReader,
    T: Deserialize<P> + crate::ttype::GetTType, // GetTType just to exclude Vec<u8>
{
    /// Vec<T> is Thrift List type
    fn read(p: &mut P) -> Result<Self> {
        let (_elem_ty, len) = p.read_list_begin()?;
        let mut list = Vec::with_capacity(len.unwrap_or_default());

        if let Some(0) = len {
            return Ok(list);
        }

        let mut idx = 0;
        loop {
            let more = p.read_list_value_begin()?;
            if !more {
                break;
            }
            let item = Deserialize::read(p)?;
            p.read_list_value_end()?;
            list.push(item);

            idx += 1;
            if should_break(len, more, idx) {
                break;
            }
        }
        p.read_list_end()?;
        Ok(list)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use anyhow::Error;
use thiserror::Error;

use crate::ttype::TType;
use crate::ApplicationException;

#[derive(Debug, Error)]
#[allow(dead_code)]
pub enum ProtocolError {
    #[error("end of file reached")]
    EOF,
    #[error("bad thrift version specified")]
    BadVersion,
    #[error("missing protocol version")]
    ProtocolVersionMissing,
    #[error("protocol skip depth exceeded")]
    SkipDepthExceeded,
    #[error("streams unsupported")]
    StreamUnsupported,
    #[error("STOP outside of struct in skip")]
    UnexpectedStopInSkip,
    #[error("Invalid type in skip {0:?}")]
    InvalidTypeInSkip(TType),
    #[error("Unknown or invalid protocol ID {0}")]
    InvalidProtocolID(i16),
    #[error("Unknown or invalid TMessage type {0}")]
    InvalidMessageType(u32),
    #[error("Unknown or invalid type tag")]
    InvalidTypeTag,
    #[error("Unknown or invalid data length")]
    InvalidDataLength,
    #[error("Invalid value for type")]
    InvalidValue,
    #[error("Unexpected trailing data after the end of a value")]
    TrailingData,
    #[error("Unwanted extra union {0} field ty {1:?} id {2}")]
    UnwantedExtraUnionField(String, TType, i32),
}

/// Error value returned by functions that do not throw any user-defined exceptions.
#[derive(Debug, Error)]
pub enum NonthrowingFunctionError {
    #[error(transparent)]
    ApplicationException(#[from] ApplicationException),
    #[error(transparent)]
    ThriftError(#[from] Error),
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

//! Helpers for generated code using exceptions

use crate::thrift_protocol::MessageType;

/// This trait should be implemented for each individual exception type. It will typically be generated.
pub trait ExceptionInfo {
    /// Exception name
    fn exn_name(&self) -> &'static str {
        std::any::type_name::<Self>()
    }

    // Exception value
    fn exn_value(&self) -> String;

    /// Is a declared exception
    fn exn_is_declared(&self) -> bool;
}

/// An extension of ExceptionInfo that also includes successful results.
/// This is implemented on generated *Exn types.
pub trait ResultInfo: ExceptionInfo {
    fn result_type(&self) -> ResultType;
}

/// Classify a result from a specific method call.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum ResultType {
    /// A successful return
    Return,
    /// A declared exception
    Error,
    /// Some other exception (eg ApplicationException)
    Exception,
}

impl ResultType {
    pub fn message_type(&self) -> MessageType {
        match self {
            ResultType::Return | ResultType::Error => MessageType::Reply,
            ResultType::Exception => MessageType::Exception,
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

// Re-exports from dependencies which are only sometimes used by generated code.
// Having the generated code access these types through fbthrift::export lets us
// avoid touching hundreds of generated Cargo.toml files each time a code
// generator feature involves a new crate.

pub use ordered_float::OrderedFloat;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::io::Cursor;

use bytes::Bytes;
use bytes::BytesMut;

use crate::bufext::BufExt;
use crate::bufext::BufMutExt;

/// Helper type alias to get encoding buffer type
pub type FramingEncoded<F> = <F as Framing>::EncBuf;

/// Helper type alias to get the type of the finalized encoded buffer
pub type FramingEncodedFinal<F> = <<F as Framing>::EncBuf as BufMutExt>::Final;

/// Helper type alias to get the buffer to use as input to decoding
pub type FramingDecoded<F> = <F as Framing>::DecBuf;

/// Trait describing the in-memory frames the transport uses for Protocol messages.
pub trait Framing {
    /// Buffer type we encode into
    type EncBuf: BufMutExt + Send + 'static;

    /// Buffer type we decode from
    type DecBuf: BufExt + Send + 'static;

    /// Allocate a new encoding buffer with a given capacity
    /// FIXME: need &self?
    fn enc_with_capacity(cap: usize) -> Self::EncBuf;
}

impl Framing for Bytes {
    type EncBuf = BytesMut;
    type DecBuf = Cursor<Bytes>;

    fn enc_with_capacity(cap: usize) -> Self::EncBuf {
        BytesMut::with_capacity(cap)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

// Functions only intended to be called from code generated by the Rust Thrift
// code generator.

use std::ffi::CStr;
use std::fmt;
use std::future::Future;
use std::pin::Pin;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use bytes::Buf;
use futures::future::FutureExt;

use crate::serialize;
use crate::ApplicationException;
use crate::BufExt;
use crate::ContextStack;
use crate::Deserialize;
use crate::MessageType;
use crate::Protocol;
use crate::ProtocolEncodedFinal;
use crate::ProtocolReader;
use crate::ProtocolWriter;
use crate::RequestContext;
use crate::ResultInfo;
use crate::ResultType;
use crate::Serialize;
use crate::SerializedMessage;

// Note: `variants_by_number` must be sorted by the i32 values.
pub fn enum_display(
    variants_by_number: &[(&str, i32)],
    formatter: &mut fmt::Formatter,
    number: i32,
) -> fmt::Result {
    match variants_by_number.binary_search_by_key(&number, |entry| entry.1) {
        Ok(i) => formatter.write_str(variants_by_number[i].0),
        Err(_) => write!(formatter, "{}", number),
    }
}

// Note: `variants_by_name` must be sorted by the string values.
pub fn enum_from_str(
    variants_by_name: &[(&str, i32)],
    value: &str,
    type_name: &'static str,
) -> Result<i32> {
    match variants_by_name.binary_search_by_key(&value, |entry| entry.0) {
        Ok(i) => Ok(variants_by_name[i].1),
        Err(_) => bail!("Unable to parse {} as {}", value, type_name),
    }
}

pub fn type_name_of_val<T>(_: &T) -> &'static str {
    std::any::type_name::<T>()
}

pub fn buf_len<B: Buf>(b: &B) -> anyhow::Result<u32> {
    let length: usize = b.remaining();
    let length: u32 = length.try_into().with_context(|| {
        format!("Unable to report a buffer length of {length} bytes as a `u32`")
    })?;
    Ok(length)
}

/// Serialize a result as encoded into a generated *Exn type, wrapped in an envelope.
pub fn serialize_result_envelope<P, CTXT, RES>(
    name: &str,
    name_cstr: &<CTXT::ContextStack as ContextStack>::Name,
    seqid: u32,
    rctxt: &CTXT,
    ctx_stack: &mut CTXT::ContextStack,
    res: RES,
) -> anyhow::Result<ProtocolEncodedFinal<P>>
where
    P: Protocol,
    RES: ResultInfo + Serialize<P::Sizer> + Serialize<P::Serializer>,
    CTXT: RequestContext<Name = CStr>,
    <CTXT as RequestContext>::ContextStack: ContextStack<Frame = P::Frame>,
    ProtocolEncodedFinal<P>: Clone + Buf + BufExt,
{
    let res_type = res.result_type();

    if matches!(res_type, ResultType::Error | ResultType::Exception) {
        assert_eq!(res.exn_is_declared(), res_type == ResultType::Error);

        rctxt.set_user_exception_header(res.exn_name(), &res.exn_value())?;
    }

    ctx_stack.pre_write()?;
    let envelope = serialize!(P, |p| {
        p.write_message_begin(name, res_type.message_type(), seqid);
        res.write(p);
        p.write_message_end();
    });

    ctx_stack.on_write_data(SerializedMessage {
        protocol: P::PROTOCOL_ID,
        method_name: name_cstr,
        buffer: envelope.clone(),
    })?;
    let bytes_written = buf_len(&envelope)?;
    ctx_stack.post_write(bytes_written)?;

    Ok(envelope)
}

pub fn serialize_stream_item<P, RES>(res: RES) -> anyhow::Result<ProtocolEncodedFinal<P>>
where
    P: Protocol,
    RES: ResultInfo + Serialize<P::Sizer> + Serialize<P::Serializer>,
{
    Ok(serialize!(P, |p| {
        res.write(p);
    }))
}

/// Serialize a request with envelope.
pub fn serialize_request_envelope<P, ARGS>(
    name: &str,
    args: &ARGS,
) -> anyhow::Result<ProtocolEncodedFinal<P>>
where
    P: Protocol,
    ARGS: Serialize<P::Sizer> + Serialize<P::Serializer>,
{
    let envelope = serialize!(P, |p| {
        // Note: we send a 0 message sequence ID from clients because
        // this field should not be used by the server (except for some
        // language implementations).
        p.write_message_begin(name, MessageType::Call, 0);
        args.write(p);
        p.write_message_end();
    });

    Ok(envelope)
}

/// Deserialize a client response. This deserializes the envelope then
/// deserializes either a reply or an ApplicationException.
pub fn deserialize_response_envelope<P, T>(
    de: &mut P::Deserializer,
) -> anyhow::Result<Result<T, ApplicationException>>
where
    P: Protocol,
    T: Deserialize<P::Deserializer>,
{
    let (_, message_type, _) = de.read_message_begin(|_| ())?;

    let res = match message_type {
        MessageType::Reply => Ok(T::read(de)?),
        MessageType::Exception => Err(ApplicationException::read(de)?),
        MessageType::Call | MessageType::Oneway | MessageType::InvalidMessageType => {
            bail!("Unwanted message type `{:?}`", message_type)
        }
    };

    de.read_message_end()?;

    Ok(res)
}

/// Abstract spawning some potentially CPU-heavy work onto a CPU thread
pub trait Spawner: 'static {
    fn spawn<F, R>(func: F) -> Pin<Box<dyn Future<Output = R> + Send>>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static;
}

/// No-op implementation of Spawner - just run on current thread
pub struct NoopSpawner;
impl Spawner for NoopSpawner {
    #[inline]
    fn spawn<F, R>(func: F) -> Pin<Box<dyn Future<Output = R> + Send>>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        async { func() }.boxed()
    }
}

pub async fn async_deserialize_response_envelope<P, T, S>(
    de: P::Deserializer,
) -> anyhow::Result<(Result<T, ApplicationException>, P::Deserializer)>
where
    P: Protocol,
    P::Deserializer: Send,
    T: Deserialize<P::Deserializer> + Send + 'static,
    S: Spawner,
{
    S::spawn(move || {
        let mut de = de;
        let res = deserialize_response_envelope::<P, T>(&mut de);
        res.map(|res| (res, de))
    })
    .await
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

#![recursion_limit = "1024"]
#![deny(warnings)]

use std::i32;

macro_rules! bail_err {
    ($e:expr) => {
        return Err(From::from($e))
    };
}

macro_rules! ensure_err {
    ($cond:expr, $e:expr) => {
        if !$cond {
            bail_err!($e);
        }
    };
}

use anyhow::Result;

#[macro_use]
pub mod protocol;

pub mod adapter;
pub mod application_exception;
pub mod binary_protocol;
pub mod binary_type;
pub mod builtin_types;
pub mod compact_protocol;
pub mod context_stack;
pub mod deserialize;
pub mod export;
pub mod framing;
pub mod metadata;
pub mod processor;
pub mod request_context;
pub mod serialize;
pub mod simplejson_protocol;
pub mod thrift_protocol;
pub mod ttype;
pub mod uri;

mod bufext;
mod client;
mod errors;
mod exceptions;
mod varint;

#[cfg(test)]
mod tests;

#[doc(hidden)]
pub mod help;

pub mod types {
    // Define ApplicationException as if it were a normal generated type to make things simpler
    // for codegen.
    pub use crate::application_exception::ApplicationException;
}

pub use crate::application_exception::ApplicationException;
pub use crate::application_exception::ApplicationExceptionErrorCode;
pub use crate::binary_protocol::BinaryProtocol;
pub use crate::bufext::BufExt;
pub use crate::bufext::BufMutExt;
pub use crate::bufext::DeserializeSource;
pub use crate::client::ClientFactory;
pub use crate::client::ClientStreamElement;
pub use crate::client::Transport;
pub use crate::compact_protocol::CompactProtocol;
pub use crate::context_stack::ContextStack;
pub use crate::context_stack::DummyContextStack;
pub use crate::context_stack::SerializedMessage;
pub use crate::deserialize::Deserialize;
pub use crate::errors::NonthrowingFunctionError;
pub use crate::errors::ProtocolError;
pub use crate::exceptions::ExceptionInfo;
pub use crate::exceptions::ResultInfo;
pub use crate::exceptions::ResultType;
pub use crate::framing::Framing;
pub use crate::framing::FramingDecoded;
pub use crate::framing::FramingEncoded;
pub use crate::framing::FramingEncodedFinal;
pub use crate::help::NoopSpawner;
pub use crate::processor::NullServiceProcessor;
pub use crate::processor::ReplyState;
pub use crate::processor::SerializedStreamElement;
pub use crate::processor::ServiceProcessor;
pub use crate::processor::ThriftService;
pub use crate::protocol::Field;
pub use crate::protocol::Protocol;
pub use crate::protocol::ProtocolDecoded;
pub use crate::protocol::ProtocolEncoded;
pub use crate::protocol::ProtocolEncodedFinal;
pub use crate::protocol::ProtocolReader;
pub use crate::protocol::ProtocolWriter;
pub use crate::request_context::DummyRequestContext;
pub use crate::request_context::RequestContext;
pub use crate::serialize::Serialize;
pub use crate::simplejson_protocol::SimpleJsonProtocol;
pub use crate::thrift_protocol::MessageType;
pub use crate::thrift_protocol::ProtocolID;
pub use crate::ttype::GetTType;
pub use crate::ttype::TType;
pub use crate::uri::GetUri;

pub trait ThriftEnum: Sized {
    fn enumerate() -> &'static [(Self, &'static str)];

    fn variants() -> &'static [&'static str];

    fn variant_values() -> &'static [Self];
}

/// Set the default ID's for unknown exceptions and fields.
/// When reading off the wire, these default values will be
/// overridden with the unrecognized id (which must be nonnegative).
// ---
// Keep in sync with the UNKNOWN_ID constant in //common/rust/thrift/ast.
pub const __UNKNOWN_ID: i32 = i32::MIN;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

pub trait ThriftAnnotations: 'static {
    /// Returns the structured annotation `T` that is defined on this Thrift type, if the
    /// annotation exists.
    fn get_structured_annotation<T: Sized + 'static>() -> Option<T> {
        None
    }

    /// Returns the structured annotation `T` that is defined on the given field with id `field_id`
    /// on this Thrift type.
    ///
    /// Returns `None` if the field ID does not exist or if the given `T` does not exist as an
    /// annotation on the given field.
    ///
    /// For example, for the given Thrift file:
    /// ```thrift
    /// struct FieldAnnotation {
    ///     1: string payload;
    /// }
    ///
    /// struct Foo {
    ///     @FieldAnnotation{ payload = "hello world" }
    ///     1: string username;
    /// }
    /// ```
    ///
    /// The call
    /// ```ignore
    /// <Foo as ThriftAnnotations>::get_field_structured_annotation::<FieldAnnotation>(1)
    /// ```
    /// will return
    /// ```ignore
    /// Some(FieldAnnotation {
    ///     payload: "hello world".to_string(),
    /// })
    /// ```
    fn get_field_structured_annotation<T: Sized + 'static>(_field_id: i16) -> Option<T> {
        None
    }
}

/// Identical to [ThriftAnnotations] but is implemented on the unit type (), for which
/// `get_structured_annotation` and `get_field_structured_annotation` will always return `None`.
///
/// This allows a method to take `<T: MaybeThriftAnnotations>` to express the logic of "give me a
/// privacy aware Thrift struct OR a 'void' type to indicate no type".
pub trait MaybeThriftAnnotations: 'static {
    fn get_structured_annotation<T: Sized + 'static>() -> Option<T> {
        None
    }

    fn get_field_structured_annotation<T: Sized + 'static>(_field_id: i16) -> Option<T> {
        None
    }
}

impl<S: ThriftAnnotations> MaybeThriftAnnotations for S {
    fn get_structured_annotation<T: Sized + 'static>() -> Option<T> {
        <S as ThriftAnnotations>::get_structured_annotation::<T>()
    }

    fn get_field_structured_annotation<T: Sized + 'static>(field_id: i16) -> Option<T> {
        <S as ThriftAnnotations>::get_field_structured_annotation::<T>(field_id)
    }
}

impl MaybeThriftAnnotations for () {}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::ffi::CStr;
use std::marker::PhantomData;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;

use crate::application_exception::ApplicationException;
use crate::application_exception::ApplicationExceptionErrorCode;
use crate::context_stack::ContextStack;
use crate::exceptions::ExceptionInfo;
use crate::exceptions::ResultInfo;
use crate::framing::Framing;
use crate::framing::FramingDecoded;
use crate::framing::FramingEncodedFinal;
use crate::protocol::Protocol;
use crate::protocol::ProtocolDecoded;
use crate::protocol::ProtocolReader;
use crate::protocol::ProtocolWriter;
use crate::request_context::RequestContext;
use crate::serialize::Serialize;
use crate::thrift_protocol::ProtocolID;
use crate::ttype::TType;

pub enum SerializedStreamElement<Payload> {
    /// A normal stream response, without any error. Contains the serialized response.
    Success(Payload),
    /// Contains the serialized declared exception.
    DeclaredException(Payload),
    /// Contains the application exception.
    ApplicationException(ApplicationException),
    /// The serialization failed. Contains the error.
    SerializationError(Error),
}

pub trait ReplyState<F>
where
    F: Framing,
{
    type RequestContext;

    fn send_reply(&mut self, reply: FramingEncodedFinal<F>);
    fn send_stream_reply(
        &mut self,
        response: FramingEncodedFinal<F>,
        stream: Option<BoxStream<'static, SerializedStreamElement<FramingEncodedFinal<F>>>>,
        protocol_id: ProtocolID,
    ) -> Result<()>;
    fn set_interaction_processor(
        &mut self,
        _processor: Arc<
            dyn ThriftService<
                    F,
                    Handler = (),
                    RequestContext = Self::RequestContext,
                    ReplyState = Self,
                > + ::std::marker::Send
                + 'static,
        >,
    ) -> Result<()> {
        bail!("Thrift server does not support interactions");
    }
}

#[async_trait]
pub trait ThriftService<F>: Send + Sync + 'static
where
    F: Framing + Send + 'static,
{
    type Handler;
    type RequestContext;
    type ReplyState;

    async fn call(
        &self,
        req: FramingDecoded<F>,
        req_ctxt: &Self::RequestContext,
        reply_state: Arc<Mutex<Self::ReplyState>>,
    ) -> Result<(), Error>;

    fn create_interaction(
        &self,
        _name: &str,
    ) -> ::anyhow::Result<
        Arc<
            dyn ThriftService<
                    F,
                    Handler = (),
                    RequestContext = Self::RequestContext,
                    ReplyState = Self::ReplyState,
                > + ::std::marker::Send
                + 'static,
        >,
    > {
        bail!("Thrift server does not support interactions");
    }

    /// Returns function names this thrift service is able to handle, similar
    /// to the keys of C++'s createMethodMetadata().
    ///
    /// Return value includes inherited functions from parent thrift services,
    /// and interactions' functions.
    fn get_method_names(&self) -> &'static [&'static str];

    /// Applies to interactions only
    ///
    /// Termination callback is invoked immediately as soon as the client's
    /// termination signal is received. This differs to the interaction service
    /// being dropped, which only happens when all outstanding requests and
    /// streams have been completed. This is not invoked if the connection
    /// closes without the signal being received.
    async fn on_termination(&self);
}

#[async_trait]
impl<F, T> ThriftService<F> for Box<T>
where
    T: ThriftService<F> + Send + Sync + ?Sized,
    F: Framing + Send + 'static,
    T::RequestContext: Send + Sync + 'static,
    T::ReplyState: Send + Sync + 'static,
{
    type Handler = T::Handler;
    type RequestContext = T::RequestContext;
    type ReplyState = T::ReplyState;

    async fn call(
        &self,
        req: FramingDecoded<F>,
        req_ctxt: &Self::RequestContext,
        reply_state: Arc<Mutex<Self::ReplyState>>,
    ) -> Result<(), Error> {
        (**self).call(req, req_ctxt, reply_state).await
    }

    fn create_interaction(
        &self,
        name: &str,
    ) -> ::anyhow::Result<
        Arc<
            dyn ThriftService<
                    F,
                    Handler = (),
                    RequestContext = Self::RequestContext,
                    ReplyState = Self::ReplyState,
                > + ::std::marker::Send
                + 'static,
        >,
    > {
        (**self).create_interaction(name)
    }

    fn get_method_names(&self) -> &'static [&'static str] {
        (**self).get_method_names()
    }

    async fn on_termination(&self) {
        (**self).on_termination().await
    }
}

#[async_trait]
impl<F, T> ThriftService<F> for Arc<T>
where
    T: ThriftService<F> + Send + Sync + ?Sized,
    F: Framing + Send + 'static,
    T::RequestContext: Send + Sync + 'static,
    T::ReplyState: Send + Sync + 'static,
{
    type Handler = T::Handler;
    type RequestContext = T::RequestContext;
    type ReplyState = T::ReplyState;

    async fn call(
        &self,
        req: FramingDecoded<F>,
        req_ctxt: &Self::RequestContext,
        reply_state: Arc<Mutex<Self::ReplyState>>,
    ) -> Result<(), Error> {
        (**self).call(req, req_ctxt, reply_state).await
    }

    fn create_interaction(
        &self,
        name: &str,
    ) -> ::anyhow::Result<
        Arc<
            dyn ThriftService<
                    F,
                    Handler = (),
                    RequestContext = Self::RequestContext,
                    ReplyState = Self::ReplyState,
                > + ::std::marker::Send
                + 'static,
        >,
    > {
        (**self).create_interaction(name)
    }

    fn get_method_names(&self) -> &'static [&'static str] {
        (**self).get_method_names()
    }

    async fn on_termination(&self) {
        (**self).on_termination().await
    }
}

/// Trait implemented by a generated type to implement a service.
#[async_trait]
pub trait ServiceProcessor<P>
where
    P: Protocol,
{
    type RequestContext;
    type ReplyState;

    /// Given a method name, return a reference to the processor for that index.
    fn method_idx(&self, name: &[u8]) -> Result<usize, ApplicationException>;

    /// Given a method index and the remains of the message input, get a future
    /// for the result of the method. This will only be called if the corresponding
    /// `method_idx()` returns an (index, ServiceProcessor) tuple.
    /// `frame` is a reference to the frame containing the request.
    /// `request` is a deserializer instance set up to decode the request.
    async fn handle_method(
        &self,
        idx: usize,
        //frame: &P::Frame,
        d: &mut P::Deserializer,
        req: ProtocolDecoded<P>,
        req_ctxt: &Self::RequestContext,
        reply_state: Arc<Mutex<Self::ReplyState>>,
        seqid: u32,
    ) -> Result<(), Error>;

    /// Given a method name, return a reference to the interaction creation fn for that index
    fn create_interaction_idx(&self, _name: &str) -> ::anyhow::Result<::std::primitive::usize> {
        bail!("Processor does not support interactions");
    }

    /// Given a creation method index, it produces a fresh interaction processor
    fn handle_create_interaction(
        &self,
        _idx: ::std::primitive::usize,
    ) -> ::anyhow::Result<
        Arc<
            dyn ThriftService<
                    P::Frame,
                    Handler = (),
                    RequestContext = Self::RequestContext,
                    ReplyState = Self::ReplyState,
                > + ::std::marker::Send
                + 'static,
        >,
    > {
        bail!("Processor does not support interactions");
    }

    /// See [ThriftService::on_termination] docs
    async fn handle_on_termination(&self);
}

/// Null processor which implements no methods - it acts as the super for any service
/// which has no super-service.
#[derive(Debug, Clone)]
pub struct NullServiceProcessor<P, R, RS> {
    _phantom: PhantomData<(P, R, RS)>,
}

impl<P, R, RS> NullServiceProcessor<P, R, RS> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData,
        }
    }
}

impl<P, R, RS> Default for NullServiceProcessor<P, R, RS> {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl<P, R, RS> ServiceProcessor<P> for NullServiceProcessor<P, R, RS>
where
    P: Protocol + Sync,
    P::Deserializer: Send,
    R: Sync,
    RS: Sync + Send,
{
    type RequestContext = R;
    type ReplyState = RS;

    #[inline]
    fn method_idx(&self, name: &[u8]) -> Result<usize, ApplicationException> {
        Err(ApplicationException::new(
            ApplicationExceptionErrorCode::UnknownMethod,
            format!("Unknown method {}", String::from_utf8_lossy(name)),
        ))
    }

    async fn handle_method(
        &self,
        _idx: usize,
        //_frame: &P::Frame,
        _d: &mut P::Deserializer,
        _req: ProtocolDecoded<P>,
        _req_ctxt: &R,
        _reply_state: Arc<Mutex<RS>>,
        _seqid: u32,
    ) -> Result<(), Error> {
        // Should never be called since method_idx() always returns an error
        unimplemented!("NullServiceProcessor implements no methods")
    }

    fn create_interaction_idx(&self, name: &str) -> ::anyhow::Result<::std::primitive::usize> {
        bail!("Unknown interaction {}", name);
    }

    fn handle_create_interaction(
        &self,
        _idx: ::std::primitive::usize,
    ) -> ::anyhow::Result<
        Arc<
            dyn ThriftService<
                    P::Frame,
                    Handler = (),
                    RequestContext = Self::RequestContext,
                    ReplyState = Self::ReplyState,
                > + ::std::marker::Send
                + 'static,
        >,
    > {
        unimplemented!("NullServiceProcessor implements no interactions")
    }

    async fn handle_on_termination(&self) {}
}

#[async_trait]
impl<P, R, RS> ThriftService<P::Frame> for NullServiceProcessor<P, R, RS>
where
    P: Protocol + Send + Sync + 'static,
    P::Frame: Send + 'static,
    R: RequestContext<Name = CStr> + Send + Sync + 'static,
    R::ContextStack: ContextStack<Name = CStr>,
    RS: ReplyState<P::Frame> + Send + Sync + 'static,
{
    type Handler = ();
    type RequestContext = R;
    type ReplyState = RS;

    async fn call(
        &self,
        req: ProtocolDecoded<P>,
        rctxt: &R,
        reply_state: Arc<Mutex<RS>>,
    ) -> Result<(), Error> {
        let mut p = P::deserializer(req);

        const SERVICE_NAME: &str = "NullService";
        let ((name, ae), _, seqid) = p.read_message_begin(|name| {
            let name = String::from_utf8_lossy(name).into_owned();
            let ae = ApplicationException::unimplemented_method(SERVICE_NAME, &name);
            (name, ae)
        })?;

        p.skip(TType::Struct)?;
        p.read_message_end()?;

        rctxt.set_user_exception_header(ae.exn_name(), &ae.exn_value())?;
        let res = serialize!(P, |p| {
            p.write_message_begin(&name, ae.result_type().message_type(), seqid);
            ae.write(p);
            p.write_message_end();
        });
        reply_state.lock().unwrap().send_reply(res);
        Ok(())
    }

    fn create_interaction(
        &self,
        name: &str,
    ) -> ::anyhow::Result<
        Arc<
            dyn ThriftService<
                    P::Frame,
                    Handler = Self::Handler,
                    RequestContext = Self::RequestContext,
                    ReplyState = Self::ReplyState,
                > + ::std::marker::Send
                + 'static,
        >,
    > {
        bail!("Unimplemented interaction {}", name);
    }

    fn get_method_names(&self) -> &'static [&'static str] {
        &[]
    }

    async fn on_termination(&self) {}
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use crate::binary_type::CopyFromBuf;
use crate::binary_type::Discard;
use crate::bufext::BufMutExt;
use crate::errors::ProtocolError;
use crate::framing::Framing;
use crate::framing::FramingDecoded;
use crate::framing::FramingEncoded;
use crate::framing::FramingEncodedFinal;
use crate::thrift_protocol::MessageType;
use crate::thrift_protocol::ProtocolID;
use crate::ttype::TType;
use crate::Result;

/// The maximum recursive depth the skip() function will traverse
pub const DEFAULT_RECURSION_DEPTH: i32 = 64;

/// Helper type alias to get the pre-finalization encoded type of a protocol frame.
pub type ProtocolEncoded<P> = FramingEncoded<<P as Protocol>::Frame>;

/// Helper type alias to get the final encoded type of a protocol frame
pub type ProtocolEncodedFinal<P> = FramingEncodedFinal<<P as Protocol>::Frame>;

/// Helper type alias to get the buffer type for a frame to be decoded.
pub type ProtocolDecoded<P> = FramingDecoded<<P as Protocol>::Frame>;

/// An instance of Protocol glues a Framing implementation to a serializer
/// (ProtocolWriter) and deserializer (ProtocolReader). It constructs, as
/// needed, a serializer to construct a frame with a given protocol, or a
/// deserializer from a frame into a stream of deserialized objects.
pub trait Protocol: 'static {
    /// Type of the framing implementation
    type Frame: Framing;

    /// Compute the size of a frame for a given protocol. This can be exact or too large, but
    /// must not be too small.
    type Sizer: ProtocolWriter<Final = usize>;

    /// Serialize into a buffer. The buffer is allocated with the size computed by Sizer, so
    /// it must be large enough.
    type Serializer: ProtocolWriter<Final = <<Self::Frame as Framing>::EncBuf as BufMutExt>::Final>;

    /// Set up a deserializer from a frame's buffer.
    type Deserializer: ProtocolReader;

    const PROTOCOL_ID: ProtocolID;

    fn serializer<SZ, SER>(sz: SZ, ser: SER) -> <Self::Serializer as ProtocolWriter>::Final
    where
        SZ: FnOnce(&mut Self::Sizer),
        SER: FnOnce(&mut Self::Serializer);

    fn deserializer(buf: <Self::Frame as Framing>::DecBuf) -> Self::Deserializer;

    fn into_buffer(_: Self::Deserializer) -> <Self::Frame as Framing>::DecBuf;
}

fn skip_inner<P: ProtocolReader + ?Sized>(
    p: &mut P,
    field_type: TType,
    max_depth: i32,
) -> Result<()> {
    if max_depth <= 0 {
        bail_err!(ProtocolError::SkipDepthExceeded)
    }

    match field_type {
        TType::Void => {}
        TType::Bool => {
            p.read_bool()?;
        }
        TType::Byte => {
            p.read_byte()?;
        }
        TType::I16 => {
            p.read_i16()?;
        }
        TType::I32 => {
            p.read_i32()?;
        }
        TType::I64 => {
            p.read_i64()?;
        }
        TType::Double => {
            p.read_double()?;
        }
        TType::Float => {
            p.read_float()?;
        }
        TType::Struct => {
            p.read_struct_begin(|_| ())?;
            loop {
                let fields = &[];
                let (_, type_id, _) = p.read_field_begin(|_| (), fields)?;
                if type_id == TType::Stop {
                    break;
                }
                skip_inner(p, type_id, max_depth - 1)?;
                p.read_field_end()?;
            }
            p.read_struct_end()?;
        }
        TType::Map => {
            let (key_type, value_type, len) = p.read_map_begin()?;
            if len != Some(0) {
                let mut idx = 0;
                loop {
                    let more = p.read_map_key_begin()?;
                    if !more {
                        break;
                    }
                    skip_inner(p, key_type, max_depth - 1)?;
                    p.read_map_value_begin()?;
                    skip_inner(p, value_type, max_depth)?;
                    p.read_map_value_end()?;

                    idx += 1;
                    if should_break(len, more, idx) {
                        break;
                    }
                }
            }
            p.read_map_end()?;
        }
        TType::Set => {
            let (elem_type, len) = p.read_set_begin()?;
            if len != Some(0) {
                let mut idx = 0;
                loop {
                    let more = p.read_set_value_begin()?;
                    if !more {
                        break;
                    }
                    skip_inner(p, elem_type, max_depth - 1)?;
                    p.read_set_value_end()?;

                    idx += 1;
                    if should_break(len, more, idx) {
                        break;
                    }
                }
            }
            p.read_set_end()?;
        }
        TType::List => {
            let (elem_type, len) = p.read_list_begin()?;
            if len != Some(0) {
                let mut idx = 0;
                loop {
                    let more = p.read_list_value_begin()?;
                    if !more {
                        break;
                    }
                    skip_inner(p, elem_type, max_depth - 1)?;
                    p.read_list_value_end()?;

                    idx += 1;
                    if should_break(len, more, idx) {
                        break;
                    }
                }
            }
            p.read_list_end()?;
        }
        TType::UTF8 => {
            p.read_string()?;
        }
        TType::UTF16 => {
            p.read_string()?;
        }
        TType::String => {
            p.read_binary::<Discard>()?;
        }
        TType::Stream => bail_err!(ProtocolError::StreamUnsupported),
        TType::Stop => bail_err!(ProtocolError::UnexpectedStopInSkip),
    };
    Ok(())
}

/// Trait for emitting a frame formatted in a given protocol
pub trait ProtocolWriter {
    type Final;

    fn write_message_begin(&mut self, name: &str, type_id: MessageType, seqid: u32);
    fn write_message_end(&mut self);
    fn write_struct_begin(&mut self, name: &str);
    fn write_struct_end(&mut self);
    fn write_field_begin(&mut self, name: &str, type_id: TType, id: i16);
    fn write_field_end(&mut self);
    fn write_field_stop(&mut self);
    fn write_map_begin(&mut self, key_type: TType, value_type: TType, size: usize);
    fn write_map_key_begin(&mut self);
    fn write_map_value_begin(&mut self);
    fn write_map_end(&mut self);
    fn write_list_begin(&mut self, elem_type: TType, size: usize);
    fn write_list_value_begin(&mut self);
    fn write_list_end(&mut self);
    fn write_set_begin(&mut self, elem_type: TType, size: usize);
    fn write_set_value_begin(&mut self);
    fn write_set_end(&mut self);
    fn write_bool(&mut self, value: bool);
    fn write_byte(&mut self, value: i8);
    fn write_i16(&mut self, value: i16);
    fn write_i32(&mut self, value: i32);
    fn write_i64(&mut self, value: i64);
    fn write_double(&mut self, value: f64);
    fn write_float(&mut self, value: f32);
    fn write_string(&mut self, value: &str);
    fn write_binary(&mut self, value: &[u8]);

    fn finish(self) -> Self::Final;
}

/// Trait for decoding a frame in a given protocol
pub trait ProtocolReader {
    fn read_message_begin<F, T>(&mut self, method: F) -> Result<(T, MessageType, u32)>
    where
        F: FnOnce(&[u8]) -> T;
    fn read_message_end(&mut self) -> Result<()>;
    fn read_struct_begin<F, T>(&mut self, strukt: F) -> Result<T>
    where
        F: FnOnce(&[u8]) -> T;
    fn read_struct_end(&mut self) -> Result<()>;
    fn read_field_begin<F, T>(&mut self, field: F, fields: &[Field]) -> Result<(T, TType, i16)>
    where
        F: FnOnce(&[u8]) -> T;
    fn read_field_end(&mut self) -> Result<()>;
    fn read_map_begin(&mut self) -> Result<(TType, TType, Option<usize>)>;
    fn read_map_key_begin(&mut self) -> Result<bool>;
    fn read_map_value_begin(&mut self) -> Result<()>;
    fn read_map_value_end(&mut self) -> Result<()>;
    fn read_map_end(&mut self) -> Result<()>;
    fn read_list_begin(&mut self) -> Result<(TType, Option<usize>)>;
    fn read_list_value_begin(&mut self) -> Result<bool>;
    fn read_list_value_end(&mut self) -> Result<()>;
    fn read_list_end(&mut self) -> Result<()>;
    fn read_set_begin(&mut self) -> Result<(TType, Option<usize>)>;
    fn read_set_value_begin(&mut self) -> Result<bool>;
    fn read_set_value_end(&mut self) -> Result<()>;
    fn read_set_end(&mut self) -> Result<()>;
    fn read_bool(&mut self) -> Result<bool>;
    fn read_byte(&mut self) -> Result<i8>;
    fn read_i16(&mut self) -> Result<i16>;
    fn read_i32(&mut self) -> Result<i32>;
    fn read_i64(&mut self) -> Result<i64>;
    fn read_double(&mut self) -> Result<f64>;
    fn read_float(&mut self) -> Result<f32>;
    fn read_string(&mut self) -> Result<String>;
    fn read_binary<V: CopyFromBuf>(&mut self) -> Result<V>;

    /// Skip over the next data element from the provided input Protocol object
    fn skip(&mut self, field_type: TType) -> Result<()> {
        skip_inner(self, field_type, DEFAULT_RECURSION_DEPTH)
    }
}

pub fn should_break(len: Option<usize>, more: bool, idx: usize) -> bool {
    match (len, more) {
        (Some(real_length), _) => idx >= real_length,
        (None, true) => false,
        (None, false) => true,
    }
}

/// Protocol::serializer wants the same serializer closure passed twice so
/// that it can specialize it for two types: once to compute the size, and a second time to
/// actually serialize the content. This macro helps by taking the factory type and
/// applying the serializer expression twice.
#[macro_export]
macro_rules! serialize {
    ($protocol:ty, $serializer:expr) => {
        <$protocol as $crate::Protocol>::serializer($serializer, $serializer)
    };
}

pub struct Field {
    pub name: &'static str,
    pub ttype: TType,
    pub id: i16,
}

impl Field {
    pub const fn new(name: &'static str, ttype: TType, id: i16) -> Self {
        Field { name, ttype, id }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::marker::PhantomData;

use anyhow::Error;

use crate::bufext::BufExt;
use crate::context_stack::ContextStack;
use crate::context_stack::DummyContextStack;
use crate::framing::Framing;
use crate::framing::FramingEncodedFinal;

pub trait RequestContext {
    type ContextStack: ContextStack;
    type Name: ?Sized;

    // This should return a ContextStack with a lifetime bounded by &self and the
    // names, but that needs GATs, so require 'static for now.
    fn get_context_stack(
        &self,
        service_name: &'static Self::Name,
        fn_name: &'static Self::Name,
    ) -> Result<Self::ContextStack, Error>;

    fn set_user_exception_header(&self, ex_type: &str, ex_reason: &str) -> Result<(), Error>;
}

pub struct DummyRequestContext<Name: ?Sized, Frame> {
    _phantom: PhantomData<(Frame, Name)>,
}

impl<Name: ?Sized, Frame> DummyRequestContext<Name, Frame> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData,
        }
    }
}

impl<Name: ?Sized, Frame: Framing> RequestContext for DummyRequestContext<Name, Frame>
where
    FramingEncodedFinal<Frame>: BufExt,
{
    type ContextStack = crate::context_stack::DummyContextStack<Name, Frame>;
    type Name = Name;

    fn get_context_stack(
        &self,
        _service_name: &Self::Name,
        _fn_name: &Self::Name,
    ) -> Result<Self::ContextStack, Error> {
        Ok(DummyContextStack::new())
    }

    fn set_user_exception_header(&self, _ex_type: &str, _ex_reason: &str) -> Result<(), Error> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use bytes::Bytes;

    use super::*;

    fn assert_context_stack(_: &impl RequestContext) {}

    #[test]
    fn check_unsized() {
        assert_context_stack(&DummyRequestContext::<std::ffi::CStr, Bytes>::new());
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::hash::Hash;
use std::sync::Arc;

use bytes::Bytes;
use ordered_float::OrderedFloat;

use crate::protocol::ProtocolWriter;
use crate::ttype::GetTType;

// Write trait. Every type that needs to be serialized will implement this trait.
pub trait Serialize<P>
where
    P: ProtocolWriter,
{
    fn write(&self, p: &mut P);
}

impl<P, T> Serialize<P> for &T
where
    P: ProtocolWriter,
    T: ?Sized + Serialize<P>,
{
    fn write(&self, p: &mut P) {
        (**self).write(p);
    }
}

impl<P, T> Serialize<P> for Box<T>
where
    P: ProtocolWriter,
    T: Serialize<P>,
{
    #[inline]
    fn write(&self, p: &mut P) {
        self.as_ref().write(p)
    }
}

impl<P, T> Serialize<P> for Arc<T>
where
    P: ProtocolWriter,
    T: Serialize<P>,
{
    fn write(&self, p: &mut P) {
        (**self).write(p);
    }
}

impl<P> Serialize<P> for ()
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, _p: &mut P) {}
}

impl<P> Serialize<P> for bool
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_bool(*self)
    }
}

impl<P> Serialize<P> for i8
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_byte(*self)
    }
}

impl<P> Serialize<P> for i16
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_i16(*self)
    }
}

impl<P> Serialize<P> for i32
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_i32(*self)
    }
}

impl<P> Serialize<P> for i64
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_i64(*self)
    }
}

impl<P> Serialize<P> for f64
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_double(*self)
    }
}

impl<P> Serialize<P> for f32
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_float(*self)
    }
}

impl<P> Serialize<P> for OrderedFloat<f64>
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_double(self.0)
    }
}

impl<P> Serialize<P> for OrderedFloat<f32>
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_float(self.0)
    }
}

impl<P> Serialize<P> for String
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_string(self.as_str())
    }
}

impl<P> Serialize<P> for str
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_string(self)
    }
}

impl<P> Serialize<P> for Bytes
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_binary(self.as_ref())
    }
}

impl<P> Serialize<P> for Vec<u8>
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_binary(self.as_ref())
    }
}

impl<P> Serialize<P> for [u8]
where
    P: ProtocolWriter,
{
    #[inline]
    fn write(&self, p: &mut P) {
        p.write_binary(self)
    }
}

impl<P, T> Serialize<P> for BTreeSet<T>
where
    P: ProtocolWriter,
    T: GetTType + Ord,
    T: Serialize<P>,
{
    fn write(&self, p: &mut P) {
        p.write_set_begin(T::TTYPE, self.len());
        for item in self.iter() {
            p.write_set_value_begin();
            item.write(p);
        }
        p.write_set_end();
    }
}

impl<P, T, S> Serialize<P> for HashSet<T, S>
where
    P: ProtocolWriter,
    T: GetTType + Hash + Eq,
    T: Serialize<P>,
    S: std::hash::BuildHasher,
{
    fn write(&self, p: &mut P) {
        p.write_set_begin(T::TTYPE, self.len());
        for item in self.iter() {
            p.write_set_value_begin();
            item.write(p);
        }
        p.write_set_end();
    }
}

impl<P, K, V> Serialize<P> for BTreeMap<K, V>
where
    P: ProtocolWriter,
    K: GetTType + Ord,
    K: Serialize<P>,
    V: GetTType,
    V: Serialize<P>,
{
    fn write(&self, p: &mut P) {
        p.write_map_begin(K::TTYPE, V::TTYPE, self.len());
        for (k, v) in self.iter() {
            p.write_map_key_begin();
            k.write(p);
            p.write_map_value_begin();
            v.write(p);
        }
        p.write_map_end();
    }
}

impl<P, K, V, S> Serialize<P> for HashMap<K, V, S>
where
    P: ProtocolWriter,
    K: GetTType + Hash + Eq,
    K: Serialize<P>,
    V: GetTType,
    V: Serialize<P>,
    S: std::hash::BuildHasher,
{
    fn write(&self, p: &mut P) {
        p.write_map_begin(K::TTYPE, V::TTYPE, self.len());
        for (k, v) in self.iter() {
            p.write_map_key_begin();
            k.write(p);
            p.write_map_value_begin();
            v.write(p);
        }
        p.write_map_end();
    }
}

impl<P, T> Serialize<P> for Vec<T>
where
    P: ProtocolWriter,
    T: GetTType,
    T: Serialize<P>,
{
    /// Vec<T> is Thrift List type
    fn write(&self, p: &mut P) {
        p.write_list_begin(T::TTYPE, self.len());
        for item in self.iter() {
            p.write_list_value_begin();
            item.write(p);
        }
        p.write_list_end();
    }
}

impl<P, T> Serialize<P> for [T]
where
    P: ProtocolWriter,
    T: GetTType,
    T: Serialize<P>,
{
    /// \[T\] is Thrift List type
    fn write(&self, p: &mut P) {
        p.write_list_begin(T::TTYPE, self.len());
        for item in self.iter() {
            p.write_list_value_begin();
            item.write(p);
        }
        p.write_list_end();
    }
}