use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use dashmap::DashMap;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use zeroize::Zeroizing;

use crate::SdkError;

/// Maximum incoming payload size (512 KB) before MLS processing.
const MAX_INCOMING_PAYLOAD_SIZE: usize = 512 * 1024;

/// Generated protobuf types from DataPlaneService.
pub mod slim_proto {
    tonic::include_proto!("dataplane.proto.v1");
}

use slim_proto::data_plane_service_client::DataPlaneServiceClient;
use slim_proto::{
    message, Content, Message, Name, Publish, SessionHeader, SessionMessageType, SessionType,
    SlimHeader, StringName, Subscribe,
};

/// Errors from transport operations.
#[derive(Debug, thiserror::Error)]
pub enum TransportError {
    #[error("gRPC transport error: {0}")]
    Grpc(#[from] tonic::transport::Error),

    #[error("gRPC status error: {0}")]
    Status(#[from] tonic::Status),

    #[error("send error: outbound channel closed")]
    SendFailed,

    #[error("mock transport error: {0}")]
    Mock(String),
}

/// An incoming message dispatched from the gRPC stream.
#[derive(Debug, Clone)]
pub struct IncomingMessage {
    pub channel_id: [u8; 32],
    pub session_message_type: i32,
    pub payload: Vec<u8>,
}

/// Trait abstracting message transport for the SDK.
///
/// Implemented by `GrpcTransport` (production) and `MockTransport` (testing).
#[async_trait]
pub trait Transport: Send + Sync + 'static {
    /// Subscribe to a channel. Returns a receiver for incoming messages.
    async fn subscribe(
        &self,
        channel_name: (&str, &str, &str),
        channel_id: [u8; 32],
    ) -> Result<mpsc::Receiver<IncomingMessage>, TransportError>;

    /// Publish a message to a channel.
    async fn publish(
        &self,
        channel_name: (&str, &str, &str),
        session_message_type: SessionMessageType,
        content: Content,
    ) -> Result<(), TransportError>;
}

/// Validate that a gRPC endpoint uses HTTPS for remote connections.
///
/// Delegates to `validate_url_security` for proper URL parsing to prevent
/// bypasses like `http://localhost.evil.com`.
pub(crate) fn validate_grpc_endpoint(endpoint: &str) -> Result<(), SdkError> {
    crate::client::validate_url_security(endpoint).map_err(SdkError::Auth)
}

/// Parse the `exp` claim from a JWT (base64url-encoded middle segment).
///
/// Returns `None` if the token format is invalid or the `exp` field is missing.
fn jwt_expiry(token: &str) -> Option<u64> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return None;
    }
    // base64url decode the payload (middle segment).
    use base64::Engine;
    let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(parts[1])
        .ok()?;
    let value: serde_json::Value = serde_json::from_slice(&payload_bytes).ok()?;
    value.get("exp")?.as_u64()
}

/// gRPC transport client for the DataPlaneService bidirectional stream.
///
/// # Example
///
/// ```no_run
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// use skytale_sdk::transport::GrpcTransport;
/// let transport = GrpcTransport::connect("https://relay.skytale.sh:5000", None).await?;
/// # Ok(())
/// # }
/// ```
pub struct GrpcTransport {
    outbound_tx: mpsc::Sender<Message>,
    channel_receivers: Arc<DashMap<[u8; 32], mpsc::Sender<IncomingMessage>>>,
}

impl GrpcTransport {
    /// Connect to the supernode's gRPC DataPlaneService.
    ///
    /// When `jwt` is `Some`, injects a `Bearer` authorization header on every
    /// outgoing gRPC request via a tonic interceptor.
    pub async fn connect(
        endpoint: &str,
        jwt: Option<Zeroizing<String>>,
    ) -> Result<Self, TransportError> {
        let (outbound_tx, outbound_rx) = mpsc::channel::<Message>(64);
        let outbound_stream = ReceiverStream::new(outbound_rx);

        let channel = tonic::transport::Channel::from_shared(endpoint.to_string())
            .map_err(|_| {
                TransportError::Status(tonic::Status::invalid_argument("invalid endpoint URI"))
            })?
            .connect()
            .await?;

        let response = if let Some(token) = jwt {
            let expiry = jwt_expiry(token.as_str());
            if let Some(exp) = expiry {
                tracing::info!(expires_at = exp, "JWT token loaded");
            } else {
                tracing::warn!("could not parse JWT expiry — token will not be refreshed");
            }

            // SECURITY NOTE: The JWT persists in the tonic MetadataValue for the
            // lifetime of this transport. tonic::metadata::MetadataValue does not
            // support Zeroize. Mitigation: JWTs should have short TTLs (5 min).
            let auth_value: tonic::metadata::MetadataValue<_> =
                format!("Bearer {}", token.as_str()).parse().map_err(|_| {
                    TransportError::Status(tonic::Status::internal("invalid JWT token format"))
                })?;
            let mut client = DataPlaneServiceClient::with_interceptor(
                channel,
                move |mut req: tonic::Request<()>| {
                    // Reject requests if the JWT has expired rather than sending
                    // a stale token that would produce opaque gRPC failures.
                    if let Some(exp) = expiry {
                        let now = SystemTime::now()
                            .duration_since(UNIX_EPOCH)
                            .unwrap_or_default()
                            .as_secs();
                        if now >= exp {
                            return Err(tonic::Status::unauthenticated(
                                "JWT has expired — reconnect with a fresh token",
                            ));
                        }
                        if now + 60 >= exp {
                            tracing::warn!(
                                expires_in_secs = exp.saturating_sub(now),
                                "JWT expires within 60 seconds"
                            );
                        }
                    }
                    req.metadata_mut()
                        .insert("authorization", auth_value.clone());
                    Ok(req)
                },
            );
            client.open_channel(outbound_stream).await?
        } else {
            let mut client = DataPlaneServiceClient::new(channel);
            client.open_channel(outbound_stream).await?
        };

        let mut inbound = response.into_inner();

        let channel_receivers: Arc<DashMap<[u8; 32], mpsc::Sender<IncomingMessage>>> =
            Arc::new(DashMap::new());

        let receivers = channel_receivers.clone();
        tokio::spawn(async move {
            while let Ok(Some(msg)) = inbound.message().await {
                if let Some(incoming) = parse_incoming_message(&msg) {
                    if let Some(tx) = receivers.get(&incoming.channel_id) {
                        if tx.try_send(incoming).is_err() {
                            tracing::warn!("incoming message channel full, message dropped");
                        }
                    }
                }
            }
        });

        Ok(Self {
            outbound_tx,
            channel_receivers,
        })
    }
}

#[async_trait]
impl Transport for GrpcTransport {
    /// Subscribe to a channel. Returns a receiver for incoming messages on that channel.
    ///
    /// `channel_name` is a 3-component string name `(org, channel, subchannel)`.
    /// `channel_id` is the 32-byte hash that identifies this channel locally.
    async fn subscribe(
        &self,
        channel_name: (&str, &str, &str),
        channel_id: [u8; 32],
    ) -> Result<mpsc::Receiver<IncomingMessage>, TransportError> {
        let (tx, rx) = mpsc::channel::<IncomingMessage>(64);
        self.channel_receivers.insert(channel_id, tx);

        let msg = Message {
            message_type: Some(message::MessageType::Subscribe(Subscribe {
                header: Some(make_header(channel_name)),
            })),
            metadata: Default::default(),
        };
        self.outbound_tx
            .send(msg)
            .await
            .map_err(|_| TransportError::SendFailed)?;

        Ok(rx)
    }

    /// Publish a message (application data or MLS operation) to a channel.
    async fn publish(
        &self,
        channel_name: (&str, &str, &str),
        session_message_type: SessionMessageType,
        content: Content,
    ) -> Result<(), TransportError> {
        let msg = Message {
            message_type: Some(message::MessageType::Publish(Publish {
                header: Some(make_header(channel_name)),
                session: Some(SessionHeader {
                    session_type: SessionType::Multicast as i32,
                    session_message_type: session_message_type as i32,
                    session_id: 0,
                    message_id: 0,
                }),
                msg: Some(content),
            })),
            metadata: Default::default(),
        };
        self.outbound_tx
            .send(msg)
            .await
            .map_err(|_| TransportError::SendFailed)?;
        Ok(())
    }
}

fn make_header(channel_name: (&str, &str, &str)) -> SlimHeader {
    SlimHeader {
        source: None,
        destination: Some(Name {
            name: None,
            str_name: Some(StringName {
                str_component_0: channel_name.0.to_string(),
                str_component_1: channel_name.1.to_string(),
                str_component_2: channel_name.2.to_string(),
            }),
        }),
        identity: String::new(),
        fanout: 0,
        recv_from: None,
        forward_to: None,
        incoming_conn: None,
        error: None,
    }
}

/// Extract an IncomingMessage from a gRPC Message, if it contains a Publish payload.
fn parse_incoming_message(msg: &Message) -> Option<IncomingMessage> {
    let publish = match &msg.message_type {
        Some(message::MessageType::Publish(p)) => p,
        _ => return None,
    };

    // Derive channel_id from destination name components via SHA-256.
    // Uses "/" separators to match the supernode's hash_name().
    let channel_id = if let Some(dest) = &publish.header.as_ref()?.destination {
        if let Some(str_name) = &dest.str_name {
            let mut hasher = <sha2::Sha256 as sha2::Digest>::new();
            sha2::Digest::update(&mut hasher, str_name.str_component_0.as_bytes());
            sha2::Digest::update(&mut hasher, b"/");
            sha2::Digest::update(&mut hasher, str_name.str_component_1.as_bytes());
            sha2::Digest::update(&mut hasher, b"/");
            sha2::Digest::update(&mut hasher, str_name.str_component_2.as_bytes());
            let hash: [u8; 32] = sha2::Digest::finalize(hasher).into();
            hash
        } else {
            return None;
        }
    } else {
        return None;
    };

    let session_message_type = publish
        .session
        .as_ref()
        .map(|s| s.session_message_type)
        .unwrap_or(0);

    // Extract payload bytes from the Content oneof.
    // The supernode appends an 8-byte LE sequence number to relayed payloads;
    // strip it so mls-rs receives clean MLS ciphertext.
    let payload = match &publish.msg {
        Some(content) => match &content.content_type {
            Some(slim_proto::content::ContentType::AppPayload(app)) => strip_seq_suffix(&app.blob),
            Some(slim_proto::content::ContentType::CommandPayload(cmd)) => {
                strip_seq_suffix(&extract_mls_content(cmd))
            }
            None => Vec::new(),
        },
        None => Vec::new(),
    };

    if payload.len() > MAX_INCOMING_PAYLOAD_SIZE {
        tracing::warn!(size = payload.len(), "incoming payload too large, dropping");
        return None;
    }

    Some(IncomingMessage {
        channel_id,
        session_message_type,
        payload,
    })
}

/// Strip the trailing 8-byte LE sequence number appended by the supernode.
///
/// Returns the original bytes unchanged if the payload is too short (≤8 bytes).
fn strip_seq_suffix(data: &[u8]) -> Vec<u8> {
    if data.len() > 8 {
        data[..data.len() - 8].to_vec()
    } else {
        data.to_vec()
    }
}

/// Extract MLS content bytes from command payloads that carry them.
fn extract_mls_content(cmd: &slim_proto::CommandPayload) -> Vec<u8> {
    use slim_proto::command_payload::CommandPayloadType;
    match &cmd.command_payload_type {
        Some(CommandPayloadType::GroupAdd(p)) => p
            .mls
            .as_ref()
            .map(|m| m.mls_content.clone())
            .unwrap_or_default(),
        Some(CommandPayloadType::GroupRemove(p)) => p
            .mls
            .as_ref()
            .map(|m| m.mls_content.clone())
            .unwrap_or_default(),
        Some(CommandPayloadType::GroupWelcome(p)) => p
            .mls
            .as_ref()
            .map(|m| m.mls_content.clone())
            .unwrap_or_default(),
        Some(CommandPayloadType::GroupProposal(p)) => p.mls_proposal.clone(),
        Some(CommandPayloadType::JoinReply(p)) => p.key_package.clone().unwrap_or_default(),
        _ => Vec::new(),
    }
}
