use super::{
    BTreeMap, BTreeSet, BtPeerInfo, Digest, PeerWireBlockRequestModel,
    PeerWireExtensionHandshakeModel, PeerWireHandshakeModel, PeerWireMessageKind,
    PeerWireMetadataMessageModel, PeerWirePieceBlockModel, PeerWireTransportRequest,
    PeerWireTransportResponse, PieceState, RequestGroup, RpcError, TorrentMessageModel,
    TransportEndpoint, TransportScheme, bt_metadata_piece_count, bt_peer_metadata_key,
    bt_piece_span_bytes, bt_runtime_total_length, hex_string, resolve_bt_info_hash,
    rpc_bt_local_node_id,
};

/// Identifies the peer selected for a peer-wire compatibility exchange.
#[derive(Clone, Debug)]
pub(in crate::dispatcher) struct PeerWirePeerTarget {
    /// Original peer index inside the runtime peer list.
    pub(in crate::dispatcher) index: usize,
    /// Peer runtime snapshot used to build the outbound request.
    pub(in crate::dispatcher) peer: BtPeerInfo,
}

/// Holds the outbound peer-wire request and bookkeeping for a compatibility probe.
#[derive(Clone, Debug)]
pub(in crate::dispatcher) struct PeerWireExchangePlan {
    /// Peer list index that should receive the parsed response data.
    pub(in crate::dispatcher) peer_index: usize,
    /// Expected info hash validated against the peer handshake.
    pub(in crate::dispatcher) info_hash: [u8; 20],
    /// Transport payload sent to the peer.
    pub(in crate::dispatcher) request: PeerWireTransportRequest,
    /// Optional block request emitted after the handshake.
    pub(in crate::dispatcher) block_request: Option<PeerWireBlockRequestModel>,
    /// Optional known `ut_metadata` extension id for the selected peer.
    pub(in crate::dispatcher) metadata_extension_id: Option<u8>,
}

/// Captures peer-wire handshake and frame state recovered from a peer response.
#[derive(Clone, Debug, Default)]
pub(in crate::dispatcher) struct PeerWireExchangeResponseModel {
    /// Remote peer ID emitted by the handshake when present.
    pub(in crate::dispatcher) peer_id: Option<String>,
    /// Parsed remote extended handshake, when observed.
    pub(in crate::dispatcher) extension_handshake: Option<PeerWireExtensionHandshakeModel>,
    /// Latest observed choke state.
    pub(in crate::dispatcher) peer_choked: Option<bool>,
    /// Latest observed interest state.
    pub(in crate::dispatcher) peer_interested: Option<bool>,
    /// Piece set explicitly advertised by a bitfield frame.
    pub(in crate::dispatcher) bitfield_pieces: Option<BTreeSet<u32>>,
    /// Aggregate set of pieces implied by bitfield, have, and piece frames.
    pub(in crate::dispatcher) available_pieces: BTreeSet<u32>,
    /// Piece payload frames recovered from the response.
    pub(in crate::dispatcher) pieces: Vec<PeerWirePieceBlockModel>,
    /// Metadata payloads recovered from BEP 9 messages.
    pub(in crate::dispatcher) metadata_messages: Vec<PeerWireMetadataMessageModel>,
}

/// Builds the outbound peer-wire request for the current BitTorrent runtime state.
pub(in crate::dispatcher) fn build_peer_wire_exchange_plan(
    group: &RequestGroup,
) -> Result<PeerWireExchangePlan, RpcError> {
    let bt = group
        .bt()
        .ok_or_else(|| RpcError::unsupported("peer-wire exchange requires bt runtime state"))?;
    let target = pick_bt_peer_target(&bt.peers)?;
    let info_hash_vec = resolve_bt_info_hash(group, bt)?;
    let info_hash: [u8; 20] = info_hash_vec
        .as_slice()
        .try_into()
        .map_err(|_| RpcError::unsupported("peer-wire exchange info hash must be 20 bytes"))?;
    let peer_id_vec = rpc_bt_local_node_id(group.gid());
    let peer_id: [u8; 20] = peer_id_vec
        .clone()
        .try_into()
        .map_err(|_| RpcError::unsupported("peer-wire exchange peer id must be 20 bytes"))?;

    let peer_key = bt_peer_metadata_key(&target.peer);
    let metadata_extension_id = bt.metadata_extension_ids.get(&peer_key).copied();
    let metadata_request_piece = if bt.metadata_only {
        match (metadata_extension_id, bt.metadata_size) {
            (Some(_), Some(metadata_size)) if metadata_size > 0 => {
                let piece_count = bt_metadata_piece_count(metadata_size);
                (0..piece_count).find(|piece| !bt.metadata_piece_payloads.contains_key(piece))
            }
            _ => None,
        }
    } else {
        None
    };

    let mut handshake = PeerWireHandshakeModel::new(info_hash, peer_id);
    handshake.reserved[5] |= 0x10;
    let mut payload = handshake.serialize();

    let mut block_request = None;
    if bt.metadata_only {
        let extension_handshake = PeerWireExtensionHandshakeModel {
            extensions: BTreeMap::from([("ut_metadata".to_owned(), 1_u8)]),
            client_name: Some("aria2-rust-pro".to_owned()),
            metadata_size: None,
            request_queue: Some(16),
        };
        payload.extend_from_slice(
            &TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Extension(
                extension_handshake.to_peer_wire_message(),
            ))
            .serialize_peer_wire_frame()
            .map_err(|error| {
                RpcError::unsupported(&format!(
                    "peer-wire extension handshake serialization failed: {error}"
                ))
            })?,
        );
        if let (Some(extension_message_id), Some(piece)) =
            (metadata_extension_id, metadata_request_piece)
        {
            payload.extend_from_slice(
                &TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Extension(
                    PeerWireMetadataMessageModel::request(piece)
                        .to_peer_wire_message(extension_message_id),
                ))
                .serialize_peer_wire_frame()
                .map_err(|error| {
                    RpcError::unsupported(&format!(
                        "peer-wire metadata request serialization failed: {error}"
                    ))
                })?,
            );
        }
    } else {
        let piece_length = group.piece_length().max(1);
        let total_length = bt_runtime_total_length(group);
        let requestable = group.bt_requestable_piece_ids(false, 8);
        let availability = group.piece_availability();
        let selected_piece = requestable
            .iter()
            .copied()
            .find(|piece| availability.get(piece).copied().unwrap_or(0) > 0)
            .or_else(|| requestable.first().copied());
        block_request = selected_piece.map(|piece| PeerWireBlockRequestModel {
            piece_index: piece.0,
            block_offset: 0,
            block_length: bt_piece_span_bytes(piece, piece_length, total_length)
                .min(16_u64 * 1024)
                .max(1) as u32,
        });
        if block_request.is_some() {
            payload.extend_from_slice(
                &TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Interested)
                    .serialize_peer_wire_frame()
                    .map_err(|error| {
                        RpcError::unsupported(&format!(
                            "peer-wire interested frame serialization failed: {error}"
                        ))
                    })?,
            );
        }
        if let Some(block_request) = &block_request {
            payload.extend_from_slice(
                &TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Request(
                    block_request.clone(),
                ))
                .serialize_peer_wire_frame()
                .map_err(|error| {
                    RpcError::unsupported(&format!(
                        "peer-wire request frame serialization failed: {error}"
                    ))
                })?,
            );
        }
    }

    Ok(PeerWireExchangePlan {
        peer_index: target.index,
        info_hash,
        request: PeerWireTransportRequest {
            endpoint: TransportEndpoint {
                scheme: TransportScheme::BitTorrent,
                address: format!("{}:{}", target.peer.ip, target.peer.port),
            },
            info_hash: info_hash_vec,
            peer_id: peer_id_vec,
            payload,
        },
        block_request,
        metadata_extension_id,
    })
}

/// Chooses the best available peer target for a peer-wire exchange.
pub(super) fn pick_bt_peer_target(peers: &[BtPeerInfo]) -> Result<PeerWirePeerTarget, RpcError> {
    if peers.is_empty() {
        return Err(RpcError::unsupported(
            "peer-wire exchange requires at least one bt peer",
        ));
    }
    let mut last_error = None;
    let mut fallback = None;
    for (index, peer) in peers.iter().enumerate() {
        if peer.ip.trim().is_empty() {
            last_error = Some("peer ip must not be empty".to_owned());
            continue;
        }
        if peer.port == 0 {
            last_error = Some("peer port must be non-zero".to_owned());
            continue;
        }
        let target = PeerWirePeerTarget {
            index,
            peer: peer.clone(),
        };
        if !peer.choked {
            return Ok(target);
        }
        if fallback.is_none() {
            fallback = Some(target);
        }
    }
    if let Some(target) = fallback {
        return Ok(target);
    }
    Err(RpcError::unsupported(&format!(
        "peer-wire exchange found no valid bt peers in runtime state{}",
        last_error
            .map(|message| format!(": {message}"))
            .unwrap_or_default()
    )))
}

/// Parses a peer-wire transport payload into normalized runtime update data.
pub(in crate::dispatcher) fn parse_peer_wire_exchange_response(
    response: &PeerWireTransportResponse,
    expected_info_hash: &[u8; 20],
    known_metadata_extension_id: Option<u8>,
) -> Result<PeerWireExchangeResponseModel, RpcError> {
    let mut parsed = PeerWireExchangeResponseModel::default();
    let mut cursor = 0;
    let mut negotiated_metadata_extension_id = known_metadata_extension_id;
    if response.payload.first().copied() == Some(19) {
        let (handshake, consumed) = PeerWireHandshakeModel::parse_prefix(&response.payload)
            .map_err(|error| {
                RpcError::unsupported(&format!("invalid peer-wire handshake: {error}"))
            })?;
        if handshake.info_hash != *expected_info_hash {
            return Err(RpcError::unsupported(
                "peer-wire handshake info hash did not match download runtime state",
            ));
        }
        parsed.peer_id = Some(hex_string(&handshake.peer_id).to_ascii_lowercase());
        cursor = consumed;
    }

    while cursor < response.payload.len() {
        let (frame, consumed) = TorrentMessageModel::parse_peer_wire_frame(
            &response.payload[cursor..],
        )
        .map_err(|error| RpcError::unsupported(&format!("invalid peer-wire frame: {error}")))?;
        cursor = cursor.saturating_add(consumed);
        match frame.peer_wire_kind().map_err(|error| {
            RpcError::unsupported(&format!("invalid peer-wire message: {error}"))
        })? {
            PeerWireMessageKind::Choke => parsed.peer_choked = Some(true),
            PeerWireMessageKind::Unchoke => parsed.peer_choked = Some(false),
            PeerWireMessageKind::Interested => parsed.peer_interested = Some(true),
            PeerWireMessageKind::NotInterested => parsed.peer_interested = Some(false),
            PeerWireMessageKind::Have(piece) => {
                parsed.available_pieces.insert(piece);
            }
            PeerWireMessageKind::Bitfield(bitfield) => {
                let available = bitfield
                    .to_piece_flags(bitfield.piece_capacity())
                    .into_iter()
                    .enumerate()
                    .filter_map(|(piece, has_piece)| has_piece.then_some(piece as u32))
                    .collect::<BTreeSet<_>>();
                parsed.available_pieces.extend(available.iter().copied());
                parsed.bitfield_pieces = Some(available);
            }
            PeerWireMessageKind::Piece(piece) => {
                parsed.available_pieces.insert(piece.piece_index);
                parsed.pieces.push(piece);
            }
            PeerWireMessageKind::Extension(message) => {
                if message.extension_message_id == 0 {
                    let handshake = PeerWireExtensionHandshakeModel::from_peer_wire_message(
                        &message,
                    )
                    .map_err(|error| {
                        RpcError::unsupported(&format!(
                            "invalid peer-wire extended handshake: {error}"
                        ))
                    })?;
                    negotiated_metadata_extension_id = handshake
                        .ut_metadata_id()
                        .or(negotiated_metadata_extension_id);
                    parsed.extension_handshake = Some(handshake);
                } else if negotiated_metadata_extension_id
                    .is_some_and(|extension_id| extension_id == message.extension_message_id)
                {
                    let metadata_message = PeerWireMetadataMessageModel::from_peer_wire_message(
                        &message,
                        negotiated_metadata_extension_id.expect("checked is_some above"),
                    )
                    .map_err(|error| {
                        RpcError::unsupported(&format!(
                            "invalid peer-wire ut_metadata payload: {error}"
                        ))
                    })?;
                    parsed.metadata_messages.push(metadata_message);
                }
            }
            PeerWireMessageKind::KeepAlive
            | PeerWireMessageKind::Request(_)
            | PeerWireMessageKind::Cancel(_)
            | PeerWireMessageKind::Port(_)
            | PeerWireMessageKind::Unknown(_) => {}
        }
    }

    Ok(parsed)
}

/// Computes how many verified bytes overlap a requested byte range.
pub(in crate::dispatcher) fn verified_length_for_range(
    group: &RequestGroup,
    range_start: u64,
    range_length: u64,
    piece_length: u64,
    total_length: u64,
) -> u64 {
    if range_length == 0 || piece_length == 0 || total_length == 0 {
        return 0;
    }
    let range_end = range_start.saturating_add(range_length).min(total_length);
    if range_end <= range_start {
        return 0;
    }
    group
        .piece_map()
        .iter()
        .filter(|(_, state)| **state == PieceState::Verified)
        .map(|(piece, _)| {
            let piece_start = u64::from(piece.0).saturating_mul(piece_length);
            let piece_end = piece_start
                .saturating_add(bt_piece_span_bytes(*piece, piece_length, total_length))
                .min(total_length);
            let overlap_start = piece_start.max(range_start);
            let overlap_end = piece_end.min(range_end);
            overlap_end.saturating_sub(overlap_start)
        })
        .sum()
}

/// Returns whether a peer-wire bitfield covers every expected piece index.
pub(in crate::dispatcher) fn peer_wire_bitfield_is_complete(
    pieces: &BTreeSet<u32>,
    expected_piece_count: usize,
) -> bool {
    expected_piece_count > 0
        && pieces.len() >= expected_piece_count
        && (0..expected_piece_count as u32).all(|piece| pieces.contains(&piece))
}