use std::collections::BTreeMap;

use anyhow::{Context, Result, bail};
use kafka_protocol::error::{ParseResponseErrorCode, ResponseError};
use kafka_protocol::messages::find_coordinator_response::FindCoordinatorResponse;
use kafka_protocol::messages::txn_offset_commit_request::{
    TxnOffsetCommitRequestPartition, TxnOffsetCommitRequestTopic,
};
use kafka_protocol::messages::{
    FindCoordinatorRequest, ProducerId, TransactionalId, TxnOffsetCommitRequest,
};
use kafka_protocol::protocol::StrBytes;

use super::state::{ProducerIdentity, TransactionCoordinator};
use crate::constants::{FIND_COORDINATOR_GROUP_KEY_TYPE, FIND_COORDINATOR_TRANSACTION_KEY_TYPE};
use crate::network::BrokerConnection;
use crate::types::{CommitOffset, ConsumerGroupMetadata};
use crate::{ConsumerGroupMetadataError, ProducerError};

pub enum TransactionFailureDisposition {
    AbortOnly,
    Fatal,
}

pub fn classify_transactional_error(error: ResponseError) -> TransactionFailureDisposition {
    match error {
        ResponseError::ProducerFenced
        | ResponseError::TransactionalIdAuthorizationFailed
        | ResponseError::InvalidProducerIdMapping
        | ResponseError::GroupAuthorizationFailed
        | ResponseError::FencedInstanceId
        | ResponseError::UnsupportedForMessageFormat => TransactionFailureDisposition::Fatal,
        _ => TransactionFailureDisposition::AbortOnly,
    }
}

pub fn ensure_transaction_v2_feature(
    connection: &BrokerConnection,
) -> std::result::Result<(), ProducerError> {
    let Some(level) = connection.finalized_feature_level("transaction.version") else {
        return Err(ProducerError::MissingTransactionVersionFeature);
    };

    if level < 2 {
        return Err(ProducerError::UnsupportedTransactionVersion { level });
    }

    Ok(())
}

pub fn build_find_coordinator_request(
    transactional_id: &str,
    version: i16,
) -> FindCoordinatorRequest {
    if version >= 4 {
        FindCoordinatorRequest::default()
            .with_key_type(FIND_COORDINATOR_TRANSACTION_KEY_TYPE)
            .with_coordinator_keys(vec![StrBytes::from_string(transactional_id.to_owned())])
    } else {
        FindCoordinatorRequest::default()
            .with_key(StrBytes::from_string(transactional_id.to_owned()))
            .with_key_type(FIND_COORDINATOR_TRANSACTION_KEY_TYPE)
    }
}

pub fn build_group_find_coordinator_request(
    group_id: &str,
    version: i16,
) -> FindCoordinatorRequest {
    if version >= 4 {
        FindCoordinatorRequest::default()
            .with_key_type(FIND_COORDINATOR_GROUP_KEY_TYPE)
            .with_coordinator_keys(vec![StrBytes::from_string(group_id.to_owned())])
    } else {
        FindCoordinatorRequest::default()
            .with_key(StrBytes::from_string(group_id.to_owned()))
            .with_key_type(FIND_COORDINATOR_GROUP_KEY_TYPE)
    }
}

pub fn validate_group_metadata(
    group_metadata: &ConsumerGroupMetadata,
) -> std::result::Result<(), ConsumerGroupMetadataError> {
    if group_metadata.group_id.trim().is_empty() {
        return Err(ConsumerGroupMetadataError::EmptyGroupId);
    }
    if group_metadata.generation_id > 0 && group_metadata.member_id.is_empty() {
        return Err(ConsumerGroupMetadataError::MissingMemberId);
    }
    Ok(())
}

pub fn build_txn_offset_commit_request(
    transactional_id: &str,
    producer: ProducerIdentity,
    offsets: &[CommitOffset],
    group_metadata: &ConsumerGroupMetadata,
) -> TxnOffsetCommitRequest {
    let mut topics = BTreeMap::<String, Vec<TxnOffsetCommitRequestPartition>>::new();
    for offset in offsets {
        topics.entry(offset.topic.clone()).or_default().push(
            TxnOffsetCommitRequestPartition::default()
                .with_partition_index(offset.partition)
                .with_committed_offset(offset.offset)
                .with_committed_leader_epoch(-1)
                .with_committed_metadata(None),
        );
    }

    TxnOffsetCommitRequest::default()
        .with_transactional_id(TransactionalId(StrBytes::from_string(
            transactional_id.to_owned(),
        )))
        .with_group_id(StrBytes::from_string(group_metadata.group_id.clone()).into())
        .with_producer_id(ProducerId(producer.id))
        .with_producer_epoch(producer.epoch)
        .with_generation_id(group_metadata.generation_id)
        .with_member_id(StrBytes::from_string(group_metadata.member_id.clone()))
        .with_group_instance_id(
            group_metadata
                .group_instance_id
                .clone()
                .map(StrBytes::from_string),
        )
        .with_topics(
            topics
                .into_iter()
                .map(|(topic, partitions)| {
                    TxnOffsetCommitRequestTopic::default()
                        .with_name(StrBytes::from_string(topic).into())
                        .with_partitions(partitions)
                })
                .collect(),
        )
}

pub fn parse_find_coordinator_response(
    response: FindCoordinatorResponse,
    version: i16,
) -> Result<TransactionCoordinator> {
    if version >= 4 {
        let coordinator = response
            .coordinators
            .into_iter()
            .next()
            .context("FindCoordinator returned no coordinators")?;
        if let Some(error) = coordinator.error_code.err() {
            bail!("FindCoordinator failed: {error}");
        }
        let port = u16::try_from(coordinator.port)
            .with_context(|| format!("invalid coordinator port {}", coordinator.port))?;
        return Ok(TransactionCoordinator {
            broker_id: *coordinator.node_id,
            address: format!("{}:{}", coordinator.host, port),
        });
    }

    if let Some(error) = response.error_code.err() {
        bail!("FindCoordinator failed: {error}");
    }
    let port = u16::try_from(response.port)
        .with_context(|| format!("invalid coordinator port {}", response.port))?;
    Ok(TransactionCoordinator {
        broker_id: *response.node_id,
        address: format!("{}:{}", response.host, port),
    })
}

pub fn find_coordinator_error(
    response: &FindCoordinatorResponse,
    version: i16,
) -> Option<ResponseError> {
    if version >= 4 {
        return response
            .coordinators
            .first()
            .and_then(|coordinator| coordinator.error_code.err());
    }

    response.error_code.err()
}

#[cfg(test)]
mod tests {
    use super::*;
    use kafka_protocol::messages::BrokerId;
    use kafka_protocol::messages::find_coordinator_response::Coordinator;

    #[test]
    fn validate_group_metadata_requires_group_id() {
        let error = validate_group_metadata(&ConsumerGroupMetadata {
            group_id: "   ".to_owned(),
            generation_id: 0,
            member_id: String::new(),
            group_instance_id: None,
        })
        .unwrap_err();

        assert!(matches!(error, ConsumerGroupMetadataError::EmptyGroupId));
    }

    #[test]
    fn validate_group_metadata_requires_member_id_for_active_generation() {
        let error = validate_group_metadata(&ConsumerGroupMetadata {
            group_id: "group-a".to_owned(),
            generation_id: 3,
            member_id: String::new(),
            group_instance_id: None,
        })
        .unwrap_err();

        assert!(matches!(error, ConsumerGroupMetadataError::MissingMemberId));
    }

    #[test]
    fn classify_transactional_errors_distinguishes_fatal_from_abort_only() {
        assert!(matches!(
            classify_transactional_error(ResponseError::ProducerFenced),
            TransactionFailureDisposition::Fatal
        ));
        assert!(matches!(
            classify_transactional_error(ResponseError::UnknownServerError),
            TransactionFailureDisposition::AbortOnly
        ));
    }

    #[test]
    fn find_coordinator_requests_follow_version_shape() {
        let old = build_find_coordinator_request("tx-a", 3);
        assert_eq!(old.key.to_string(), "tx-a");
        assert!(old.coordinator_keys.is_empty());
        assert_eq!(old.key_type, FIND_COORDINATOR_TRANSACTION_KEY_TYPE);

        let modern = build_find_coordinator_request("tx-a", 4);
        assert!(modern.key.is_empty());
        assert_eq!(modern.coordinator_keys[0].to_string(), "tx-a");

        let group = build_group_find_coordinator_request("group-a", 4);
        assert_eq!(group.key_type, FIND_COORDINATOR_GROUP_KEY_TYPE);
        assert_eq!(group.coordinator_keys[0].to_string(), "group-a");
    }

    #[test]
    fn txn_offset_commit_request_groups_offsets_by_topic() {
        let request = build_txn_offset_commit_request(
            "tx-a",
            ProducerIdentity { id: 42, epoch: 3 },
            &[
                CommitOffset {
                    topic: "topic-a".to_owned(),
                    partition: 0,
                    offset: 7,
                },
                CommitOffset {
                    topic: "topic-b".to_owned(),
                    partition: 1,
                    offset: 11,
                },
                CommitOffset {
                    topic: "topic-a".to_owned(),
                    partition: 2,
                    offset: 13,
                },
            ],
            &ConsumerGroupMetadata {
                group_id: "group-a".to_owned(),
                generation_id: 5,
                member_id: "member-a".to_owned(),
                group_instance_id: Some("instance-a".to_owned()),
            },
        );

        assert_eq!(request.transactional_id.0.to_string(), "tx-a");
        assert_eq!(request.producer_id.0, 42);
        assert_eq!(request.producer_epoch, 3);
        assert_eq!(request.group_id.0.to_string(), "group-a");
        assert_eq!(request.group_instance_id.unwrap().to_string(), "instance-a");
        assert_eq!(request.topics.len(), 2);
        let topic_a = request
            .topics
            .iter()
            .find(|topic| topic.name.to_string() == "topic-a")
            .unwrap();
        assert_eq!(topic_a.partitions.len(), 2);
    }

    #[test]
    fn parse_find_coordinator_response_handles_old_and_modern_errors() {
        let old = FindCoordinatorResponse::default()
            .with_node_id(BrokerId(2))
            .with_host(StrBytes::from_static_str("broker-a"))
            .with_port(9092);
        let coordinator = parse_find_coordinator_response(old.clone(), 3).unwrap();
        assert_eq!(coordinator.broker_id, 2);
        assert_eq!(coordinator.address, "broker-a:9092");
        assert!(parse_find_coordinator_response(old.with_port(-1), 3).is_err());

        let modern = FindCoordinatorResponse::default().with_coordinators(vec![
            Coordinator::default()
                .with_node_id(BrokerId(3))
                .with_host(StrBytes::from_static_str("broker-b"))
                .with_port(9093),
        ]);
        let coordinator = parse_find_coordinator_response(modern, 4).unwrap();
        assert_eq!(coordinator.broker_id, 3);
        assert_eq!(coordinator.address, "broker-b:9093");

        let error = FindCoordinatorResponse::default()
            .with_error_code(ResponseError::CoordinatorNotAvailable.code());
        assert!(find_coordinator_error(&error, 3).is_some());
        assert!(parse_find_coordinator_response(error, 3).is_err());

        let error = FindCoordinatorResponse::default().with_coordinators(vec![
            Coordinator::default().with_error_code(ResponseError::CoordinatorNotAvailable.code()),
        ]);
        assert!(find_coordinator_error(&error, 4).is_some());
        assert!(parse_find_coordinator_response(error, 4).is_err());
        assert!(parse_find_coordinator_response(FindCoordinatorResponse::default(), 4).is_err());
    }
}
