batata-client 0.0.2

use std::collections::HashMap;

use prost_types::Any;
use serde::{Deserialize, Serialize};

use crate::api::remote::{Request, RequestTrait, Response, ResponseTrait};
use crate::api::Payload;
use crate::common::{NAMING_MODULE, REGISTER_INSTANCE, DEREGISTER_INSTANCE, UPDATE_INSTANCE};

// ==================== Instance ====================

/// Service instance information
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Instance {
    pub instance_id: String,
    pub ip: String,
    pub port: i32,
    pub weight: f64,
    pub healthy: bool,
    pub enabled: bool,
    pub ephemeral: bool,
    pub cluster_name: String,
    pub service_name: String,
    pub metadata: HashMap<String, String>,
    pub instance_heart_beat_interval: i64,
    pub instance_heart_beat_time_out: i64,
    pub ip_delete_timeout: i64,
    pub instance_id_generator: String,
}

impl Instance {
    pub fn new(ip: &str, port: i32) -> Self {
        Self {
            ip: ip.to_string(),
            port,
            weight: 1.0,
            healthy: true,
            enabled: true,
            ephemeral: true,
            cluster_name: crate::common::DEFAULT_CLUSTER.to_string(),
            metadata: HashMap::new(),
            instance_heart_beat_interval: 5000,
            instance_heart_beat_time_out: 15000,
            ip_delete_timeout: 30000,
            ..Default::default()
        }
    }

    pub fn with_weight(mut self, weight: f64) -> Self {
        self.weight = weight;
        self
    }

    pub fn with_cluster(mut self, cluster: &str) -> Self {
        self.cluster_name = cluster.to_string();
        self
    }

    pub fn with_metadata(mut self, key: &str, value: &str) -> Self {
        self.metadata.insert(key.to_string(), value.to_string());
        self
    }

    pub fn with_ephemeral(mut self, ephemeral: bool) -> Self {
        self.ephemeral = ephemeral;
        self
    }

    pub fn with_enabled(mut self, enabled: bool) -> Self {
        self.enabled = enabled;
        self
    }

    /// Generate instance ID if not set
    pub fn generate_instance_id(&mut self) {
        if self.instance_id.is_empty() {
            self.instance_id = format!("{}#{}#{}", self.ip, self.port, self.cluster_name);
        }
    }

    /// Get the instance key
    pub fn key(&self) -> String {
        format!("{}#{}#{}", self.ip, self.port, self.cluster_name)
    }
}

// ==================== Service ====================

/// Service information
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Service {
    pub name: String,
    pub group_name: String,
    pub clusters: String,
    pub cache_millis: i64,
    pub hosts: Vec<Instance>,
    pub last_ref_time: i64,
    pub checksum: String,
    pub all_ips: bool,
    pub reach_protection_threshold: bool,
}

impl Service {
    pub fn new(name: &str, group_name: &str) -> Self {
        Self {
            name: name.to_string(),
            group_name: group_name.to_string(),
            cache_millis: 10000,
            hosts: Vec::new(),
            ..Default::default()
        }
    }

    pub fn key(&self) -> String {
        format!("{}@@{}", self.group_name, self.name)
    }

    /// Get healthy instances
    pub fn healthy_instances(&self) -> Vec<&Instance> {
        self.hosts.iter().filter(|i| i.healthy && i.enabled).collect()
    }

    /// Get all enabled instances
    pub fn enabled_instances(&self) -> Vec<&Instance> {
        self.hosts.iter().filter(|i| i.enabled).collect()
    }
}

// ==================== Naming Request Base ====================

/// Base naming request
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct NamingRequest {
    #[serde(flatten)]
    pub request: Request,
    pub namespace: String,
    pub service_name: String,
    pub group_name: String,
    pub module: String,
}

impl NamingRequest {
    pub fn new(namespace: &str, service_name: &str, group_name: &str) -> Self {
        Self {
            request: Request::new(),
            namespace: namespace.to_string(),
            service_name: service_name.to_string(),
            group_name: group_name.to_string(),
            module: NAMING_MODULE.to_string(),
        }
    }
}

impl RequestTrait for NamingRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.request.headers()
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.request.request_id.clone()
    }
}

// ==================== Instance Request ====================

/// Instance registration/deregistration request
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InstanceRequest {
    #[serde(flatten)]
    pub naming_request: NamingRequest,
    pub r#type: String,
    pub instance: Instance,
}

impl InstanceRequest {
    pub fn register(namespace: &str, service_name: &str, group_name: &str, instance: Instance) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            r#type: REGISTER_INSTANCE.to_string(),
            instance,
        }
    }

    pub fn deregister(namespace: &str, service_name: &str, group_name: &str, instance: Instance) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            r#type: DEREGISTER_INSTANCE.to_string(),
            instance,
        }
    }

    pub fn update(namespace: &str, service_name: &str, group_name: &str, instance: Instance) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            r#type: UPDATE_INSTANCE.to_string(),
            instance,
        }
    }
}

impl RequestTrait for InstanceRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.naming_request.headers()
    }

    fn request_type(&self) -> &'static str {
        "InstanceRequest"
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.naming_request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.naming_request.request_id()
    }
}

/// Instance response
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InstanceResponse {
    #[serde(flatten)]
    pub response: Response,
    pub r#type: String,
}

impl ResponseTrait for InstanceResponse {
    fn response_type(&self) -> &'static str {
        "InstanceResponse"
    }

    fn set_request_id(&mut self, request_id: String) {
        self.response.request_id = request_id;
    }

    fn error_code(&self) -> i32 {
        self.response.error_code
    }

    fn result_code(&self) -> i32 {
        self.response.result_code
    }

    fn message(&self) -> String {
        self.response.message.clone()
    }
}

// ==================== Batch Instance Request ====================

/// Batch instance request
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BatchInstanceRequest {
    #[serde(flatten)]
    pub naming_request: NamingRequest,
    pub r#type: String,
    pub instances: Vec<Instance>,
}

impl BatchInstanceRequest {
    pub fn register(namespace: &str, service_name: &str, group_name: &str, instances: Vec<Instance>) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            r#type: REGISTER_INSTANCE.to_string(),
            instances,
        }
    }

    pub fn deregister(namespace: &str, service_name: &str, group_name: &str, instances: Vec<Instance>) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            r#type: DEREGISTER_INSTANCE.to_string(),
            instances,
        }
    }
}

impl RequestTrait for BatchInstanceRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.naming_request.headers()
    }

    fn request_type(&self) -> &'static str {
        "BatchInstanceRequest"
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.naming_request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.naming_request.request_id()
    }
}

/// Batch instance response
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BatchInstanceResponse {
    #[serde(flatten)]
    pub response: Response,
    pub r#type: String,
}

impl ResponseTrait for BatchInstanceResponse {
    fn response_type(&self) -> &'static str {
        "BatchInstanceResponse"
    }

    fn set_request_id(&mut self, request_id: String) {
        self.response.request_id = request_id;
    }

    fn error_code(&self) -> i32 {
        self.response.error_code
    }

    fn result_code(&self) -> i32 {
        self.response.result_code
    }

    fn message(&self) -> String {
        self.response.message.clone()
    }
}

// ==================== Service Query ====================

/// Service query request
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServiceQueryRequest {
    #[serde(flatten)]
    pub naming_request: NamingRequest,
    pub cluster: String,
    pub healthy_only: bool,
    pub udp_port: i32,
}

impl ServiceQueryRequest {
    pub fn new(namespace: &str, service_name: &str, group_name: &str) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            cluster: String::new(),
            healthy_only: false,
            udp_port: 0,
        }
    }

    pub fn with_cluster(mut self, cluster: &str) -> Self {
        self.cluster = cluster.to_string();
        self
    }

    pub fn with_healthy_only(mut self, healthy_only: bool) -> Self {
        self.healthy_only = healthy_only;
        self
    }
}

impl RequestTrait for ServiceQueryRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.naming_request.headers()
    }

    fn request_type(&self) -> &'static str {
        "ServiceQueryRequest"
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.naming_request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.naming_request.request_id()
    }
}

/// Query service response
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct QueryServiceResponse {
    #[serde(flatten)]
    pub response: Response,
    pub service_info: Service,
}

impl ResponseTrait for QueryServiceResponse {
    fn response_type(&self) -> &'static str {
        "QueryServiceResponse"
    }

    fn set_request_id(&mut self, request_id: String) {
        self.response.request_id = request_id;
    }

    fn error_code(&self) -> i32 {
        self.response.error_code
    }

    fn result_code(&self) -> i32 {
        self.response.result_code
    }

    fn message(&self) -> String {
        self.response.message.clone()
    }
}

// ==================== Service List ====================

/// Service list request
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServiceListRequest {
    #[serde(flatten)]
    pub naming_request: NamingRequest,
    pub page_no: i32,
    pub page_size: i32,
    pub selector: String,
}

impl ServiceListRequest {
    pub fn new(namespace: &str, group_name: &str) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, "", group_name),
            page_no: 1,
            page_size: 100,
            selector: String::new(),
        }
    }

    pub fn with_page(mut self, page_no: i32, page_size: i32) -> Self {
        self.page_no = page_no;
        self.page_size = page_size;
        self
    }
}

impl RequestTrait for ServiceListRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.naming_request.headers()
    }

    fn request_type(&self) -> &'static str {
        "ServiceListRequest"
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.naming_request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.naming_request.request_id()
    }
}

/// Service list response
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServiceListResponse {
    #[serde(flatten)]
    pub response: Response,
    pub count: i32,
    pub service_names: Vec<String>,
}

impl ResponseTrait for ServiceListResponse {
    fn response_type(&self) -> &'static str {
        "ServiceListResponse"
    }

    fn set_request_id(&mut self, request_id: String) {
        self.response.request_id = request_id;
    }

    fn error_code(&self) -> i32 {
        self.response.error_code
    }

    fn result_code(&self) -> i32 {
        self.response.result_code
    }

    fn message(&self) -> String {
        self.response.message.clone()
    }
}

// ==================== Subscribe Service ====================

/// Subscribe service request
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SubscribeServiceRequest {
    #[serde(flatten)]
    pub naming_request: NamingRequest,
    pub subscribe: bool,
    pub clusters: String,
}

impl SubscribeServiceRequest {
    pub fn subscribe(namespace: &str, service_name: &str, group_name: &str) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            subscribe: true,
            clusters: String::new(),
        }
    }

    pub fn unsubscribe(namespace: &str, service_name: &str, group_name: &str) -> Self {
        Self {
            naming_request: NamingRequest::new(namespace, service_name, group_name),
            subscribe: false,
            clusters: String::new(),
        }
    }

    pub fn with_clusters(mut self, clusters: &str) -> Self {
        self.clusters = clusters.to_string();
        self
    }
}

impl RequestTrait for SubscribeServiceRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.naming_request.headers()
    }

    fn request_type(&self) -> &'static str {
        "SubscribeServiceRequest"
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.naming_request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.naming_request.request_id()
    }
}

/// Subscribe service response
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SubscribeServiceResponse {
    #[serde(flatten)]
    pub response: Response,
    pub service_info: Service,
}

impl ResponseTrait for SubscribeServiceResponse {
    fn response_type(&self) -> &'static str {
        "SubscribeServiceResponse"
    }

    fn set_request_id(&mut self, request_id: String) {
        self.response.request_id = request_id;
    }

    fn error_code(&self) -> i32 {
        self.response.error_code
    }

    fn result_code(&self) -> i32 {
        self.response.result_code
    }

    fn message(&self) -> String {
        self.response.message.clone()
    }
}

// ==================== Notify Subscriber ====================

/// Notify subscriber request (from server)
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct NotifySubscriberRequest {
    #[serde(flatten)]
    pub request: Request,
    pub namespace: String,
    pub service_name: String,
    pub group_name: String,
    pub service_info: Service,
    pub module: String,
}

impl RequestTrait for NotifySubscriberRequest {
    fn headers(&self) -> HashMap<String, String> {
        self.request.headers()
    }

    fn request_type(&self) -> &'static str {
        "NotifySubscriberRequest"
    }

    fn insert_headers(&mut self, headers: HashMap<String, String>) {
        self.request.insert_headers(headers);
    }

    fn request_id(&self) -> String {
        self.request.request_id()
    }
}

/// Notify subscriber response
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct NotifySubscriberResponse {
    #[serde(flatten)]
    pub response: Response,
}

impl NotifySubscriberResponse {
    pub fn new() -> Self {
        Self {
            response: Response::new(),
        }
    }

    pub fn to_payload(&self, request_id: &str) -> Payload {
        let mut headers = HashMap::new();
        headers.insert("requestId".to_string(), request_id.to_string());

        let body = serde_json::to_vec(self).unwrap_or_default();

        Payload {
            metadata: Some(crate::api::Metadata {
                r#type: "NotifySubscriberResponse".to_string(),
                client_ip: String::new(),
                headers,
            }),
            body: Some(Any {
                type_url: String::new(),
                value: body,
            }),
        }
    }
}

impl ResponseTrait for NotifySubscriberResponse {
    fn response_type(&self) -> &'static str {
        "NotifySubscriberResponse"
    }

    fn set_request_id(&mut self, request_id: String) {
        self.response.request_id = request_id;
    }

    fn error_code(&self) -> i32 {
        self.response.error_code
    }

    fn result_code(&self) -> i32 {
        self.response.result_code
    }
}