use tonic::{Request, Response, Status};
use crate::coordinator::coordinator_client::CoordinatorClient as GeneratedCoordinatorClient;
use crate::coordinator::{RegisterRequest, RegisterResponse, HeartbeatRequest, HeartbeatResponse, PeerRequest, PeerResponse};

pub struct CoordinatorClient {
    client: GeneratedCoordinatorClient<tonic::transport::Channel>,
}

impl CoordinatorClient {
    pub async fn new(coordinator_addr: String) -> Result<Self, tonic::transport::Error> {
        let client = GeneratedCoordinatorClient::connect(coordinator_addr).await?;
        Ok(Self { client })
    }

    pub async fn register_node(
        &mut self,
        network_key: String,
        node_id: String,
        public_key: String,
        local_ip: String,
    ) -> Result<RegisterResponse, Status> {
        let request = Request::new(RegisterRequest {
            network_key,
            node_id,
            public_key,
            local_ip,
        });
        let response = self.client.register_node(request).await?;
        Ok(response.into_inner())
    }

    pub async fn heartbeat(
        &mut self,
        node_id: String,
        session_token: String,
        public_endpoint: String,
    ) -> Result<HeartbeatResponse, Status> {
        let request = Request::new(HeartbeatRequest {
            node_id,
            session_token,
            public_endpoint,
        });
        let response = self.client.heartbeat(request).await?;
        Ok(response.into_inner())
    }

    pub async fn get_peer_endpoint(
        &mut self,
        target_node_id: String,
        session_token: String,
    ) -> Result<PeerResponse, Status> {
        let request = Request::new(PeerRequest {
            target_node_id,
            session_token,
        });
        let response = self.client.get_peer_endpoint(request).await?;
        Ok(response.into_inner())
    }
}

// Optional: You can add a method to update node name if needed in the future.
// pub async fn update_node_name(&mut self, ...) -> ... { ... }