// Copyright 2020 Andy Grove
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Support for distributed schedulers, such as Kubernetes

pub mod execution_plans;
pub mod planner;
pub mod state;

use std::convert::TryInto;
use std::ffi::OsStr;
use std::fmt;

use crate::serde::protobuf::{
    execute_query_params::Query, job_status, scheduler_grpc_server::SchedulerGrpc, CompletedJob,
    ExecuteQueryParams, ExecuteQueryResult, ExecuteSqlParams, ExecutorMetadata, FailedJob,
    FilePartitionMetadata, FileType, GetExecutorMetadataParams, GetExecutorMetadataResult,
    GetFileMetadataParams, GetFileMetadataResult, GetJobStatusParams, GetJobStatusResult,
    JobStatus, PartitionLocation, QueuedJob, RegisterExecutorParams, RegisterExecutorResult,
    RunningJob,
};
use crate::serde::scheduler::ExecutorMeta;

use clap::arg_enum;
use datafusion::physical_plan::ExecutionPlan;

// an enum used to configure the backend
// needs to be visible to code generated by configure_me
arg_enum! {
    #[derive(Debug, serde::Deserialize)]
    pub enum ConfigBackend {
        Etcd,
        Standalone
    }
}

impl parse_arg::ParseArgFromStr for ConfigBackend {
    fn describe_type<W: fmt::Write>(mut writer: W) -> fmt::Result {
        write!(writer, "The configuration backend for the scheduler")
    }
}

use crate::{client::BallistaClient, error::Result, serde::scheduler::Action};
use crate::{prelude::BallistaError, scheduler::planner::DistributedPlanner};
use execution_plans::ShuffleReaderExec;

use arrow::datatypes::{Schema, SchemaRef};
use datafusion::execution::context::{ExecutionConfig, ExecutionContext};
use log::{debug, error, info, warn};
use rand::{distributions::Alphanumeric, thread_rng, Rng};
use tonic::{Request, Response};

use self::state::{ConfigBackendClient, SchedulerState};
use crate::utils::format_plan;
use datafusion::physical_plan::parquet::ParquetExec;
use std::time::Instant;

pub struct SchedulerServer<Config: ConfigBackendClient> {
    state: SchedulerState<Config>,
    namespace: String,
}

impl<Config: ConfigBackendClient> SchedulerServer<Config> {
    pub fn new(config: Config, namespace: String) -> Self {
        Self {
            state: SchedulerState::new(config),
            namespace,
        }
    }
}

#[tonic::async_trait]
impl<T: ConfigBackendClient + Send + Sync + 'static> SchedulerGrpc for SchedulerServer<T> {
    async fn get_executors_metadata(
        &self,
        _request: Request<GetExecutorMetadataParams>,
    ) -> std::result::Result<Response<GetExecutorMetadataResult>, tonic::Status> {
        info!("Received get_executors_metadata request");
        let result = self
            .state
            .get_executors_metadata(self.namespace.as_str())
            .await
            .map_err(|e| {
                let msg = format!("Error reading executors metadata: {}", e);
                error!("{}", msg);
                tonic::Status::internal(msg)
            })?
            .into_iter()
            .map(|meta| meta.into())
            .collect();
        Ok(Response::new(GetExecutorMetadataResult {
            metadata: result,
        }))
    }

    async fn register_executor(
        &self,
        request: Request<RegisterExecutorParams>,
    ) -> std::result::Result<Response<RegisterExecutorResult>, tonic::Status> {
        if let RegisterExecutorParams {
            metadata: Some(metadata),
        } = request.into_inner()
        {
            info!("Received register_executor request for {:?}", metadata);
            self.state
                .save_executor_metadata(&self.namespace, metadata.into())
                .await
                .map_err(|e| {
                    let msg = format!("Could not save executor metadata: {}", e);
                    error!("{}", msg);
                    tonic::Status::internal(msg)
                })?;
            Ok(Response::new(RegisterExecutorResult {}))
        } else {
            warn!("Received invalid executor registration request");
            Err(tonic::Status::invalid_argument(
                "Missing metadata in request",
            ))
        }
    }

    async fn get_file_metadata(
        &self,
        request: Request<GetFileMetadataParams>,
    ) -> std::result::Result<Response<GetFileMetadataResult>, tonic::Status> {
        let GetFileMetadataParams { path, file_type } = request.into_inner();

        let file_type: FileType = file_type.try_into().map_err(|e| {
            let msg = format!("Error reading request: {}", e);
            error!("{}", msg);
            tonic::Status::internal(msg)
        })?;

        match file_type {
            FileType::Parquet => {
                let parquet_exec =
                    ParquetExec::try_from_path(&path, None, None, 1024, 1).map_err(|e| {
                        let msg = format!("Error opening parquet files: {}", e);
                        error!("{}", msg);
                        tonic::Status::internal(msg)
                    })?;

                //TODO include statistics and any other info needed to reconstruct ParquetExec
                Ok(Response::new(GetFileMetadataResult {
                    schema: Some(parquet_exec.schema().as_ref().into()),
                    partitions: parquet_exec
                        .partitions()
                        .iter()
                        .map(|part| FilePartitionMetadata {
                            filename: part.filenames().to_vec(),
                        })
                        .collect(),
                }))
            }
            //TODO implement for CSV
            _ => Err(tonic::Status::unimplemented(
                "get_file_metadata unsupported file type",
            )),
        }
    }

    async fn execute_query(
        &self,
        request: Request<ExecuteQueryParams>,
    ) -> std::result::Result<Response<ExecuteQueryResult>, tonic::Status> {
        if let ExecuteQueryParams { query: Some(query) } = request.into_inner() {
            let plan = match query {
                Query::LogicalPlan(logical_plan) => {
                    // parse protobuf
                    (&logical_plan).try_into().map_err(|e| {
                        let msg = format!("Could not parse logical plan protobuf: {}", e);
                        error!("{}", msg);
                        tonic::Status::internal(msg)
                    })?
                }
                Query::Sql(sql) => {
                    //TODO we can't just create a new context because we need a context that has
                    // tables registered from previous SQL statements that have been executed
                    let mut ctx = ExecutionContext::new();
                    let df = ctx.sql(&sql).map_err(|e| {
                        let msg = format!("Error parsing SQL: {}", e);
                        error!("{}", msg);
                        tonic::Status::internal(msg)
                    })?;
                    df.to_logical_plan()
                }
            };
            debug!("Received plan for execution: {:?}", plan);
            let executors = self
                .state
                .get_executors_metadata(&self.namespace)
                .await
                .map_err(|e| {
                    let msg = format!("Error reading executors metadata: {}", e);
                    error!("{}", msg);
                    tonic::Status::internal(msg)
                })?;
            debug!("Found executors: {:?}", executors);

            let job_id: String = {
                let mut rng = thread_rng();
                std::iter::repeat(())
                    .map(|()| rng.sample(Alphanumeric))
                    .map(char::from)
                    .take(7)
                    .collect()
            };

            // Save placeholder job metadata
            self.state
                .save_job_metadata(
                    &self.namespace,
                    &job_id,
                    &JobStatus {
                        status: Some(job_status::Status::Queued(QueuedJob {})),
                    },
                )
                .await
                .map_err(|e| {
                    tonic::Status::internal(format!("Could not save job metadata: {}", e))
                })?;

            // TODO: handle errors once we have more job metadata
            let namespace = self.namespace.to_owned();
            let state = self.state.clone();
            let job_id_spawn = job_id.clone();
            tokio::spawn(async move {
                // create physical plan using DataFusion
                let datafusion_ctx = ExecutionContext::new();
                macro_rules! fail_job {
                    ($code :expr) => {{
                        match $code {
                            Err(error) => {
                                warn!("Job {} failed with {}", job_id_spawn, error);
                                state
                                    .save_job_metadata(
                                        &namespace,
                                        &job_id_spawn,
                                        &JobStatus {
                                            status: Some(job_status::Status::Failed(FailedJob {
                                                error: format!("{}", error),
                                            })),
                                        },
                                    )
                                    .await
                                    .unwrap();
                                return;
                            }
                            Ok(value) => value,
                        }
                    }};
                };

                let start = Instant::now();

                let plan = fail_job!(datafusion_ctx
                    .optimize(&plan)
                    .and_then(|plan| datafusion_ctx.create_physical_plan(&plan))
                    .map_err(|e| {
                        let msg = format!("Could not create physical plan: {}", e);
                        error!("{}", msg);
                        tonic::Status::internal(msg)
                    }));

                info!(
                    "DataFusion created physical plan in {} seconds",
                    start.elapsed().as_secs(),
                );

                // create distributed physical plan using Ballista
                if let Err(e) = state
                    .save_job_metadata(
                        &namespace,
                        &job_id_spawn,
                        &JobStatus {
                            status: Some(job_status::Status::Running(RunningJob {})),
                        },
                    )
                    .await
                {
                    warn!(
                        "Could not update job {} status to running: {}",
                        job_id_spawn, e
                    );
                }
                let mut planner = fail_job!(DistributedPlanner::try_new(executors).map_err(|e| {
                    let msg = format!("Could not create distributed planner: {}", e);
                    error!("{}", msg);
                    tonic::Status::internal(msg)
                }));
                let plan = fail_job!(planner.execute_distributed_query(plan).await.map_err(|e| {
                    let msg = format!("Could not execute distributed plan: {}", e);
                    error!("{}", msg);
                    tonic::Status::internal(msg)
                }));

                // save partition info into job's state
                let plan = plan
                    .as_any()
                    .downcast_ref::<ShuffleReaderExec>()
                    .expect("Expected plan final operator to be ShuffleReaderExec");
                let mut partition_location = vec![];
                for loc in &plan.partition_location {
                    partition_location.push(loc.clone().try_into().unwrap());
                }
                state
                    .save_job_metadata(
                        &namespace,
                        &job_id_spawn,
                        &JobStatus {
                            status: Some(job_status::Status::Completed(CompletedJob {
                                partition_location,
                            })),
                        },
                    )
                    .await
                    .unwrap();
            });

            Ok(Response::new(ExecuteQueryResult { job_id }))
        } else {
            Err(tonic::Status::internal("Error parsing request"))
        }
    }

    async fn get_job_status(
        &self,
        request: Request<GetJobStatusParams>,
    ) -> std::result::Result<Response<GetJobStatusResult>, tonic::Status> {
        let job_id = request.into_inner().job_id;
        debug!("Received get_job_status request for job {}", job_id);
        let job_meta = self
            .state
            .get_job_metadata(&self.namespace, &job_id)
            .await
            .map_err(|e| {
                let msg = format!("Error reading job metadata: {}", e);
                error!("{}", msg);
                tonic::Status::internal(msg)
            })?;
        Ok(Response::new(GetJobStatusResult {
            status: Some(job_meta),
        }))
    }
}
