use std::{
    collections::BTreeMap,
    sync::{Arc, Mutex},
};

use iroh::{
    EndpointAddr,
    endpoint::Connection,
    protocol::{AcceptError, ProtocolHandler},
};
use tokio::sync::RwLock;

use crate::{
    domain::FolderId,
    iroh::{
        CONTROL_REQUEST_TIMEOUT, DiscoveredPeer, HostedFolder, MAX_AUDIT_EVENTS_PER_RESPONSE,
        MAX_CONTROL_MESSAGE_BYTES, recover_mutex,
    },
    protocol::{ControlMessage, ManifestRequestKind, ManifestSummary, validate_invite},
};

#[derive(Clone, Debug)]
pub(super) struct AppaProtocol {
    pub(super) folders: Arc<RwLock<BTreeMap<FolderId, HostedFolder>>>,
    pub(super) discovered_peers: Arc<Mutex<BTreeMap<(FolderId, iroh::EndpointId), DiscoveredPeer>>>,
    pub(super) announced_folders: Arc<Mutex<BTreeMap<FolderId, String>>>,
}

impl ProtocolHandler for AppaProtocol {
    async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
        while let Ok((mut send, mut receive)) = connection.accept_bi().await {
            if let Err(error) = self
                .handle_request(&connection, &mut send, &mut receive)
                .await
            {
                tracing::warn!(peer = %connection.remote_id(), error = %error, "Rejected Appa control request");
            }
        }
        Ok(())
    }
}

impl AppaProtocol {
    async fn handle_request(
        &self,
        connection: &Connection,
        send: &mut iroh::endpoint::SendStream,
        receive: &mut iroh::endpoint::RecvStream,
    ) -> Result<(), AcceptError> {
        let request = read_control_request(receive).await?;
        let response = self.request_response(connection, request).await?;
        send_response(send, &response).await
    }

    async fn request_response(
        &self,
        connection: &Connection,
        request: ControlMessage,
    ) -> Result<ControlMessage, AcceptError> {
        match request {
            ControlMessage::Request {
                folder_id,
                capability,
                invite,
                requester_endpoint,
                request: request_kind,
            } => {
                self.handle_folder_request(
                    connection,
                    folder_id,
                    capability,
                    invite,
                    requester_endpoint,
                    request_kind,
                )
                .await
            }
            ControlMessage::Announcement {
                folder_id,
                capability,
                root_hash,
                sender_endpoint,
            } => {
                self.handle_announcement(
                    connection,
                    folder_id,
                    capability,
                    root_hash,
                    sender_endpoint,
                )
                .await
            }
            _ => Err(protocol_error("unsupported Appa control message")),
        }
    }

    async fn handle_folder_request(
        &self,
        connection: &Connection,
        folder_id: FolderId,
        capability: String,
        invite: Option<Box<crate::protocol::Invite>>,
        requester_endpoint: EndpointAddr,
        request_kind: ManifestRequestKind,
    ) -> Result<ControlMessage, AcceptError> {
        verify_remote_identity(connection, &requester_endpoint, "requester")?;
        let response = self
            .authorize_request(
                folder_id,
                &capability,
                invite.as_deref(),
                &requester_endpoint,
                request_kind,
            )
            .await?;
        recover_mutex(&self.discovered_peers).insert(
            (folder_id, requester_endpoint.id),
            DiscoveredPeer {
                folder_id,
                endpoint: requester_endpoint,
            },
        );
        Ok(response)
    }

    async fn handle_announcement(
        &self,
        connection: &Connection,
        folder_id: FolderId,
        capability: String,
        root_hash: String,
        sender_endpoint: EndpointAddr,
    ) -> Result<ControlMessage, AcceptError> {
        verify_remote_identity(connection, &sender_endpoint, "announcer")?;
        self.authorize_announcement(folder_id, &capability, &sender_endpoint)
            .await?;
        recover_mutex(&self.announced_folders).insert(folder_id, root_hash);
        tracing::debug!(peer = %connection.remote_id(), folder = %folder_id, "Received folder change announcement");
        Ok(ControlMessage::AnnouncementAccepted)
    }
}

impl AppaProtocol {
    async fn authorize_announcement(
        &self,
        folder_id: FolderId,
        capability: &str,
        sender: &EndpointAddr,
    ) -> Result<(), AcceptError> {
        let folders = self.folders.read().await;
        let folder = folders
            .get(&folder_id)
            .ok_or_else(|| protocol_error("folder is not hosted by this Appa node"))?;
        if capability != folder.capability || !folder.roster.can_publish(&sender.id.to_string()) {
            return Err(protocol_error("unauthorized Appa folder announcement"));
        }
        Ok(())
    }
}

impl AppaProtocol {
    async fn authorize_request(
        &self,
        folder_id: FolderId,
        capability: &str,
        invite: Option<&crate::protocol::Invite>,
        requester_endpoint: &EndpointAddr,
        request_kind: ManifestRequestKind,
    ) -> Result<ControlMessage, AcceptError> {
        let folders = self.folders.read().await;
        let folder = folders
            .get(&folder_id)
            .ok_or_else(|| protocol_error("folder is not hosted by this Appa node"))?;
        let is_current_capability = capability == folder.capability;
        let is_active_member = folder
            .active_member_ids
            .contains(&requester_endpoint.id.to_string());
        if !is_current_capability {
            return Err(protocol_error("invalid Appa folder capability"));
        }
        if !is_active_member
            && !invite.is_some_and(|invite| {
                validate_invite(invite, time::OffsetDateTime::now_utc()).is_ok()
                    && invite.folder_id == folder_id
                    && invite.capability == folder.capability
            })
        {
            return Err(protocol_error("requester lacks a valid Appa invitation"));
        }
        match request_kind {
            ManifestRequestKind::Summary => Ok(ControlMessage::Summary {
                summary: ManifestSummary {
                    folder_id,
                    root_hash: folder.merkle_index.root_hash().to_owned(),
                },
                roster: folder.roster.clone(),
                peer_endpoints: folder.peer_endpoints.clone(),
                capability: folder.capability.clone(),
            }),
            ManifestRequestKind::MerkleNode { prefix } => Ok(ControlMessage::MerkleNode {
                node: folder
                    .merkle_index
                    .node(&prefix)
                    .map_err(protocol_error)?
                    .ok_or_else(|| protocol_error("unknown Merkle prefix"))?,
            }),
            ManifestRequestKind::AuditEvents { after_sequences } => {
                Ok(ControlMessage::AuditEvents {
                    events: folder
                        .audit_events
                        .iter()
                        .filter(|event| {
                            event.sequence
                                > after_sequences
                                    .get(&event.author_device_id)
                                    .copied()
                                    .unwrap_or_default()
                        })
                        .take(MAX_AUDIT_EVENTS_PER_RESPONSE)
                        .cloned()
                        .collect(),
                })
            }
        }
    }
}

async fn read_control_request(
    receive: &mut iroh::endpoint::RecvStream,
) -> Result<ControlMessage, AcceptError> {
    tokio::time::timeout(CONTROL_REQUEST_TIMEOUT, read_request_payload(receive))
        .await
        .map_err(|_| protocol_error("Appa control request timed out"))?
}

async fn read_request_payload(
    receive: &mut iroh::endpoint::RecvStream,
) -> Result<ControlMessage, AcceptError> {
    let payload = receive
        .read_to_end(MAX_CONTROL_MESSAGE_BYTES)
        .await
        .map_err(protocol_error)?;
    serde_json::from_slice(&payload).map_err(protocol_error)
}

async fn send_response(
    send: &mut iroh::endpoint::SendStream,
    response: &ControlMessage,
) -> Result<(), AcceptError> {
    let payload = serde_json::to_vec(response).map_err(protocol_error)?;
    send.write_all(&payload).await.map_err(protocol_error)?;
    send.finish()?;
    Ok(())
}

fn verify_remote_identity(
    connection: &Connection,
    endpoint: &EndpointAddr,
    peer_role: &str,
) -> Result<(), AcceptError> {
    if endpoint.id == connection.remote_id() {
        return Ok(());
    }
    Err(protocol_error(format!(
        "Appa {peer_role} identity does not match its endpoint"
    )))
}

fn protocol_error(error: impl std::fmt::Display) -> AcceptError {
    AcceptError::from_err(std::io::Error::other(error.to_string()))
}