use std::sync::Arc;

use async_trait::async_trait;
use dagger_sdk::{Container, Socket};

use crate::{dagger_middleware::DaggerMiddleware, leptos_service::LeptosService};

use super::RustService;

pub struct SshAgent {
    socket: Socket,
}

impl SshAgent {
    pub fn new(socket: Socket) -> Self {
        Self { socket }
    }
}

#[async_trait]
impl DaggerMiddleware for SshAgent {
    async fn handle(&self, container: Container) -> eyre::Result<Container> {
        let c = container
            .with_unix_socket("/tmp/ssh_socket", self.socket.clone())
            .with_env_variable("SSH_AUTH_SOCK", "/tmp/ssh_socket");

        Ok(c)
    }
}

pub trait SshAgentExt {
    fn with_ssh_agent(&mut self, socket: dagger_sdk::Socket) -> &mut Self {
        self
    }
}

impl SshAgentExt for RustService {
    fn with_ssh_agent(&mut self, socket: dagger_sdk::Socket) -> &mut Self {
        self.with_stage(super::RustServiceStage::BeforeDeps(Arc::new(
            SshAgent::new(socket.clone()),
        )))
        .with_stage(super::RustServiceStage::AfterBase(Arc::new(SshAgent::new(
            socket.clone(),
        ))))
    }
}

impl SshAgentExt for LeptosService {
    fn with_ssh_agent(&mut self, socket: dagger_sdk::Socket) -> &mut Self {
        self.with_stage(super::RustServiceStage::BeforeDeps(Arc::new(
            SshAgent::new(socket.clone()),
        )))
        .with_stage(super::RustServiceStage::AfterBase(Arc::new(SshAgent::new(
            socket.clone(),
        ))))
    }
}