use anyhow::{bail, Context};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use zerocopy::{IntoBytes, TryFromBytes};

use crate::packet::{
  Packet, PacketData,
  PacketData_variants::{Command, Disconnect, HolePunch},
  PacketHeader, PacketKind,
};

pub async fn read_packet<R: AsyncReadExt + Unpin>(reader: &mut R, udp: bool) -> anyhow::Result<Packet> {
  let mut header = [0; size_of::<PacketHeader>()];
  reader.read_exact(&mut header).await.context("reading data")?;
  let Ok(header) = PacketHeader::try_read_from_bytes(&header) else {
    bail!("parsing packet buffer")
  };

  macro_rules! read_data {
    ($read: expr, $size: expr, $ty: ident $(, $field:ident => $assert: expr)*) => {{
      async fn read_data<R: AsyncReadExt + Unpin>(reader: &mut R, size: u16) -> anyhow::Result<PacketData> {
        type T = crate::packet::PacketData_variants::$ty;
        let size = size as usize;
        if size < size_of::<T>() {
          bail!("buffer too small for packet: expected {}, got {size}", size_of::<T>())
        }
        let mut data = [0; u16::MAX as usize];
        reader.read_exact(&mut data[..size]).await.context("reading data")?;
        let packet_data = match T::try_read_from_bytes(&data[..size_of::<T>()]) {
          Ok(data) => data,
          Err(error) => {
            bail!(concat!("interpreting ", stringify!($ty), ": {:?}"), error)
          }
        };

        $(
          {
            let $field = packet_data.$field;
            $assert
          }?;
        )*
        Ok(PacketData::$ty(packet_data))
      }
      type _FixIntellisense = crate::packet::PacketData_variants::$ty;
      read_data($read, $size).await?
    }};
  }

  let data: PacketData = match header.kind {
    PacketKind::Unknown => {
      let mut data = vec![0; header.size.into()];
      reader.read_exact(&mut data).await.context("reading unknown data")?;
      PacketData::Unknown(data)
    }
    PacketKind::Init => read_data!(reader, header.size, Init),
    PacketKind::Player => read_data!(reader, header.size, Player),
    PacketKind::Cap => read_data!(reader, header.size, Cap, anim => anim.assert_valid()),
    PacketKind::Game => read_data!(reader, header.size, Game, stage => stage.assert_valid()),
    PacketKind::Tag => read_data!(reader, header.size, Tag),
    PacketKind::Connect => read_data!(reader, header.size, Connect),
    PacketKind::Disconnect => PacketData::Disconnect(Disconnect),
    PacketKind::Costume => read_data!(reader, header.size, Costume),
    PacketKind::Shine => read_data!(reader, header.size, Shine),
    PacketKind::Capture => read_data!(reader, header.size, Capture),
    PacketKind::ChangeStage => read_data!(reader, header.size, ChangeStage),
    PacketKind::Command => PacketData::Command(Command),
    PacketKind::UdpInit => read_data!(reader, header.size, UdpInit),
    PacketKind::HolePunch => PacketData::HolePunch(HolePunch),
  };

  Ok(Packet {
    user_id: header.user_id,
    udp,
    data,
  })
}

pub async fn write_packet<W: AsyncWriteExt + Unpin>(writer: &mut W, id: u128, data: PacketData) -> anyhow::Result<()> {
  let (kind, slice) = match &data {
    PacketData::Unknown(vec) => {
      if vec.len() >= (256 as usize) {
        bail!("unknown packet vec too large")
      }
      (PacketKind::Unknown, vec.as_slice())
    }
    PacketData::Init(init) => (PacketKind::Init, init.as_bytes()),
    PacketData::Player(player) => (PacketKind::Player, player.as_bytes()),
    PacketData::Cap(cap) => (PacketKind::Cap, cap.as_bytes()),
    PacketData::Game(game) => (PacketKind::Game, game.as_bytes()),
    PacketData::Tag(tag) => (PacketKind::Tag, tag.as_bytes()),
    PacketData::Connect(connect) => (PacketKind::Connect, connect.as_bytes()),
    PacketData::Disconnect(..) => (PacketKind::Disconnect, [].as_slice()),
    PacketData::Costume(costume) => (PacketKind::Costume, costume.as_bytes()),
    PacketData::Shine(shine) => (PacketKind::Shine, shine.as_bytes()),
    PacketData::Capture(capture) => (PacketKind::Capture, capture.as_bytes()),
    PacketData::ChangeStage(change_stage) => (PacketKind::ChangeStage, change_stage.as_bytes()),
    PacketData::Command(..) => (PacketKind::Command, [].as_slice()),
    PacketData::UdpInit(udp_init) => (PacketKind::UdpInit, udp_init.as_bytes()),
    PacketData::HolePunch(..) => (PacketKind::HolePunch, [].as_slice()),
  };

  writer
    .write_all(
      PacketHeader {
        kind,
        size: slice.len() as u16,
        user_id: id,
      }
      .as_bytes(),
    )
    .await
    .context("writing header")?;
  writer.write_all(slice).await.context("writing data")?;
  writer.flush().await.context("flushing writer")
}
