mdrs-client-rust/src/commands/shared.rs

use crate::models::file::File;
use crate::models::folder::{FolderDetail, FolderSimple};
use crate::connection::MDRSConnection;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs;
use std::sync::{Arc, LazyLock, Mutex};
use unicode_normalization::UnicodeNormalization;

// ---------------------------------------------------------------------------
// Cache structs — matching Python's cache format exactly
// ---------------------------------------------------------------------------

#[derive(Deserialize, Clone)]
pub struct CacheToken {
    pub access: String,
    pub refresh: String,
}

/// Minimal user fields stored in cache — matches Python's `User` dataclass.
#[derive(Deserialize, Clone)]
pub struct CacheUser {
    pub id: u32,
    pub username: String,
    pub laboratory_ids: Vec<u32>,
    pub is_reviewer: bool,
}

/// All four fields of a laboratory, needed for digest computation.
#[derive(Deserialize, Clone)]
pub struct CacheLaboratory {
    pub id: u32,
    pub name: String,
    #[serde(default)]
    pub pi_name: String,
    #[serde(default)]
    pub full_name: String,
}

/// Wrapper matching Python's `Laboratories` serialization: `{"items": [...]}`.
#[derive(Deserialize, Clone, Default)]
pub struct CacheLabsWrapper {
    pub items: Vec<CacheLaboratory>,
}

#[derive(Deserialize, Clone)]
pub struct Cache {
    pub user: Option<CacheUser>,
    pub token: CacheToken,
    pub laboratories: CacheLabsWrapper,
}

// ---------------------------------------------------------------------------
// Digest computation — must produce exactly the same hash as Python's
// `CacheData.__calc_digest()`:
//   hashlib.sha256(
//       json.dumps([user_asdict, token_asdict, labs_asdict]).encode("utf-8")
//   ).hexdigest()
//
// Python's default json.dumps uses separators=(', ', ': ') and
// ensure_ascii=True.  Field order follows dataclass definition order.
// ---------------------------------------------------------------------------

/// Escape a string in Python json.dumps style:
/// - Special chars: ", \, and control chars -> standard JSON escapes
/// - Non-ASCII chars -> \uXXXX  (matches Python ensure_ascii=True default)
fn python_json_string(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('"');
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => {
                out.push_str(&format!("\\u{:04x}", c as u32));
            }
            c if c.is_ascii() => out.push(c),
            c => {
                // Non-ASCII: encode as \uXXXX (BMP) or surrogate pair (outside BMP)
                let code = c as u32;
                if code <= 0xFFFF {
                    out.push_str(&format!("\\u{:04x}", code));
                } else {
                    let code = code - 0x10000;
                    let high = 0xD800 + (code >> 10);
                    let low = 0xDC00 + (code & 0x3FF);
                    out.push_str(&format!("\\u{:04x}\\u{:04x}", high, low));
                }
            }
        }
    }
    out.push('"');
    out
}

/// Serialize a list of u32 as a Python-style JSON array: `[1, 2, 3]`
fn python_json_u32_array(items: &[u32]) -> String {
    if items.is_empty() {
        return "[]".to_string();
    }
    let inner: Vec<String> = items.iter().map(|x| x.to_string()).collect();
    format!("[{}]", inner.join(", "))
}

/// Build the JSON array string that Python's `__calc_digest` hashes:
///   [user_asdict_or_null, token_asdict, labs_asdict]
///
/// Field order matches each Python dataclass definition:
///   User:         id, username, laboratory_ids, is_reviewer
///   Token:        access, refresh
///   Laboratories: items
///   Laboratory:   id, name, pi_name, full_name
pub fn python_digest_json(
    user: Option<&CacheUser>,
    access: &str,
    refresh: &str,
    labs: &CacheLabsWrapper,
) -> String {
    let user_str = match user {
        None => "null".to_string(),
        Some(u) => format!(
            "{{\"id\": {}, \"username\": {}, \"laboratory_ids\": {}, \"is_reviewer\": {}}}",
            u.id,
            python_json_string(&u.username),
            python_json_u32_array(&u.laboratory_ids),
            if u.is_reviewer { "true" } else { "false" }
        ),
    };

    let token_str = format!(
        "{{\"access\": {}, \"refresh\": {}}}",
        python_json_string(access),
        python_json_string(refresh)
    );

    let items: Vec<String> = labs
        .items
        .iter()
        .map(|lab| {
            format!(
                "{{\"id\": {}, \"name\": {}, \"pi_name\": {}, \"full_name\": {}}}",
                lab.id,
                python_json_string(&lab.name),
                python_json_string(&lab.pi_name),
                python_json_string(&lab.full_name)
            )
        })
        .collect();

    let items_str = if items.is_empty() {
        "[]".to_string()
    } else {
        format!("[{}]", items.join(", "))
    };
    let labs_str = format!("{{\"items\": {}}}", items_str);

    format!("[{}, {}, {}]", user_str, token_str, labs_str)
}

/// Compute the cache digest compatible with Python's `CacheData.__calc_digest`.
pub fn compute_digest(
    user: Option<&CacheUser>,
    access: &str,
    refresh: &str,
    labs: &CacheLabsWrapper,
) -> String {
    let json_str = python_digest_json(user, access, refresh, labs);
    let mut hasher = Sha256::new();
    hasher.update(json_str.as_bytes());
    format!("{:x}", hasher.finalize())
}

// ---------------------------------------------------------------------------
// Per-remote async mutex map (in-process serialization)
// ---------------------------------------------------------------------------

static REMOTE_LOCKS: LazyLock<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

fn get_remote_lock(remote: &str) -> Arc<tokio::sync::Mutex<()>> {
    let mut map = REMOTE_LOCKS.lock().unwrap();
    map.entry(remote.to_string())
        .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
        .clone()
}

// ---------------------------------------------------------------------------
// Cache file path helper
// ---------------------------------------------------------------------------

fn cache_file_path(remote: &str) -> std::path::PathBuf {
    crate::settings::SETTINGS
        .config_dirname
        .join("cache")
        .join(format!("{}.json", remote))
}

// ---------------------------------------------------------------------------
// Load cache (low-level, no token refresh)
// ---------------------------------------------------------------------------

/// Load token and laboratories from the login cache file (no token refresh check).
pub fn load_cache(remote: &str) -> Result<Cache, Box<dyn std::error::Error>> {
    let cache_path = cache_file_path(remote);
    if !cache_path.exists() {
        return Err(format!(
            "Not logged in to `{}`. Run `mdrs login {}` first.",
            remote, remote
        )
        .into());
    }
    let data = fs::read_to_string(&cache_path)?;
    serde_json::from_str::<Cache>(&data).map_err(|e| {
        format!(
            "Cache for `{}` is invalid or outdated ({}). Run `mdrs login {}` to refresh it.",
            remote, e, remote
        )
        .into()
    })
}

// ---------------------------------------------------------------------------
// Token-aware cache load with refresh and locking
// ---------------------------------------------------------------------------

/// Load cache, check token expiry, and refresh the access token if needed.
///
/// Locking strategy:
/// - Per-remote `tokio::sync::Mutex` serializes concurrent async tasks within
///   the same process.
/// - `flock(LOCK_EX)` on the cache file serializes across separate processes
///   on the same host.
pub async fn load_cache_with_token_refresh(
    remote: &str,
) -> Result<Cache, Box<dyn std::error::Error>> {
    // Acquire the in-process async mutex for this remote
    let lock = get_remote_lock(remote);
    let _guard = lock.lock().await;

    // Re-read the cache inside the lock (another task may have already refreshed)
    let mut cache = load_cache(remote)?;

    if crate::token::is_expired(&cache.token.refresh) {
        return Err(format!(
            "Session for `{}` has expired. Please run `mdrs login {}` again.",
            remote, remote
        )
        .into());
    }

    if crate::token::is_refresh_required(&cache.token.access, &cache.token.refresh) {
        let new_access = refresh_and_persist(remote, &cache).await?;
        cache.token.access = new_access;
    }

    Ok(cache)
}

/// Call the token-refresh endpoint and write the new access token back to the
/// cache file.  The caller must already hold the per-remote async mutex.
/// Also recomputes the digest so Python can verify the cache.
async fn refresh_and_persist(
    remote: &str,
    cache: &Cache,
) -> Result<String, Box<dyn std::error::Error>> {
    // Build a connection without Bearer token just to reach the refresh endpoint
    let url = crate::commands::config::get_remote_url(remote)?
        .ok_or_else(|| format!("Remote `{}` is not configured.", remote))?;
    let conn = MDRSConnection::new(&url);

    let new_access = conn.token_refresh(&cache.token.refresh).await?;

    // Recompute the digest with the new access token so the Python client
    // can still verify the cache after a token refresh.
    let new_digest = compute_digest(
        cache.user.as_ref(),
        &new_access,
        &cache.token.refresh,
        &cache.laboratories,
    );

    // Persist the updated access token to the cache file with an exclusive file
    // lock so that other processes do not read a partially written file.
    let cache_path = cache_file_path(remote);

    let raw = fs::read_to_string(&cache_path)?;
    let mut obj: serde_json::Value = serde_json::from_str(&raw)?;

    obj["token"]["access"] = serde_json::Value::String(new_access.clone());
    obj["digest"] = serde_json::Value::String(new_digest);

    // Write atomically: write to .tmp then rename, while holding an exclusive
    // flock on the .tmp file for cross-process safety.
    let tmp_path = cache_path.with_extension("tmp");
    {
        use fs2::FileExt;
        use std::io::Write;
        let mut tmp_file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&tmp_path)?;
        tmp_file.lock_exclusive()?;
        tmp_file.write_all(serde_json::to_string(&obj)?.as_bytes())?;
        tmp_file.flush()?;
        tmp_file.unlock()?;
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let mut perms = fs::metadata(&tmp_path)?.permissions();
        perms.set_mode(0o600);
        fs::set_permissions(&tmp_path, perms)?;
    }
    fs::rename(&tmp_path, &cache_path)?;

    Ok(new_access)
}

// ---------------------------------------------------------------------------
// Connection helpers
// ---------------------------------------------------------------------------

/// Create an authenticated MDRSConnection for the given remote label
pub fn create_authenticated_conn(
    remote: &str,
    cache: &Cache,
) -> Result<MDRSConnection, Box<dyn std::error::Error>> {
    let url = crate::commands::config::get_remote_url(remote)?
        .ok_or_else(|| format!("Remote `{}` is not configured.", remote))?;
    let mut conn = MDRSConnection::new(&url);
    conn.token = Some(cache.token.access.clone());
    Ok(conn)
}

// ---------------------------------------------------------------------------
// Path and lab helpers
// ---------------------------------------------------------------------------

/// Parse "remote:/labname/path/" into (remote, labname, folder_path)
pub fn parse_remote_path(
    remote_path: &str,
) -> Result<(String, String, String), Box<dyn std::error::Error>> {
    let parts: Vec<&str> = remote_path.splitn(2, ':').collect();
    if parts.len() != 2 {
        return Err("remote_path must be in the form 'remote:/labname/path/'".into());
    }
    let remote = parts[0].to_string();
    let rest = parts[1];
    if !rest.starts_with('/') {
        return Err("Path must be absolute (start with '/')".into());
    }
    let folder_parts: Vec<&str> = rest.trim_start_matches('/').splitn(2, '/').collect();
    let labname = folder_parts[0].to_string();
    let path = if folder_parts.len() > 1 && !folder_parts[1].is_empty() {
        format!("/{}", folder_parts[1].trim_end_matches('/'))
    } else {
        "/".to_string()
    };
    Ok((remote, labname, path))
}

/// Look up a laboratory by name in the cache
pub fn find_lab_in_cache<'a>(
    cache: &'a Cache,
    labname: &str,
) -> Result<&'a CacheLaboratory, Box<dyn std::error::Error>> {
    cache
        .laboratories
        .items
        .iter()
        .find(|l| l.name == labname)
        .ok_or_else(|| format!("Laboratory `{}` not found.", labname).into())
}

/// Apply Unicode NFC normalization to a string.
pub fn nfc(s: &str) -> String {
    s.chars().nfc().collect()
}

/// Resolve a folder by path using the API (GET v3/folders/?path=...&laboratory_id=...)
pub async fn find_folder(
    conn: &MDRSConnection,
    lab_id: u32,
    path: &str,
    password: Option<&str>,
) -> Result<FolderDetail, Box<dyn std::error::Error>> {
    let normalized_path = nfc(path);
    let folders = conn.list_folders_by_path(lab_id, &normalized_path).await?;
    if folders.is_empty() {
        return Err(format!("Folder `{}` not found.", path).into());
    }
    if folders.len() != 1 {
        return Err(
            format!("Ambiguous path `{}`: {} folders matched.", path, folders.len()).into(),
        );
    }
    let folder_simple = &folders[0];
    if folder_simple.lock {
        match password {
            None => {
                return Err(format!(
                    "Folder `{}` is locked. Use -p/--password to provide a password.",
                    path
                )
                .into())
            }
            Some(pw) => conn.folder_auth(&folder_simple.id, pw).await?,
        }
    }
    let folder = conn.retrieve_folder(&folder_simple.id).await?;
    Ok(folder)
}

/// Find a file by name (NFC-normalized, case-insensitive) in a file list.
pub fn find_file_by_name<'a>(files: &'a [File], name: &str) -> Option<&'a File> {
    let name_lower = nfc(name).to_lowercase();
    files.iter().find(|f| nfc(&f.name).to_lowercase() == name_lower)
}

/// Find a sub-folder by name (NFC-normalized, case-insensitive).
pub fn find_subfolder_by_name<'a>(subfolders: &'a [FolderSimple], name: &str) -> Option<&'a FolderSimple> {
    let name_lower = nfc(name).to_lowercase();
    subfolders.iter().find(|f| nfc(&f.name).to_lowercase() == name_lower)
}

/// Format an ISO 8601 timestamp as "YYYY/MM/DD HH:MM:SS"
pub fn fmt_datetime(iso: &str) -> String {
    let s = iso.trim();
    let s = if let Some(pos) = s[10..].find(|c: char| c == '+' || c == '-') {
        &s[..10 + pos]
    } else {
        s.trim_end_matches('Z')
    };
    if s.len() >= 19 {
        let date = s[..10].replace('-', "/");
        let time = &s[11..19];
        format!("{} {}", date, time)
    } else {
        iso.to_string()
    }
}