use bytes::{BufMut, BytesMut};
use tokio_util::codec::Encoder;

use crate::error::Error;

use super::MessageCodec;
use super::frame::Frame;

/// Write a length-prefixed string (u16 BE length + UTF-8 bytes).
fn write_string(dst: &mut BytesMut, s: &str) -> Result<(), Error> {
    let len = s.len();
    if len > u16::MAX as usize {
        return Err(Error::Codec(format!(
            "string too long for wire format: {len} bytes (max {})",
            u16::MAX
        )));
    }
    dst.put_u16(len as u16);
    dst.put_slice(s.as_bytes());
    Ok(())
}

/// Write a length-prefixed payload (u32 BE length + raw bytes).
fn write_payload(dst: &mut BytesMut, p: &[u8]) {
    dst.put_u32(p.len() as u32);
    dst.put_slice(p);
}

/// Write a header section: [u16 header_count] followed by each header as
/// [u16 key_len][key bytes][u32 value_len][value bytes].
fn write_headers(dst: &mut BytesMut, headers: &[crate::types::Header]) {
    dst.put_u16(headers.len() as u16);
    for header in headers {
        dst.put_u16(header.key.len() as u16);
        dst.put_slice(&header.key);
        dst.put_u32(header.value.len() as u32);
        dst.put_slice(&header.value);
    }
}

/// Write topic + group fields shared by Subscribe/Unsubscribe/UnsubscribeAck.
fn write_topic_group(
    dst: &mut BytesMut,
    topic: &crate::types::TopicName,
    group: &Option<crate::types::GroupName>,
) -> Result<(), Error> {
    write_string(dst, &topic.0)?;
    match group {
        Some(g) => write_string(dst, &g.0)?,
        None => write_string(dst, "")?,
    }
    Ok(())
}

impl Encoder<Frame> for MessageCodec {
    type Error = Error;

    fn encode(&mut self, item: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
        // Reserve space for the 4-byte frame length (backfilled later).
        let length_offset = dst.len();
        dst.put_u32(0);

        // Write message type byte.
        let mt: u8 = item.message_type().into();
        dst.put_u8(mt);

        // Write payload fields per message type.
        match item {
            Frame::Connect {
                version,
                namespace,
                api_key,
                capabilities,
            } => {
                dst.put_u8(version);
                write_string(dst, &namespace.0)?;
                write_string(dst, &api_key.0)?;
                dst.put_u32(capabilities);
            }
            Frame::ConnectAck {
                success,
                message,
                server_capabilities,
            } => {
                dst.put_u8(if success { 0x00 } else { 0x01 });
                write_string(dst, &message)?;
                dst.put_u32(server_capabilities);
            }
            Frame::Subscribe { topic, group } => {
                write_topic_group(dst, &topic, &group)?;
            }
            Frame::SubscribeAck {
                topic,
                group,
                start_offset,
                matched_count,
            } => {
                write_string(dst, &topic.0)?;
                dst.put_u64(start_offset.0);
                match &group {
                    Some(g) => write_string(dst, &g.0)?,
                    None => write_string(dst, "")?,
                }
                dst.put_u16(matched_count);
            }
            Frame::Unsubscribe { topic, group } => {
                write_topic_group(dst, &topic, &group)?;
            }
            Frame::UnsubscribeAck { topic, group } => {
                write_topic_group(dst, &topic, &group)?;
            }
            Frame::Publish {
                topic,
                headers,
                ttl_ms,
                payload,
            } => {
                write_string(dst, &topic.0)?;
                if self.ttl_enabled {
                    dst.put_u32(ttl_ms);
                }
                if self.headers_enabled {
                    write_headers(dst, &headers);
                }
                write_payload(dst, &payload);
            }
            Frame::PublishAck { message_id } => {
                dst.put_u64(message_id.0);
            }
            Frame::Deliver {
                message_id,
                offset,
                timestamp_millis,
                ttl_ms,
                topic,
                headers,
                payload,
            } => {
                dst.put_u64(message_id.0);
                dst.put_u64(offset.0);
                dst.put_u64(timestamp_millis);
                if self.ttl_enabled {
                    dst.put_u32(ttl_ms);
                }
                write_string(dst, &topic.0)?;
                if self.headers_enabled {
                    write_headers(dst, &headers);
                }
                write_payload(dst, &payload);
            }
            Frame::Ack { message_id } => {
                dst.put_u64(message_id.0);
            }
            Frame::Error { code, message } => {
                dst.put_u16(code.into());
                write_string(dst, &message)?;
            }
            Frame::SubscribeAt {
                topic,
                offset,
                group,
            } => {
                write_string(dst, &topic.0)?;
                dst.put_u64(offset.0);
                match &group {
                    Some(g) => write_string(dst, &g.0)?,
                    None => write_string(dst, "")?,
                }
            }
            Frame::SubscribeFrom {
                topic,
                timestamp_millis,
                group,
            } => {
                write_string(dst, &topic.0)?;
                dst.put_u64(timestamp_millis);
                match &group {
                    Some(g) => write_string(dst, &g.0)?,
                    None => write_string(dst, "")?,
                }
            }
            Frame::GoingAway {
                reason,
                deadline_ms,
                message,
            } => {
                dst.put_u8(reason.0);
                dst.put_u32(deadline_ms);
                write_string(dst, &message)?;
            }
            Frame::Reconnect {
                version,
                namespace,
                api_key,
                capabilities,
                subscriptions,
            } => {
                dst.put_u8(version);
                write_string(dst, &namespace.0)?;
                write_string(dst, &api_key.0)?;
                dst.put_u32(capabilities);
                dst.put_u16(subscriptions.len() as u16);
                for entry in &subscriptions {
                    write_string(dst, &entry.topic.0)?;
                    match &entry.group {
                        Some(g) => write_string(dst, &g.0)?,
                        None => write_string(dst, "")?,
                    }
                    dst.put_u64(entry.last_acked_offset.0);
                }
            }
            Frame::ReconnectAck {
                status,
                message,
                server_capabilities,
                results,
            } => {
                dst.put_u8(status as u8);
                write_string(dst, &message)?;
                dst.put_u32(server_capabilities);
                dst.put_u16(results.len() as u16);
                for result in results {
                    write_string(dst, &result.topic.0)?;
                    match &result.group {
                        Some(g) => write_string(dst, &g.0)?,
                        None => write_string(dst, "")?,
                    }
                    dst.put_u8(result.status as u8);
                    dst.put_u64(result.actual_start_offset.0);
                }
            }
            Frame::Fetch {
                topic,
                offset,
                max_count,
            } => {
                write_string(dst, &topic.0)?;
                dst.put_u64(offset.0);
                dst.put_u32(max_count);
            }
            Frame::FetchFrom {
                topic,
                timestamp_millis,
                max_count,
            } => {
                write_string(dst, &topic.0)?;
                dst.put_u64(timestamp_millis);
                dst.put_u32(max_count);
            }
            Frame::FetchResponse {
                topic,
                next_offset,
                head_offset,
                messages,
            } => {
                write_string(dst, &topic.0)?;
                dst.put_u64(next_offset.0);
                dst.put_u64(head_offset.0);
                dst.put_u32(messages.len() as u32);
                for msg in &messages {
                    dst.put_u64(msg.offset.0);
                    dst.put_u64(msg.timestamp_millis);
                    if self.ttl_enabled {
                        dst.put_u32(msg.ttl_ms);
                    }
                    if self.headers_enabled {
                        write_headers(dst, &msg.headers);
                    }
                    write_payload(dst, &msg.payload);
                }
            }
            Frame::Nack { message_id, reason } => {
                dst.put_u64(message_id.0);
                write_string(dst, &reason)?;
            }
            Frame::NackAck { message_id, status } => {
                dst.put_u64(message_id.0);
                dst.put_u8(status);
            }
            Frame::BatchAck { message_ids } => {
                dst.put_u32(message_ids.len() as u32);
                for id in message_ids {
                    dst.put_u64(id.0);
                }
            }
            Frame::BatchAckResponse {
                acked_count,
                unknown_count,
            } => {
                dst.put_u32(acked_count);
                dst.put_u32(unknown_count);
            }
            Frame::NotLeader { leader_addr } => {
                write_string(dst, &leader_addr)?;
            }
            Frame::ClusterUnavailable {
                retry_after_ms,
                message,
            } => {
                dst.put_u32(retry_after_ms);
                write_string(dst, &message)?;
            }
            Frame::Ping | Frame::Pong => {
                // Empty payload — message type byte only.
            }
        }

        // Backfill frame length: total bytes after the length field = message type + payload.
        let frame_length = (dst.len() - length_offset - 4) as u32;
        let length_bytes = frame_length.to_be_bytes();
        dst[length_offset..length_offset + 4].copy_from_slice(&length_bytes);

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codec::ErrorCode;
    use crate::types::{ApiKey, GroupName, MessageId, NamespaceId, TopicName, TopicOffset};
    use bytes::Bytes;

    fn encode(frame: Frame) -> Result<BytesMut, Error> {
        let mut codec = MessageCodec::new(1_048_576);
        let mut buf = BytesMut::new();
        codec.encode(frame, &mut buf)?;
        Ok(buf)
    }

    #[test]
    fn encode_ping() {
        let buf = encode(Frame::Ping).unwrap();
        assert_eq!(&buf[..], &[0, 0, 0, 1, 0x0C]);
    }

    #[test]
    fn encode_pong() {
        let buf = encode(Frame::Pong).unwrap();
        assert_eq!(&buf[..], &[0, 0, 0, 1, 0x0D]);
    }

    #[test]
    fn encode_connect() {
        let buf = encode(Frame::Connect {
            version: 1,
            namespace: NamespaceId("ns".into()),
            api_key: ApiKey("key".into()),
            capabilities: 0,
        })
        .unwrap();
        // frame_length = 1(type) + 1(ver) + 2+2(ns) + 2+3(key) + 4(caps) = 15
        assert_eq!(&buf[0..4], &[0, 0, 0, 15]);
        assert_eq!(buf[4], 0x01); // Connect
        assert_eq!(buf[5], 1); // version
        assert_eq!(&buf[6..8], &[0, 2]); // ns len
        assert_eq!(&buf[8..10], b"ns");
        assert_eq!(&buf[10..12], &[0, 3]); // key len
        assert_eq!(&buf[12..15], b"key");
        assert_eq!(&buf[15..19], &[0, 0, 0, 0]); // capabilities
    }

    #[test]
    fn encode_connect_ack_success() {
        let buf = encode(Frame::ConnectAck {
            success: true,
            message: String::new(),
            server_capabilities: 0,
        })
        .unwrap();
        // frame_length = 1(type) + 1(status) + 2+0(msg) + 4(caps) = 8
        assert_eq!(&buf[0..4], &[0, 0, 0, 8]);
        assert_eq!(buf[4], 0x02);
        assert_eq!(buf[5], 0x00); // success
        assert_eq!(&buf[6..8], &[0, 0]); // empty message
        assert_eq!(&buf[8..12], &[0, 0, 0, 0]); // server_capabilities
    }

    #[test]
    fn encode_connect_ack_failure() {
        let buf = encode(Frame::ConnectAck {
            success: false,
            message: "bad key".into(),
            server_capabilities: 0,
        })
        .unwrap();
        assert_eq!(buf[4], 0x02);
        assert_eq!(buf[5], 0x01); // failure
        assert_eq!(&buf[6..8], &[0, 7]); // message length
        assert_eq!(&buf[8..15], b"bad key");
        assert_eq!(&buf[15..19], &[0, 0, 0, 0]); // server_capabilities
    }

    #[test]
    fn encode_subscribe_with_group() {
        let buf = encode(Frame::Subscribe {
            topic: TopicName("events".into()),
            group: Some(GroupName("workers".into())),
        })
        .unwrap();
        assert_eq!(buf[4], 0x03);
        // topic: 2+6=8, group: 2+7=9, type: 1 => 18
        assert_eq!(&buf[0..4], &[0, 0, 0, 18]);
    }

    #[test]
    fn encode_subscribe_without_group() {
        let buf = encode(Frame::Subscribe {
            topic: TopicName("events".into()),
            group: None,
        })
        .unwrap();
        assert_eq!(buf[4], 0x03);
        // topic: 2+6=8, group: 2+0=2, type: 1 => 11
        assert_eq!(&buf[0..4], &[0, 0, 0, 11]);
        // Check group is empty string
        let group_offset = 4 + 1 + 2 + 6; // length_prefix + type + topic_len_prefix + topic
        assert_eq!(&buf[group_offset..group_offset + 2], &[0, 0]);
    }

    #[test]
    fn encode_publish() {
        let buf = encode(Frame::Publish {
            topic: TopicName("t".into()),
            payload: Bytes::from_static(b"hello"),
            headers: vec![],
            ttl_ms: 0,
        })
        .unwrap();
        assert_eq!(buf[4], 0x07);
        // type: 1, topic: 2+1=3, payload: 4+5=9 => 13
        assert_eq!(&buf[0..4], &[0, 0, 0, 13]);
    }

    #[test]
    fn encode_publish_empty_payload() {
        let buf = encode(Frame::Publish {
            topic: TopicName("t".into()),
            payload: Bytes::new(),
            headers: vec![],
            ttl_ms: 0,
        })
        .unwrap();
        // type: 1, topic: 2+1=3, payload: 4+0=4 => 8
        assert_eq!(&buf[0..4], &[0, 0, 0, 8]);
    }

    #[test]
    fn encode_publish_ack() {
        let buf = encode(Frame::PublishAck {
            message_id: MessageId(42),
        })
        .unwrap();
        assert_eq!(buf[4], 0x08);
        // type: 1, id: 8 => 9
        assert_eq!(&buf[0..4], &[0, 0, 0, 9]);
        assert_eq!(&buf[5..13], &42u64.to_be_bytes());
    }

    #[test]
    fn encode_deliver() {
        let buf = encode(Frame::Deliver {
            message_id: MessageId(7),
            offset: TopicOffset(100),
            timestamp_millis: 1_700_000_000_000,
            ttl_ms: 0,
            topic: TopicName("t".into()),
            payload: Bytes::from_static(b"data"),
            headers: vec![],
        })
        .unwrap();
        assert_eq!(buf[4], 0x09);
        // type: 1, id: 8, offset: 8, timestamp: 8, topic: 2+1=3, payload: 4+4=8 => 36
        assert_eq!(&buf[0..4], &[0, 0, 0, 36]);
    }

    #[test]
    fn encode_ack() {
        let buf = encode(Frame::Ack {
            message_id: MessageId(99),
        })
        .unwrap();
        assert_eq!(buf[4], 0x0A);
        assert_eq!(&buf[0..4], &[0, 0, 0, 9]);
    }

    #[test]
    fn encode_error() {
        let buf = encode(Frame::Error {
            code: ErrorCode::PermissionDenied,
            message: "nope".into(),
        })
        .unwrap();
        assert_eq!(buf[4], 0x0B);
        // type: 1, code: 2, msg: 2+4=6 => 9
        assert_eq!(&buf[0..4], &[0, 0, 0, 9]);
        assert_eq!(&buf[5..7], &0x0002u16.to_be_bytes());
    }

    #[test]
    fn encode_going_away() {
        use crate::codec::frame::GoingAwayReason;
        let buf = encode(Frame::GoingAway {
            reason: GoingAwayReason::UPGRADE,
            deadline_ms: 30000,
            message: "upgrade".into(),
        })
        .unwrap();
        assert_eq!(buf[4], 0x10); // GoingAway
        assert_eq!(buf[5], 0x01); // reason: UPGRADE
        assert_eq!(&buf[6..10], &30000u32.to_be_bytes()); // deadline_ms
        assert_eq!(&buf[10..12], &[0, 7]); // message length
        assert_eq!(&buf[12..19], b"upgrade");
        // frame_length = 1(type) + 1(reason) + 4(deadline) + 2+7(msg) = 15
        assert_eq!(&buf[0..4], &[0, 0, 0, 15]);
    }

    #[test]
    fn encode_going_away_zero_deadline() {
        use crate::codec::frame::GoingAwayReason;
        let buf = encode(Frame::GoingAway {
            reason: GoingAwayReason::SHUTDOWN,
            deadline_ms: 0,
            message: String::new(),
        })
        .unwrap();
        assert_eq!(buf[4], 0x10);
        assert_eq!(buf[5], 0x04); // SHUTDOWN
        assert_eq!(&buf[6..10], &0u32.to_be_bytes());
        assert_eq!(&buf[10..12], &[0, 0]); // empty message
        // frame_length = 1(type) + 1(reason) + 4(deadline) + 2+0(msg) = 8
        assert_eq!(&buf[0..4], &[0, 0, 0, 8]);
    }

    #[test]
    fn encode_going_away_max_deadline() {
        use crate::codec::frame::GoingAwayReason;
        let buf = encode(Frame::GoingAway {
            reason: GoingAwayReason::MAINTENANCE,
            deadline_ms: u32::MAX,
            message: "long drain".into(),
        })
        .unwrap();
        assert_eq!(buf[5], 0x02); // MAINTENANCE
        assert_eq!(&buf[6..10], &u32::MAX.to_be_bytes());
    }

    #[test]
    fn encode_string_too_long() {
        let long_string = "x".repeat(65536);
        let result = encode(Frame::ConnectAck {
            success: true,
            message: long_string,
            server_capabilities: 0,
        });
        assert!(result.is_err());
    }

    #[test]
    fn encode_nack() {
        let buf = encode(Frame::Nack {
            message_id: MessageId(42),
            reason: "bad data".into(),
        })
        .unwrap();
        assert_eq!(buf[4], 0x16); // Nack
        // message_id (8B) + reason_len (2B) + reason (8B) = 18 + type = 19
        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
        assert_eq!(frame_len, 19);
    }

    #[test]
    fn encode_nack_preserves_full_reason() {
        // Encoder does not truncate — truncation is a decoder responsibility
        // (defense in depth against non-standard clients).
        let long_reason = "x".repeat(2000);
        let buf = encode(Frame::Nack {
            message_id: MessageId(1),
            reason: long_reason,
        })
        .unwrap();
        // frame = 1(type) + 8(msg_id) + 2(reason_len) + 2000(reason) = 2011
        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
        assert_eq!(frame_len, 2011);
    }

    #[test]
    fn encode_nack_ack() {
        let buf = encode(Frame::NackAck {
            message_id: MessageId(42),
            status: 0x02,
        })
        .unwrap();
        assert_eq!(buf[4], 0x17); // NackAck
        // frame = 1(type) + 8(msg_id) + 1(status) = 10
        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
        assert_eq!(frame_len, 10);
        assert_eq!(buf[13], 0x02); // status
    }

    #[test]
    fn encode_batch_ack() {
        let buf = encode(Frame::BatchAck {
            message_ids: vec![MessageId(1), MessageId(2), MessageId(3)],
        })
        .unwrap();
        assert_eq!(buf[4], 0x18); // BatchAck
        // frame = 1(type) + 4(count) + 3*8(ids) = 29
        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
        assert_eq!(frame_len, 29);
        // count = 3
        let count = u32::from_be_bytes([buf[5], buf[6], buf[7], buf[8]]);
        assert_eq!(count, 3);
    }

    #[test]
    fn encode_batch_ack_empty() {
        let buf = encode(Frame::BatchAck {
            message_ids: vec![],
        })
        .unwrap();
        assert_eq!(buf[4], 0x18);
        // frame = 1(type) + 4(count) = 5
        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
        assert_eq!(frame_len, 5);
    }

    #[test]
    fn encode_batch_ack_response() {
        let buf = encode(Frame::BatchAckResponse {
            acked_count: 10,
            unknown_count: 3,
        })
        .unwrap();
        assert_eq!(buf[4], 0x19); // BatchAckResponse
        // frame = 1(type) + 4(acked) + 4(unknown) = 9
        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
        assert_eq!(frame_len, 9);
        let acked = u32::from_be_bytes([buf[5], buf[6], buf[7], buf[8]]);
        let unknown = u32::from_be_bytes([buf[9], buf[10], buf[11], buf[12]]);
        assert_eq!(acked, 10);
        assert_eq!(unknown, 3);
    }
}
