sproto 0.1.0

use std::path::PathBuf;
use std::pin::Pin;
use std::task::{Context, Poll};

use anyhow::{Context as _, Result, bail};
use indicatif::ProgressBar;
use tokio::io::AsyncWrite;

use sproto::TlsMode;
use sproto::client::{self, Credentials};

/// Common connection arguments shared across all examples.
/// Implement this on your clap `Args` struct via `#[command(flatten)]`.
#[derive(clap::Args)]
pub struct ConnectionArgs {
    /// NAS hostname or IP (required unless --from-sqlite)
    #[arg(long)]
    pub host: Option<String>,

    /// Port
    #[arg(long, default_value = "6690")]
    pub port: u16,

    /// Username for authentication
    #[arg(long)]
    pub username: Option<String>,

    /// Password for authentication
    #[arg(long)]
    pub password: Option<String>,

    /// Use existing session token directly
    #[arg(long)]
    pub session: Option<String>,

    /// `restore_id` (required with `--session`)
    #[arg(long)]
    pub restore_id: Option<String>,

    /// Read credentials from Synology Drive Client's sys.sqlite (macOS only)
    #[cfg(target_os = "macos")]
    #[arg(long)]
    pub from_sqlite: bool,

    /// Accept self-signed TLS certificates
    #[arg(long)]
    pub allow_untrusted_cert: bool,

    /// Share name to use
    #[arg(long, default_value = "home")]
    pub view: String,
}

impl ConnectionArgs {
    /// Build a connected `SprotoClient` from these args.
    pub async fn connect(&self) -> Result<sproto::Client> {
        let (host, port, credentials, tls) = self.resolve()?;

        sproto::Client::connect(
            client::Config::builder()
                .host(host)
                .port(port)
                .tls(tls)
                .credentials(credentials)
                .build(),
        )
        .await
        .context("failed to connect")
    }

    /// Resolve the `view_id` for the configured share name.
    pub async fn resolve_view(&self, client: &sproto::Client) -> Result<u64> {
        let shares = client
            .list_shares()
            .await
            .context("failed to list shares")?;

        if let Some(view) = shares.iter().find(|s| s.name == self.view) {
            Ok(view.view_id)
        } else if shares.is_empty() {
            Ok(1)
        } else {
            let names: Vec<_> = shares.iter().map(|s| s.name.as_str()).collect();
            bail!(
                "view '{}' not found. Available: {}",
                self.view,
                names.join(", ")
            );
        }
    }

    fn resolve(&self) -> Result<(String, u16, Credentials, TlsMode)> {
        #[cfg(target_os = "macos")]
        if self.from_sqlite {
            return read_from_sqlite();
        }

        let host = self
            .host
            .clone()
            .context("--host is required (or use --from-sqlite)")?;

        let tls = if self.allow_untrusted_cert {
            TlsMode::Insecure
        } else {
            TlsMode::Verified
        };

        let credentials = if let Some(session) = &self.session {
            Credentials::Session {
                session: session.clone(),
                restore_id: self
                    .restore_id
                    .clone()
                    .context("--restore-id is required with --session")?,
            }
        } else if let (Some(username), Some(password)) = (&self.username, &self.password) {
            Credentials::Password {
                username: username.clone(),
                password: password.clone(),
                otp: None,
            }
        } else {
            bail!("provide --username/--password, --session/--restore-id, or --from-sqlite");
        };

        Ok((host, self.port, credentials, tls))
    }
}

#[cfg(target_os = "macos")]
fn read_from_sqlite() -> Result<(String, u16, Credentials, TlsMode)> {
    let home = std::env::var("HOME").context("HOME not set")?;
    let db_path = format!("{home}/Library/Application Support/SynologyDrive/data/db/sys.sqlite");

    let conn =
        rusqlite::Connection::open_with_flags(&db_path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)
            .with_context(|| format!("failed to open {db_path}"))?;

    let (session, restore_id, hostname, port, ssl_allow_untrust): (
        String,
        String,
        String,
        u16,
        i32,
    ) = conn
        .query_row(
            "SELECT session, restore_id, COALESCE(NULLIF(server_ip,''), host_name), server_port, ssl_allow_untrust FROM connection_table LIMIT 1",
            [],
            |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?, row.get(4)?)),
        )
        .context("failed to read connection from sys.sqlite")?;

    let tls = if ssl_allow_untrust != 0 {
        TlsMode::Insecure
    } else {
        TlsMode::Verified
    };

    Ok((
        hostname,
        port,
        Credentials::Session {
            session,
            restore_id,
        },
        tls,
    ))
}

/// Initialize tracing with env-filter (enable with `RUST_LOG=debug`).
pub fn init_tracing() {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .with_target(false)
        .init();
}

/// An `AsyncWrite` wrapper that updates progress bars on each write
/// and defers file creation until the first write arrives.
#[allow(dead_code, reason = "not all examples use ProgressWriter")]
pub struct ProgressWriter {
    path: PathBuf,
    file: Option<tokio::fs::File>,
    bars: Vec<ProgressBar>,
}

#[allow(dead_code, reason = "not all examples use ProgressWriter")]
impl ProgressWriter {
    pub fn new(path: PathBuf, bar: ProgressBar) -> Self {
        Self {
            path,
            file: None,
            bars: vec![bar],
        }
    }

    pub fn with_overall(mut self, bar: ProgressBar) -> Self {
        self.bars.push(bar);
        self
    }

    fn ensure_file(&mut self) -> std::io::Result<&mut tokio::fs::File> {
        if self.file.is_none() {
            if let Some(parent) = self.path.parent() {
                std::fs::create_dir_all(parent)?;
            }

            let std_file = std::fs::File::create(&self.path)?;
            self.file = Some(tokio::fs::File::from_std(std_file));
        }

        Ok(self.file.as_mut().unwrap())
    }
}

impl AsyncWrite for ProgressWriter {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        let file = self.ensure_file()?;
        let poll = Pin::new(file).poll_write(cx, buf);
        if let Poll::Ready(Ok(n)) = &poll {
            for bar in &self.bars {
                bar.inc(*n as u64);
            }
        }
        poll
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        let file = self.ensure_file()?;
        Pin::new(file).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        self.file
            .as_mut()
            .map_or(Poll::Ready(Ok(())), |f| Pin::new(f).poll_shutdown(cx))
    }
}