// Copyright 2024 Golem Cloud
//
// 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.

use async_trait::async_trait;
use clap::builder::ValueParser;
use clap::Subcommand;
use golem_client::model::InvokeParameters;

use crate::clients::worker::WorkerClient;
use crate::model::{
    GolemError, GolemResult, InvocationKey, JsonValueParser, TemplateIdOrName, WorkerName,
};
use crate::parse_key_val;
use crate::template::TemplateHandler;

#[derive(Subcommand, Debug)]
#[command()]
pub enum WorkerSubcommand {
    #[command()]
    Add {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,

        #[arg(short, long, value_parser = parse_key_val, value_name = "ENV=VAL")]
        env: Vec<(String, String)>,

        #[arg(value_name = "args")]
        args: Vec<String>,
    },
    #[command()]
    InvocationKey {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,
    },
    #[command()]
    InvokeAndAwait {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,

        #[arg(short = 'k', long)]
        invocation_key: Option<InvocationKey>,

        #[arg(short, long)]
        function: String,

        #[arg(short = 'j', long, value_name = "json", value_parser = ValueParser::new(JsonValueParser))]
        parameters: serde_json::value::Value,

        #[arg(short = 's', long, default_value_t = false)]
        use_stdio: bool,
    },
    #[command()]
    Invoke {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,

        #[arg(short, long)]
        function: String,

        #[arg(short = 'j', long, value_name = "json", value_parser = ValueParser::new(JsonValueParser))]
        parameters: serde_json::value::Value,
    },
    #[command()]
    Connect {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,
    },
    #[command()]
    Interrupt {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,
    },
    #[command()]
    SimulatedCrash {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,
    },
    #[command()]
    Delete {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,
    },
    #[command()]
    Get {
        #[command(flatten)]
        template_id_or_name: TemplateIdOrName,

        #[arg(short, long)]
        worker_name: WorkerName,
    },
}

#[async_trait]
pub trait WorkerHandler {
    async fn handle(&self, subcommand: WorkerSubcommand) -> Result<GolemResult, GolemError>;
}

pub struct WorkerHandlerLive<'r, C: WorkerClient + Send + Sync, R: TemplateHandler + Send + Sync> {
    pub client: C,
    pub templates: &'r R,
}

#[async_trait]
impl<'r, C: WorkerClient + Send + Sync, R: TemplateHandler + Send + Sync> WorkerHandler
    for WorkerHandlerLive<'r, C, R>
{
    async fn handle(&self, subcommand: WorkerSubcommand) -> Result<GolemResult, GolemError> {
        match subcommand {
            WorkerSubcommand::Add {
                template_id_or_name,
                worker_name,
                env,
                args,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                let inst = self
                    .client
                    .new_worker(worker_name, template_id, args, env)
                    .await?;

                Ok(GolemResult::Ok(Box::new(inst)))
            }
            WorkerSubcommand::InvocationKey {
                template_id_or_name,
                worker_name,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                let key = self
                    .client
                    .get_invocation_key(&worker_name, &template_id)
                    .await?;

                Ok(GolemResult::Ok(Box::new(key)))
            }
            WorkerSubcommand::InvokeAndAwait {
                template_id_or_name,
                worker_name,
                invocation_key,
                function,
                parameters,
                use_stdio,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                let invocation_key = match invocation_key {
                    None => {
                        self.client
                            .get_invocation_key(&worker_name, &template_id)
                            .await?
                    }
                    Some(key) => key,
                };

                let res = self
                    .client
                    .invoke_and_await(
                        worker_name,
                        template_id,
                        function,
                        InvokeParameters { params: parameters },
                        invocation_key,
                        use_stdio,
                    )
                    .await?;

                Ok(GolemResult::Json(res.result))
            }
            WorkerSubcommand::Invoke {
                template_id_or_name,
                worker_name,
                function,
                parameters,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                self.client
                    .invoke(
                        worker_name,
                        template_id,
                        function,
                        InvokeParameters { params: parameters },
                    )
                    .await?;

                Ok(GolemResult::Str("Invoked".to_string()))
            }
            WorkerSubcommand::Connect {
                template_id_or_name,
                worker_name,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                let result = self.client.connect(worker_name, template_id).await;

                match result {
                    Ok(_) => Err(GolemError("Unexpected connection closure".to_string())),
                    Err(err) => Err(GolemError(err.to_string())),
                }
            }
            WorkerSubcommand::Interrupt {
                template_id_or_name,
                worker_name,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                self.client.interrupt(worker_name, template_id).await?;

                Ok(GolemResult::Str("Interrupted".to_string()))
            }
            WorkerSubcommand::SimulatedCrash {
                template_id_or_name,
                worker_name,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                self.client
                    .simulated_crash(worker_name, template_id)
                    .await?;

                Ok(GolemResult::Str("Done".to_string()))
            }
            WorkerSubcommand::Delete {
                template_id_or_name,
                worker_name,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                self.client.delete(worker_name, template_id).await?;

                Ok(GolemResult::Str("Deleted".to_string()))
            }
            WorkerSubcommand::Get {
                template_id_or_name,
                worker_name,
            } => {
                let template_id = self.templates.resolve_id(template_id_or_name).await?;

                let mata = self.client.get_metadata(worker_name, template_id).await?;

                Ok(GolemResult::Ok(Box::new(mata)))
            }
        }
    }
}