pub mod session;

use std::env;
use std::time::Duration;

use anyhow::{bail, Context, Result};
use clap::Parser;
use futures_util::stream::SplitSink;
use futures_util::{SinkExt, Stream, StreamExt};
use serde::{Deserialize, Serialize};
use session::{Capabilities, CommandName, TxMessage};
use tokio::net::TcpStream;
use tokio::sync::mpsc::{self, Sender};
use tokio::time;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};

/// Incoming messages from the Relay service.
#[derive(Serialize, Deserialize)]
#[serde(tag = "t", content = "c")]
enum RecvMessage {
    #[serde(rename = "create_session")]
    CreateSession { id: String },
    #[serde(rename = "end_session")]
    EndSession { id: String },
    #[serde(rename = "command")]
    Command { session_id: String, command: String },
    #[serde(rename = "request_capabilities")]
    RequestCapabilities {},
}

async fn handle_message(
    manager: &mut session::Manager,
    tx: &Sender<TxMessage>,
    msg: String,
) -> anyhow::Result<()> {
    let payload = serde_json::from_str::<RecvMessage>(&msg)?;
    match payload {
        RecvMessage::CreateSession { id } => {
            println!("creating session: {id}");
            manager.create_session(id).await
        }
        RecvMessage::EndSession { id } => {
            println!("ending session: {id}");
            manager.end_session(id).await
        }
        RecvMessage::RequestCapabilities {} => {
            println!("received capabilities request");
            send_capabilities(tx).await
        }
        RecvMessage::Command {
            session_id,
            command,
        } => manager.send_command(session_id, command).await,
    }
}

/// Sends the capabilities that this node has to the Relay service.
async fn send_capabilities(tx: &Sender<TxMessage>) -> Result<()> {
    let cap = TxMessage::Capabilities {
        shell: true,
        actions: None,
    };
    tx.send(cap).await?;
    Ok(())
}

/// reads input from websockets and handles launching a shell process for a session.
async fn input_loop(
    mut rx: impl Stream<Item = Result<Message, tokio_tungstenite::tungstenite::Error>> + Unpin,
    tx: &Sender<TxMessage>,
    mut manager: session::Manager,
) {
    loop {
        let msg = rx.next().await;
        match msg {
            Some(Ok(Message::Text(m))) => {
                println!("got a string: {m}");
                if let Err(e) = handle_message(&mut manager, tx, m).await {
                    println!("error handling message: {e}")
                }
            }
            Some(Ok(Message::Close(_))) => {
                println!("closing connection");
                break;
            }
            Some(Ok(v)) => {
                println!("unhandled type: {v}");
                break;
            }
            Some(Err(e)) => {
                println!("recv error: {}", e)
            }
            None => {
                println!("nothing to read");
                break;
            }
        };
    }

    println!("write_input loop complete")
}

async fn send_output(
    tx: &mut SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
    m: session::TxMessage,
) -> anyhow::Result<()> {
    let m_str = serde_json::to_string(&m)?;
    let msg = Message::Text(m_str);
    tx.send(msg).await?;
    Ok(())
}

fn make_command(command: String) -> Result<session::CommandName> {
    let commandsplit: Vec<&str> = command.split_whitespace().collect();
    match commandsplit.split_first() {
        Some((name, args)) => Ok(session::CommandName {
            name: name.to_string(),
            args: args.iter().map(|a| a.to_string()).collect(),
        }),
        None => bail!("provided command was empty"),
    }
}

/// Run an audited shell command
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None, trailing_var_arg=true)]
struct Args {
    command: Option<Vec<String>>,
    /// Name of the service to register with the relay
    #[arg(short, long)]
    service: String,
}

type CapabilitiesFetcher = fn() -> Result<Capabilities>;

pub struct Agent {
    pub service: String,
    pub url: String,
    pub command: CommandName,
    pub capabilities_fetcher: CapabilitiesFetcher,
}

impl Agent {
    pub fn from_env(capabilities_fetcher: CapabilitiesFetcher) -> Result<Self> {
        let service =
            env::var("COMMONFATE_SERVICE").context("COMMONFATE_SERVICE must be provided")?;
        let url = env::var("COMMONFATE_URL").context("COMMONFATE_URL must be provided")?;
        let commandstr =
            env::var("COMMONFATE_COMMAND").context("COMMONFATE_COMMAND must be provided")?;

        let command = make_command(commandstr)?;

        Ok(Self {
            service,
            url,
            command,
            capabilities_fetcher,
        })
    }

    pub async fn run(&self) {
        let url = format!("{}?service={}", self.url, self.service);

        let (tx, mut rx) = mpsc::channel(100);
        let tx2 = tx.clone();
        let tx3 = tx.clone();

        // connect to the websocket server
        let (ws_stream, _) = connect_async(url).await.unwrap();
        let (mut ws_write, ws_read) = ws_stream.split();

        let manager = session::Manager::new(tx, self.command.clone());

        // write received messages back to the websocket
        tokio::spawn(async move {
            while let Some(m) = rx.recv().await {
                println!("sending message to relay: {:?}", m);
                if let Err(e) = send_output(&mut ws_write, m).await {
                    println!("send error: {e}");
                }
            }
        });

        // report capabilities after a short delay
        tokio::spawn(async move {
            time::sleep(Duration::from_secs(1)).await;
            println!("sending capabilities");
            if let Err(e) = send_capabilities(&tx3).await {
                println!("error sending capabilities: {e}");
            }
        });

        tokio::join!(input_loop(ws_read, &tx2, manager));
    }
}