use std::sync::Arc;

use anyhow::Result;
use dashmap::DashMap;
use log::info;
use tici::{
    ShardInfo,
    config::MetaServiceConfig,
    meta::{
        Bootstrap, MetaServiceCore,
        server::MetaServiceServer,
        writer::{ShardEvent, ShardEventSubscriber, ShardEventSubscriberTrait},
    },
};
use tici_grpc::{GrpcMetaServiceClient, GrpcMetaServiceInternalClient};
use tonic::transport::Channel;

#[derive(Debug, Default)]
struct MockEventSub {
    pub shards: Arc<DashMap<u64, ShardInfo>>,
}

impl ShardEventSubscriberTrait for MockEventSub {
    async fn on_event(&mut self, e: ShardEvent) {
        use ShardEvent::*;
        match e {
            FullSync(_, vs) => {
                self.shards.clear();
                log::info!("ShardEventSubscriberTrait restored {}", vs.len());
                for s in vs {
                    self.shards.insert(s.shard_id, s);
                }
            }
            Update(_, s) => {
                let shard_id = s.shard_id;
                self.shards.insert(shard_id, s);
            }
            Append(_, append) => {
                let shard_id = append.shard_id;
                self.shards.get_mut(&shard_id).unwrap().apply_append(append);
            }
            Add(_, s) => {
                let shard_id = s.shard_id;
                self.shards.insert(shard_id, s);
            }
            Remove(_, s) => {
                for shard_id in s {
                    self.shards.remove(&shard_id);
                }
            }
            Split(_, o, vs) => {
                self.shards.remove(&o.shard_id);
                for s in vs {
                    let shard_id = s.shard_id;
                    self.shards.insert(shard_id, s);
                }
            }
            Merge(_, vs, n) => {
                for s in vs {
                    self.shards.remove(&s.shard_id);
                }
                let shard_id = n.shard_id;
                self.shards.insert(shard_id, n);
            }
            Stop => {}
        }
    }
}

pub struct MetaServiceHandle {
    pub config: MetaServiceConfig,
    server: MetaServiceServer,
    client: GrpcMetaServiceClient<Channel>,
    internal_client: GrpcMetaServiceInternalClient<Channel>,
    grpc_server_close_tx: Option<tokio::sync::oneshot::Sender<()>>,
    pub meta_service: Arc<MetaServiceCore>,
    pub subscriber: Option<ShardEventSubscriber>,
    sub_shards: Option<Arc<DashMap<u64, ShardInfo>>>,
}

impl MetaServiceHandle {
    pub async fn new(config: MetaServiceConfig, reset_meta: bool) -> Result<Self> {
        let config_clone = config.clone();
        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        let mut server = MetaServiceServer::new(config_clone).await.unwrap();
        info!("wait become leader");
        server.wait_become_leader().await;
        let meta_service = server.meta_service.clone();
        if reset_meta {
            Bootstrap::new_with_config(&config)
                .clear_all_meta(true)
                .await?;
        }

        let mock_event_sub = MockEventSub::default();
        let sub_shards = mock_event_sub.shards.clone();
        let subscriber = meta_service
            .shard_task_manager
            .event_pubsub
            .resubscribe(mock_event_sub, "test".to_string());

        info!("MetaServiceHandle: start server");
        server.start(rx).await?;
        let client = MetaServiceServer::get_client(
            config.server.get_advertise_addr(),
            &config.server.grpc_client,
        )
        .await;
        let internal_client = MetaServiceServer::get_internal_client(
            config.server.get_advertise_addr(),
            &config.server.grpc_client,
        )
        .await;
        Ok(Self {
            config,
            server,
            client,
            internal_client,
            grpc_server_close_tx: Some(tx),
            meta_service,
            subscriber: Some(subscriber),
            sub_shards: Some(sub_shards),
        })
    }

    pub async fn stop(mut self) -> Result<()> {
        self.server
            .stop(self.grpc_server_close_tx.take().unwrap())
            .await?;
        Ok(())
    }

    pub fn get_meta_internal_client(&self) -> GrpcMetaServiceInternalClient<Channel> {
        self.internal_client.clone()
    }

    pub fn get_meta_client(&self) -> GrpcMetaServiceClient<Channel> {
        self.client.clone()
    }

    pub fn get_subscribed_shards(&self) -> Arc<DashMap<u64, ShardInfo>> {
        self.sub_shards.as_ref().unwrap().clone()
    }
}