//! This example demonstrates how to use Ditto's Data Streams feature with arguments.
//!
//! It's a very naive implementation of a file server and client using Data Streams.
//!
//! The server serves files from a specified directory, and the client requests a file by name.
//!
//! The client requests a file by setting its name as the arguments for the stream it tries to
//! receive the file contents and either prints it to stdout or saves it to a specified path.
//!
//! The server reads the requested file and sends its contents to the client.

#![allow(clippy::disallowed_macros)]

use std::{io::Read, path::PathBuf, str::FromStr, sync::Arc};

use anyhow::Result;
use clap::Parser;
use dittolive_ditto::{
    prelude::*,
    preview::{
        datastreams::{Payload, Topic},
        peer_pubkey::PeerPubkey,
    },
};

#[derive(Parser, Debug)]
enum Command {
    /// Start a Data Stream server
    Serve(Serve),
    /// Fetch a file from a Data Stream server
    Fetch(Fetch),
}

#[derive(clap::Args, Debug)]
struct Serve {
    /// The path to the directory to serve.
    #[arg(value_name = "PATH")]
    path: PathBuf,
    /// The topic to bind the Data Stream Server onto.
    #[arg(long, short, default_value = "fs")]
    topic: String,
}

#[derive(clap::Args, Debug)]
struct Fetch {
    /// The file to request from the Data Stream server (as a path from the server's served
    /// directory).
    #[arg(value_name = "REQUESTED_FILE")]
    file: String,
    /// The topic to connect the Data Stream Client to.
    #[arg(long, short, default_value = "fs")]
    topic: String,
    /// The peer to connect to (if not specified, will connect to the first available peer).
    #[arg(long, short)]
    peer: Option<PeerPubkey>,
    /// The path to save the fetched file as (if not specified, the file will be printed to
    /// stdout).
    #[arg(long, short)]
    save_as: Option<PathBuf>,
}

#[tokio::main]
async fn main() -> Result<()> {
    let command = Command::parse();

    let db_id = "d0e9a40e-9beb-4c90-bb38-769269cbb6fc";

    let tempdir = tempfile::tempdir()?;
    let ditto = Ditto::open(
        DittoConfig::new(
            db_id,
            DittoConfigConnect::SmallPeersOnly { private_key: None },
        )
        .with_persistence_directory(tempdir.path()),
    )
    .await?;
    ditto.set_license_from_env("DITTO_LICENSE")?;

    ditto.sync().start()?;
    let ditto = Arc::new(ditto);

    match command {
        Command::Serve(args) => serve(ditto, args).await,
        Command::Fetch(args) => fetch(ditto, args).await,
    }
}

async fn serve(ditto: Arc<Ditto>, args: Serve) -> Result<()> {
    let Serve { path, topic } = args;

    let path = path.canonicalize()?;

    let topic = Topic::try_from(topic.as_str()).map_err(|()| anyhow::anyhow!("Invalid topic"))?;
    println!("Serving files from {path:?} on topic '{topic:?}'",);

    let _acceptor = ditto
        .datastreams()
        .bind_topic(topic)
        .finish_with(move |candidate| {
            let Some(args) = candidate.arguments() else {
                eprintln!("No arguments provided by client, ignoring the request.");
                return;
            };

            let Ok(filename) = std::str::from_utf8(args.as_slice()) else {
                eprintln!("Invalid filename provided by client, ignoring the request.");
                return;
            };
            let Ok(file_path) = path.join(&filename).canonicalize() else {
                eprintln!(
                    "Invalid file path '{filename}' provided by client, ignoring the request."
                );
                return;
            };
            if !file_path.starts_with(&path) {
                eprintln!(
                    "Attempt to access file '{file_path:?}' outside of served directory \
                     '{path:?}', ignoring the request."
                );
                return;
            }

            let Ok(mut file) = std::fs::File::open(&file_path) else {
                eprintln!("File '{filename}' not found, ignoring the request.");
                return;
            };

            let size = file.metadata().map(|m| m.len()).unwrap_or(0);

            let stream = candidate
                .with_arguments(Payload::copied_from_slice(&size.to_le_bytes()))
                .open_write_only();

            let mut buf = [0u8; 1024];
            let mut sends = 0;
            while let Ok(n @ 1..) = file.read(&mut buf) {
                _ = stream.message(Payload::copied_from_slice(&buf[..n])).send();
                sends += 1;
            }
            println!("Sent file '{filename}' in {sends} chunks.");
        });

    std::future::pending::<()>().await;
    Ok(())
}

async fn fetch(ditto: Arc<Ditto>, args: Fetch) -> Result<()> {
    let Fetch {
        file,
        topic,
        peer,
        save_as,
    } = args;

    let topic = Topic::try_from(topic.as_str()).map_err(|()| anyhow::anyhow!("Invalid topic"))?;

    let remote_peer = match peer {
        Some(peer) => peer,
        None => loop {
            let peers = ditto.presence().graph().remote_peers;

            if let Some(peer) = peers.iter().next() {
                break PeerPubkey::from_str(&peer.peer_key_string).unwrap();
            }

            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        },
    };

    let candidate = ditto
        .datastreams()
        .connect(remote_peer.clone(), topic)
        .with_arguments(Payload::copied_from_slice(file.as_bytes()))
        .finish_async()
        .await?;

    let Some(expected_size) = candidate
        .arguments()
        .and_then(|args| args.as_slice().try_into().ok().map(u64::from_le_bytes))
    else {
        anyhow::bail!("Server did not provide expected file size");
    };

    let mut stream = candidate.open(tokio::sync::mpsc::unbounded_channel());

    let mut contents = Vec::with_capacity(expected_size as usize);
    while let Some(chunk) = stream.recv().await {
        contents.extend_from_slice(chunk.payload());
        eprintln!("Received {}/{} bytes", contents.len(), expected_size,);
    }

    if let Some(save_path) = save_as {
        std::fs::write(&save_path, &contents)?;
        println!("File saved to {:?}", save_path);
    } else {
        println!("{}", String::from_utf8_lossy(&contents));
    }

    Ok(())
}
