use std::{
    sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
    },
    time::Duration,
};

use mtp::codec::{CommunicationType, CommunicationValue, RelayError, forward_relay_frame};
use thiserror::Error;

use crate::{
    app_state::AppState, omega::omega_connection::OmegaConnection, rho::rho_manager::RhoManager,
};

const TARGET_KIND_MASK: u64 = 0xC000_0000_0000_0000;
const TARGET_ID_MASK: u64 = (1_u64 << 48) - 1;
const USER_TARGET_KIND: u64 = 0x4000_0000_0000_0000;
const IOTA_TARGET_KIND: u64 = 0x8000_0000_0000_0000;
static NEXT_RELAY_FRAME_ID: AtomicU32 = AtomicU32::new(1);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RouteTarget {
    User(u64),
    Iota(u64),
}

impl RouteTarget {
    pub fn wire_id(self) -> Option<u64> {
        let (kind, id) = match self {
            Self::User(id) => (USER_TARGET_KIND, id),
            Self::Iota(id) => (IOTA_TARGET_KIND, id),
        };
        (id > 0 && id <= TARGET_ID_MASK).then_some(kind | id)
    }

    pub fn from_wire_id(value: u64) -> Option<Self> {
        let id = value & TARGET_ID_MASK;
        if id == 0 || value & !(TARGET_KIND_MASK | TARGET_ID_MASK) != 0 {
            return None;
        }

        match value & TARGET_KIND_MASK {
            USER_TARGET_KIND => Some(Self::User(id)),
            IOTA_TARGET_KIND => Some(Self::Iota(id)),
            _ => None,
        }
    }

    pub const fn id(self) -> u64 {
        match self {
            Self::User(id) | Self::Iota(id) => id,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RelaySource {
    Client { iota_id: u64 },
    Iota { iota_id: u64 },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageSecurityClass {
    RelayOnly,
    AuthenticatedPeerControl,
    AuthenticatedLocalRequest,
}

pub fn message_security_class(frame: &CommunicationValue) -> MessageSecurityClass {
    const RELAY_ONLY_TYPES: &[CommunicationType] = &[
        CommunicationType::MessageSend,
        CommunicationType::MessageLive,
        CommunicationType::MessageState,
        CommunicationType::MessageEdit,
        CommunicationType::MessageEditLive,
        CommunicationType::MessageReactionAdd,
        CommunicationType::MessageReactionRemove,
        CommunicationType::MessageReactionLive,
        CommunicationType::MessageDelete,
        CommunicationType::MessageDeleteLive,
        CommunicationType::MessageOtherIota,
        CommunicationType::SetChatSecret,
        CommunicationType::SendChat,
        CommunicationType::SettingsSave,
        CommunicationType::GlobalSettingsSave,
        CommunicationType::AddConversation,
        CommunicationType::AddCommunity,
        CommunicationType::RemoveCommunity,
    ];

    if RELAY_ONLY_TYPES.iter().any(|kind| frame.is_type(*kind)) {
        MessageSecurityClass::RelayOnly
    } else if frame.is_type(CommunicationType::GetChatSecret)
        || frame.is_type(CommunicationType::MessageGet)
        || frame.is_type(CommunicationType::MessagesGet)
    {
        MessageSecurityClass::AuthenticatedPeerControl
    } else {
        MessageSecurityClass::AuthenticatedLocalRequest
    }
}

#[derive(Debug, Error)]
pub enum RelayRouteError {
    #[error("relay has no next-hop receiver")]
    MissingReceiver,
    #[error("relay has invalid route target {0}")]
    InvalidRouteTarget(u64),
    #[error("client relay destination {actual} is not its associated Iota {expected}")]
    InvalidClientRoute { expected: u64, actual: u64 },
    #[error("relay destination Iota is not connected to this Omikron")]
    DestinationIotaNotLocal,
    #[error("relay destination client is offline")]
    ClientOffline,
    #[error("relay route resolves back to its source Iota")]
    RouteLoop,
    #[error("relay send failed: {0}")]
    Send(String),
    #[error(transparent)]
    Relay(#[from] RelayError),
}

pub fn error_response_type(error: &RelayRouteError) -> CommunicationType {
    match error {
        RelayRouteError::DestinationIotaNotLocal | RelayRouteError::ClientOffline => {
            CommunicationType::ErrorNoIota
        }
        RelayRouteError::Send(_) => CommunicationType::ErrorInternal,
        RelayRouteError::MissingReceiver
        | RelayRouteError::InvalidRouteTarget(_)
        | RelayRouteError::InvalidClientRoute { .. }
        | RelayRouteError::RouteLoop
        | RelayRouteError::Relay(_) => CommunicationType::ErrorInvalidData,
    }
}

pub async fn route_relay(
    state: &Arc<AppState>,
    source: RelaySource,
    frame: CommunicationValue,
) -> Result<(), RelayRouteError> {
    let (next_hop, frame) = prepare_frame(ensure_relay_frame_id(frame))?;
    validate_source_next_hop(source, next_hop)?;

    match source {
        RelaySource::Client { iota_id } => {
            let rho = state
                .rho
                .get_by_iota(
                    i64::try_from(iota_id)
                        .map_err(|_| RelayRouteError::InvalidRouteTarget(iota_id))?,
                )
                .await
                .ok_or(RelayRouteError::DestinationIotaNotLocal)?;
            let response = rho
                .await_relay_to_iota(&frame)
                .await
                .map_err(RelayRouteError::Send)?;
            route_response(response)
        }
        RelaySource::Iota { iota_id } => {
            route_from_iota(&state.rho, &state.omega, iota_id, next_hop, frame).await
        }
    }
}

pub async fn route_from_omega(
    rho: &RhoManager,
    frame: CommunicationValue,
) -> Result<(), RelayRouteError> {
    let (next_hop, frame) = prepare_frame(ensure_relay_frame_id(frame))?;
    let RouteTarget::Iota(iota_id) = next_hop else {
        return Err(RelayRouteError::InvalidRouteTarget(next_hop.id()));
    };
    let target = rho
        .get_by_iota(
            i64::try_from(iota_id).map_err(|_| RelayRouteError::InvalidRouteTarget(iota_id))?,
        )
        .await
        .ok_or(RelayRouteError::DestinationIotaNotLocal)?;
    let response = target
        .await_relay_to_iota(&frame)
        .await
        .map_err(RelayRouteError::Send)?;
    route_response(response)
}

async fn route_from_iota(
    rho: &RhoManager,
    omega: &OmegaConnection,
    source_iota_id: u64,
    next_hop: RouteTarget,
    frame: CommunicationValue,
) -> Result<(), RelayRouteError> {
    match next_hop {
        RouteTarget::User(user_id) => {
            let target = rho
                .get_for_user(
                    i64::try_from(user_id)
                        .map_err(|_| RelayRouteError::InvalidRouteTarget(user_id))?,
                )
                .await
                .ok_or(RelayRouteError::ClientOffline)?;
            if !target.has_local_client(user_id).await {
                return Err(RelayRouteError::ClientOffline);
            }
            target.send_relay_to_client(&frame).await.map_err(|error| {
                if error == "client offline" {
                    RelayRouteError::ClientOffline
                } else {
                    RelayRouteError::Send(error)
                }
            })
        }
        RouteTarget::Iota(iota_id) => {
            if iota_id == source_iota_id {
                return Err(RelayRouteError::RouteLoop);
            }
            if let Some(target) = rho
                .get_by_iota(
                    i64::try_from(iota_id)
                        .map_err(|_| RelayRouteError::InvalidRouteTarget(iota_id))?,
                )
                .await
            {
                let response = target
                    .await_relay_to_iota(&frame)
                    .await
                    .map_err(RelayRouteError::Send)?;
                return route_response(response);
            }
            let response = omega
                .await_response(&frame, Some(Duration::from_secs(20)))
                .await
                .map_err(RelayRouteError::Send)?;
            route_response(response)
        }
    }
}

fn route_response(response: CommunicationValue) -> Result<(), RelayRouteError> {
    if response.is_type(CommunicationType::Success) {
        Ok(())
    } else {
        Err(RelayRouteError::Send(format!(
            "next Relay hop rejected the frame with {}",
            response.get_type()
        )))
    }
}

pub fn ensure_relay_frame_id(frame: CommunicationValue) -> CommunicationValue {
    if frame.id().is_some_and(|id| id != 0) {
        return frame;
    }
    let id = NEXT_RELAY_FRAME_ID.fetch_add(1, Ordering::Relaxed).max(1);
    frame.with_id(id)
}

fn prepare_frame(
    frame: CommunicationValue,
) -> Result<(RouteTarget, CommunicationValue), RelayRouteError> {
    let next_hop = frame.receiver().ok_or(RelayRouteError::MissingReceiver)?;
    let target =
        RouteTarget::from_wire_id(next_hop).ok_or(RelayRouteError::InvalidRouteTarget(next_hop))?;
    let frame = forward_relay_frame(&frame, next_hop)?;
    Ok((target, frame))
}

fn validate_source_next_hop(
    source: RelaySource,
    next_hop: RouteTarget,
) -> Result<(), RelayRouteError> {
    match source {
        RelaySource::Client { iota_id } => match next_hop {
            RouteTarget::Iota(actual) if actual == iota_id => Ok(()),
            _ => Err(RelayRouteError::InvalidClientRoute {
                expected: iota_id,
                actual: next_hop.id(),
            }),
        },
        RelaySource::Iota { iota_id } => {
            if next_hop == RouteTarget::Iota(iota_id) {
                Err(RelayRouteError::RouteLoop)
            } else {
                Ok(())
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use mtp::codec::{CommunicationType, CommunicationValue, DataValue};

    use super::{
        RelayRouteError, RelaySource, RouteTarget, prepare_frame, validate_source_next_hop,
    };

    fn wire(target: RouteTarget) -> u64 {
        let Some(value) = target.wire_id() else {
            panic!("valid route target was rejected");
        };
        value
    }

    fn relay(receiver: Option<u64>) -> CommunicationValue {
        let frame = CommunicationValue::new(CommunicationType::Relay)
            .without_sender()
            .with_payload(DataValue::Bytes(vec![1, 2, 3]));
        receiver.map_or(frame.clone(), |id| frame.with_receiver(id))
    }

    #[test]
    fn relay_requires_a_next_hop() {
        assert!(matches!(
            prepare_frame(relay(None)),
            Err(RelayRouteError::MissingReceiver)
        ));
    }

    #[test]
    fn relay_rejects_an_outer_sender() {
        let frame = relay(Some(wire(RouteTarget::Iota(7)))).with_sender(9);
        assert!(matches!(
            prepare_frame(frame),
            Err(RelayRouteError::Relay(
                mtp::codec::RelayError::OuterSenderPresent
            ))
        ));
    }

    #[test]
    fn relay_forwarding_preserves_payload_and_next_hop() {
        let frame = relay(Some(wire(RouteTarget::Iota(7))));
        let (next_hop, forwarded) = match prepare_frame(frame.clone()) {
            Ok(value) => value,
            Err(error) => panic!("valid relay was rejected: {error}"),
        };
        assert_eq!(next_hop, RouteTarget::Iota(7));
        assert_eq!(forwarded.receiver(), Some(wire(RouteTarget::Iota(7))));
        assert_eq!(forwarded.sender(), None);
        assert_eq!(forwarded.payload(), frame.payload());
    }

    #[test]
    fn non_relay_frames_are_rejected() {
        let frame = CommunicationValue::new(CommunicationType::Ping)
            .with_receiver(wire(RouteTarget::Iota(7)));
        assert!(matches!(
            prepare_frame(frame),
            Err(RelayRouteError::Relay(mtp::codec::RelayError::NotRelay))
        ));
    }

    #[test]
    fn client_cannot_address_another_iota() {
        assert!(matches!(
            validate_source_next_hop(RelaySource::Client { iota_id: 7 }, RouteTarget::Iota(8)),
            Err(RelayRouteError::InvalidClientRoute {
                expected: 7,
                actual: 8
            })
        ));
    }

    #[test]
    fn client_cannot_address_a_user() {
        assert!(matches!(
            validate_source_next_hop(RelaySource::Client { iota_id: 7 }, RouteTarget::User(8)),
            Err(RelayRouteError::InvalidClientRoute {
                expected: 7,
                actual: 8
            })
        ));
    }

    #[test]
    fn iota_cannot_route_to_itself() {
        assert!(matches!(
            validate_source_next_hop(RelaySource::Iota { iota_id: 7 }, RouteTarget::Iota(7)),
            Err(RelayRouteError::RouteLoop)
        ));
    }

    #[tokio::test]
    async fn omega_does_not_forward_to_another_omikron() {
        let rho = crate::rho::rho_manager::RhoManager::new();
        let result = super::route_from_omega(&rho, relay(Some(wire(RouteTarget::Iota(7))))).await;
        assert!(matches!(
            result,
            Err(RelayRouteError::DestinationIotaNotLocal)
        ));
    }

    #[test]
    fn route_target_namespace_is_explicit() {
        assert_eq!(
            RouteTarget::from_wire_id(wire(RouteTarget::User(7))),
            Some(RouteTarget::User(7))
        );
        assert_eq!(
            RouteTarget::from_wire_id(wire(RouteTarget::Iota(7))),
            Some(RouteTarget::Iota(7))
        );
        assert_eq!(RouteTarget::from_wire_id(7), None);
    }
}