use crate::cluster::Cluster;
use crate::tso::TsoOracle;
use flaredb_proto::pdpb::pd_server::Pd;
use flaredb_proto::pdpb::{
    GetRegionRequest, GetRegionResponse, ListRegionsRequest, ListRegionsResponse,
    MoveRegionRequest, MoveRegionResponse, RegisterStoreRequest, RegisterStoreResponse,
};
use std::sync::Arc;
use tonic::{Request, Response, Status};

pub struct PdServiceImpl {
    _tso_oracle: Arc<TsoOracle>,
    cluster: Cluster,
}

impl PdServiceImpl {
    pub fn new(tso_oracle: Arc<TsoOracle>) -> Self {
        Self {
            _tso_oracle: tso_oracle,
            cluster: Cluster::new(),
        }
    }

    #[allow(dead_code)]
    async fn list_regions(
        &self,
        _request: Request<ListRegionsRequest>,
    ) -> Result<Response<ListRegionsResponse>, Status> {
        let (regions, stores) = self.cluster.list_regions();
        Ok(Response::new(ListRegionsResponse { regions, stores }))
    }

    #[allow(dead_code)]
    async fn move_region(
        &self,
        _request: Request<MoveRegionRequest>,
    ) -> Result<Response<MoveRegionResponse>, Status> {
        Ok(Response::new(MoveRegionResponse { accepted: true }))
    }
}

#[tonic::async_trait]
impl Pd for PdServiceImpl {
    async fn register_store(
        &self,
        request: Request<RegisterStoreRequest>,
    ) -> Result<Response<RegisterStoreResponse>, Status> {
        let req = request.into_inner();
        let requested_store_id = (req.store_id != 0).then_some(req.store_id);
        let store_id = self.cluster.register_store(req.addr, requested_store_id);
        Ok(Response::new(RegisterStoreResponse {
            store_id,
            cluster_id: 1, // fixed for now
        }))
    }

    async fn get_region(
        &self,
        request: Request<GetRegionRequest>,
    ) -> Result<Response<GetRegionResponse>, Status> {
        let req = request.into_inner();
        if let Some((region, leader)) = self.cluster.get_region_for_key(&req.key) {
            Ok(Response::new(GetRegionResponse {
                region: Some(region),
                leader: Some(leader),
            }))
        } else {
            Err(Status::not_found("Region not initialized yet"))
        }
    }

    async fn list_regions(
        &self,
        _request: Request<ListRegionsRequest>,
    ) -> Result<Response<ListRegionsResponse>, Status> {
        let (regions, stores) = self.cluster.list_regions();
        Ok(Response::new(ListRegionsResponse { regions, stores }))
    }
}