//! Shared privileged executor (port of `sessions/privilegedExec.js`).
//!
//! Bounded privileged operations over an existing SSH session: callers pass a fixed script
//! plus validated arguments — there is no API for arbitrary commands. The sudo password is
//! written to stdin only after sudo prints our prompt marker; it is never stored on disk,
//! logged or placed in argv. Payload frames are sent only after the script's READY marker.

use crate::core::error::{CoreError, CoreResult};
use crate::core::ssh::exec::ExecEvent;
use crate::core::ssh::session::{shell_quote, SessionManager, SessionRecord};
use base64::Engine;
use once_cell::sync::Lazy;
use parking_lot::Mutex;
use regex::Regex;
use serde_json::{json, Value};
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, Notify};
use zeroize::Zeroizing;

pub const SUDO_PROMPT: &str = "[chj-sudo-auth]";
pub const READY_MARKER: &str = "\x1eCHJ_READY\n";
static ERROR_MARKER: Lazy<Regex> = Lazy::new(|| Regex::new(r"\x1eCHJ_ERR=([A-Z0-9_]{2,64})\r?\n").unwrap());
static FRAME_TAG: Lazy<Regex> = Lazy::new(|| Regex::new(r"^[A-Z]{2,16}$").unwrap());
const MAX_PASSWORD_LENGTH: usize = 1024;
pub const SUDO_CACHE_IDLE: Duration = Duration::from_secs(10 * 60);

pub const PROBE_SCRIPT: &str = r#"LC_ALL=C; export LC_ALL
PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin; export PATH
printf 'user=%s\nuid=%s\ngroups=%s\nclient=%s\n' "$(id -un)" "$(id -u)" "$(id -Gn 2>/dev/null | tr ' ' ',')" "${SSH_CLIENT%% *}"
if command -v sudo >/dev/null 2>&1; then
  out=$(sudo -n -p '' true 2>&1 </dev/null); code=$?
  printf 'sudoInstalled=1\nsudoNonInteractive=%s\nsudoMessage=%s\n' "$code" "$(printf '%s' "$out" | head -n 1 | tr -d '\r\n' | cut -c1-300)"
else
  printf 'sudoInstalled=0\n'
fi"#;

pub fn sudo_message(code: &str) -> &'static str {
    match code {
        "SUDO_PASSWORD_REQUIRED" => "This operation requires the sudo password of the SSH account.",
        "SUDO_PASSWORD_INCORRECT" => "The sudo password was rejected.",
        "SUDO_NOT_PERMITTED" => "The SSH account is not permitted to run this operation with sudo.",
        "SUDO_TTY_REQUIRED" => "The sudo policy requires a terminal (requiretty); CH-J does not bypass it.",
        "SUDO_AUTH_FAILED" => "sudo authentication failed or the account is locked.",
        "SUDO_UNAVAILABLE" => "sudo is not installed and the SSH account is not root.",
        _ => "sudo did not start the operation.",
    }
}

pub fn normalize_password(value: Option<&str>) -> CoreResult<String> {
    let password = value.unwrap_or("");
    if password.chars().count() > MAX_PASSWORD_LENGTH || password.contains(['\0', '\r', '\n']) {
        return Err(CoreError::new("SUDO_PASSWORD_INVALID", "Invalid sudo password."));
    }
    Ok(password.to_string())
}

pub fn clean_remote_message(text: &str, limit: usize) -> String {
    let without_prompt = text.split(SUDO_PROMPT).collect::<String>();
    let cleaned: String = without_prompt.chars().map(|c| if (c as u32) < 0x20 && c != '\n' && c != '\t' || c as u32 == 0x7f { ' ' } else { c }).collect();
    let lines: Vec<&str> = cleaned.split('\n').map(str::trim).filter(|l| !l.is_empty()).collect();
    let tail = &lines[lines.len().saturating_sub(6)..];
    tail.join(" | ").chars().take(limit).collect()
}

pub fn classify_sudo_failure(stderr: &str, code: Option<i64>, auth_failure: Option<&str>) -> String {
    if let Some(failure) = auth_failure {
        return failure.to_string();
    }
    let test = |pattern: &str| Regex::new(pattern).map(|r| r.is_match(stderr)).unwrap_or(false);
    if test(r"(?i)incorrect password attempt|Sorry, try again") {
        "SUDO_PASSWORD_INCORRECT"
    } else if test(r"(?i)is not in the sudoers file|is not allowed to run sudo|may not run sudo|not allowed to execute|is not allowed to run") {
        "SUDO_NOT_PERMITTED"
    } else if test(r"(?i)must have a tty|a terminal is required|no tty present") {
        "SUDO_TTY_REQUIRED"
    } else if test(r"(?i)a password is required|no password was provided") {
        "SUDO_PASSWORD_REQUIRED"
    } else if test(r"(?i)account.*(locked|expired)|authentication (failure|token)|maximum number of tries") {
        "SUDO_AUTH_FAILED"
    } else if code == Some(127) || test(r"(?i)sudo: (command )?not found") {
        "SUDO_UNAVAILABLE"
    } else {
        "SUDO_FAILED"
    }
    .to_string()
}

fn parse_key_values(text: &str) -> std::collections::HashMap<String, String> {
    text.lines().filter_map(|l| l.split_once('=').filter(|(k, _)| !k.is_empty()).map(|(k, v)| (k.to_string(), v.to_string()))).collect()
}

fn public_privilege(p: &Value) -> Value {
    json!({
        "mode": p["mode"], "verified": p["verified"] == true, "username": p["username"],
        "sudoInstalled": p["sudoInstalled"] == true, "groups": p["groups"], "clientAddress": p["clientAddress"],
        "message": p["message"], "checkedAt": p["checkedAt"]
    })
}

pub struct Frame {
    pub tag: String,
    pub data: Vec<u8>,
}

impl Frame {
    pub fn text(tag: &str, data: &str) -> Self {
        Self { tag: tag.into(), data: data.as_bytes().to_vec() }
    }
}

pub struct RunOptions {
    pub script: String,
    pub args: Vec<String>,
    pub frames: Vec<Frame>,
    pub privileged: bool,
    pub sudo_password: Option<String>,
    pub remember_password: bool,
    pub use_cached_password: bool,
    pub purpose: String,
    pub error_prefix: String,
    pub timeout: Duration,
    pub output_limit: usize,
}

impl Default for RunOptions {
    fn default() -> Self {
        Self {
            script: String::new(),
            args: Vec::new(),
            frames: Vec::new(),
            privileged: true,
            sudo_password: None,
            remember_password: false,
            use_cached_password: false,
            purpose: String::new(),
            error_prefix: "REMOTE_".into(),
            timeout: Duration::from_secs(60),
            output_limit: 4 * 1024 * 1024,
        }
    }
}

#[derive(Debug, Clone)]
pub struct RunResult {
    pub code: Option<i64>,
    pub stdout: String,
    pub stderr: String,
    pub privileged: bool,
}

// ---------------------------------------------------------------------------------------------
// Interaction: wait for markers in stdout, answer on stdin.
// ---------------------------------------------------------------------------------------------

struct InteractionShared {
    stdout: Mutex<String>,
    consumed: Mutex<usize>,
    done: Mutex<bool>,
    notify: Notify,
}

enum InteractionCommand {
    Send(String),
    End,
}

#[derive(Clone)]
pub struct Interaction {
    shared: Arc<InteractionShared>,
    commands: mpsc::UnboundedSender<InteractionCommand>,
}

pub struct Matched {
    pub groups: Vec<Option<String>>,
    pub before: String,
}

impl Interaction {
    pub async fn wait_for(&self, pattern: &Regex, timeout: Duration) -> CoreResult<Matched> {
        let deadline = tokio::time::Instant::now() + timeout;
        loop {
            let notified = self.shared.notify.notified();
            {
                let stdout = self.shared.stdout.lock();
                let mut consumed = self.shared.consumed.lock();
                let text = &stdout[(*consumed).min(stdout.len())..];
                if let Some(caps) = pattern.captures(text) {
                    let whole = caps.get(0).expect("match");
                    let matched = Matched { groups: caps.iter().map(|g| g.map(|m| m.as_str().to_string())).collect(), before: text[..whole.start()].to_string() };
                    *consumed += whole.end();
                    return Ok(matched);
                }
                if *self.shared.done.lock() {
                    return Err(CoreError::new("PROTOCOL_ERROR", "The remote operation ended unexpectedly."));
                }
            }
            if tokio::time::timeout_at(deadline, notified).await.is_err() {
                return Err(CoreError::new("REMOTE_COMMAND_TIMEOUT", "The remote operation did not respond in time."));
            }
        }
    }

    pub fn send(&self, line: &str) {
        let _ = self.commands.send(InteractionCommand::Send(format!("{line}\n")));
    }

    pub fn end(&self) {
        let _ = self.commands.send(InteractionCommand::End);
    }
}

pub struct PrivilegedExecutor {
    sessions: Arc<SessionManager>,
}

impl PrivilegedExecutor {
    pub fn new(sessions: Arc<SessionManager>) -> Self {
        Self { sessions }
    }

    /// root / sudo-nopasswd / sudo-password (unverified) / none; cached per live session.
    pub async fn probe(&self, session_id: &str, refresh: bool) -> CoreResult<Value> {
        let record = self.sessions.require_connected(session_id)?;
        if !refresh {
            if let Some(p) = record.privilege.lock().clone() {
                return Ok(public_privilege(&p));
            }
        }
        let channel = self.sessions.exec_channel(&record, &format!("sh -c {}", shell_quote(PROBE_SCRIPT)), Duration::from_secs(15), 512 * 1024).await?;
        let mut channel = channel;
        channel.eof().await;
        let result = channel.finish().await?;
        let values = parse_key_values(&result.stdout);
        let get = |k: &str| values.get(k).cloned().unwrap_or_default();
        let sudo_message = get("sudoMessage");
        let mut privilege = json!({
            "username": if get("user").is_empty() { record.username.clone() } else { get("user") },
            "sudoInstalled": get("sudoInstalled") == "1",
            "groups": get("groups").split(',').filter(|s| !s.is_empty()).take(256).collect::<Vec<_>>(),
            "clientAddress": get("client").trim().chars().take(64).collect::<String>(),
            "message": clean_remote_message(&sudo_message, 300),
            "checkedAt": chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
            "mode": "unknown",
            "verified": false
        });
        let set = |p: &mut Value, mode: &str, verified: bool| {
            p["mode"] = json!(mode);
            p["verified"] = json!(verified);
        };
        if get("uid") == "0" {
            set(&mut privilege, "root", true);
        } else if get("sudoInstalled") != "1" {
            set(&mut privilege, "none", true);
        } else if get("sudoNonInteractive") == "0" {
            set(&mut privilege, "sudo-nopasswd", true);
        } else if Regex::new(r"(?i)password is required").unwrap().is_match(&sudo_message) {
            privilege["mode"] = json!("sudo-password");
        } else if Regex::new(r"(?i)not allowed|may not run sudo|not in the sudoers").unwrap().is_match(&sudo_message) {
            set(&mut privilege, "none", true);
        }
        *record.privilege.lock() = Some(privilege.clone());
        Ok(public_privilege(&privilege))
    }

    pub fn cached_password(&self, record: &SessionRecord) -> Option<Zeroizing<String>> {
        let mut cache = record.sudo_cache.lock();
        match cache.as_mut() {
            Some((password, expires)) if *expires > Instant::now() && record.is_connected() => {
                *expires = Instant::now() + SUDO_CACHE_IDLE;
                Some(password.clone())
            }
            _ => {
                *cache = None;
                None
            }
        }
    }

    pub fn has_cached_password(&self, session_id: &str) -> bool {
        self.sessions.get(session_id).map(|r| self.cached_password(&r).is_some()).unwrap_or(false)
    }

    pub fn forget_password(&self, session_id: &str) {
        if let Some(record) = self.sessions.get(session_id) {
            *record.sudo_cache.lock() = None;
        }
    }

    fn remember_failure(&self, record: &SessionRecord, code: &str) {
        let mut guard = record.privilege.lock();
        let Some(p) = guard.as_mut() else { return };
        match code {
            "SUDO_NOT_PERMITTED" => {
                p["mode"] = json!("none");
                p["verified"] = json!(true);
            }
            "SUDO_UNAVAILABLE" => {
                p["mode"] = json!("none");
                p["verified"] = json!(true);
                p["sudoInstalled"] = json!(false);
            }
            "SUDO_PASSWORD_REQUIRED" | "SUDO_PASSWORD_INCORRECT" if p["mode"] != "sudo-password" => p["mode"] = json!("sudo-password"),
            _ => {}
        }
    }

    pub async fn run(&self, session_id: &str, options: RunOptions) -> CoreResult<RunResult> {
        self.run_interactive(session_id, options, None::<fn(Interaction) -> std::future::Ready<CoreResult<Value>>>).await.map(|(result, _)| result)
    }

    /// Runs a fixed script; `interact` (if any) drives a stdin line protocol after READY.
    pub async fn run_interactive<F, Fut>(&self, session_id: &str, options: RunOptions, interact: Option<F>) -> CoreResult<(RunResult, Option<Value>)>
    where
        F: FnOnce(Interaction) -> Fut,
        Fut: Future<Output = CoreResult<Value>> + Send,
    {
        let record = self.sessions.require_connected(session_id)?;
        let mut password = Zeroizing::new(normalize_password(options.sudo_password.as_deref())?);
        let mut used_cache = false;
        if password.is_empty() && options.use_cached_password {
            if let Some(cached) = self.cached_password(&record) {
                password = cached;
                used_cache = true;
            }
        }
        let privilege = if options.privileged { Some(self.probe(session_id, false).await?) } else { None };
        if let Some(p) = &privilege {
            if p["mode"] == "none" && p["verified"] == true {
                let code = if p["sudoInstalled"] == true { "SUDO_NOT_PERMITTED" } else { "SUDO_UNAVAILABLE" };
                return Err(CoreError::new(code, sudo_message(code)));
            }
        }
        let use_sudo = privilege.as_ref().map(|p| p["mode"] != "root").unwrap_or(false);
        let body = format!("sh -c {} chj {}", shell_quote(&options.script), options.args.iter().map(|a| shell_quote(a)).collect::<Vec<_>>().join(" "));
        let command = if use_sudo { format!("LC_ALL=C sudo -S -p {} -- {body}", shell_quote(SUDO_PROMPT)) } else { format!("LC_ALL=C {body}") };
        let mut frames = Vec::new();
        for frame in &options.frames {
            if !FRAME_TAG.is_match(&frame.tag) {
                return Err(CoreError::new("PROTOCOL_ERROR", "Invalid frame tag."));
            }
            frames.push(format!("CHJ{}:{}\n", frame.tag, base64::engine::general_purpose::STANDARD.encode(&frame.data)));
        }
        let mut channel = self.sessions.exec_channel(&record, &command, options.timeout, options.output_limit).await?;
        let shared = Arc::new(InteractionShared { stdout: Mutex::new(String::new()), consumed: Mutex::new(0), done: Mutex::new(false), notify: Notify::new() });
        let (command_tx, mut command_rx) = mpsc::unbounded_channel();
        let mut answered = 0usize;
        let mut auth_failure: Option<&'static str> = None;
        let mut started = false;
        let mut interact = interact;
        let mut pending: Option<std::pin::Pin<Box<dyn Future<Output = CoreResult<Value>> + Send>>> = None;
        let mut interaction_result: Option<CoreResult<Value>> = None;
        let mut channel_done = false;
        // Drive the channel and the optional interaction concurrently on this task.
        while !channel_done || pending.is_some() {
            tokio::select! {
                event = channel.next(), if !channel_done => {
                    let event = event?;
                    if use_sudo && !started {
                        let prompts = channel.stderr.matches(SUDO_PROMPT).count();
                        if prompts > answered {
                            if answered == 0 && !password.is_empty() {
                                answered = 1;
                                channel.write(format!("{}\n", password.as_str()).as_bytes()).await;
                            } else {
                                answered = prompts;
                                auth_failure = Some(if password.is_empty() { "SUDO_PASSWORD_REQUIRED" } else { "SUDO_PASSWORD_INCORRECT" });
                                channel.eof().await;
                                channel.close().await;
                                break;
                            }
                        }
                    }
                    if !started && channel.stdout.contains(READY_MARKER) {
                        started = true;
                        for frame in &frames {
                            channel.write(frame.as_bytes()).await;
                        }
                        *shared.consumed.lock() = channel.stdout.find(READY_MARKER).map(|i| i + READY_MARKER.len()).unwrap_or(0);
                        match interact.take() {
                            Some(callback) => pending = Some(Box::pin(callback(Interaction { shared: shared.clone(), commands: command_tx.clone() }))),
                            None => channel.eof().await,
                        }
                    }
                    *shared.stdout.lock() = channel.stdout.clone();
                    if event == ExecEvent::Done {
                        channel_done = true;
                        *shared.done.lock() = true;
                    }
                    shared.notify.notify_waiters();
                }
                command = command_rx.recv() => match command {
                    Some(InteractionCommand::Send(line)) => channel.write(line.as_bytes()).await,
                    Some(InteractionCommand::End) | None => channel.eof().await,
                },
                result = async { pending.as_mut().expect("pending interaction").await }, if pending.is_some() => {
                    pending = None;
                    if result.is_err() {
                        // Any failure closes stdin; the remote script treats EOF as "roll back".
                        channel.eof().await;
                    }
                    interaction_result = Some(result);
                }
            }
        }
        let stderr_raw = channel.stderr.clone();
        let stdout_raw = channel.stdout.clone();
        let exit = if auth_failure.is_some() { None } else { channel.finish().await.map(|r| r.code).unwrap_or(None) };
        if !started {
            let code = classify_sudo_failure(&stderr_raw, exit, auth_failure);
            self.remember_failure(&record, &code);
            if code == "SUDO_PASSWORD_INCORRECT" || code == "SUDO_PASSWORD_REQUIRED" {
                *record.sudo_cache.lock() = None;
            }
            self.sessions.logger.warn(
                "Privileged remote operation was not started.",
                json!({ "sessionId": record.session_id, "host": record.host, "purpose": options.purpose, "code": code }),
            );
            return Err(CoreError::new(code.clone(), sudo_message(&code)).with("remoteMessage", clean_remote_message(&stderr_raw, 600)));
        }
        if use_sudo {
            if let Some(p) = record.privilege.lock().as_mut() {
                let mode = if (!password.is_empty() && answered > 0) || p["mode"] == "sudo-password" { "sudo-password" } else { "sudo-nopasswd" };
                p["verified"] = json!(true);
                p["mode"] = json!(mode);
            }
            if !password.is_empty() && answered > 0 && (options.remember_password || used_cache) {
                *record.sudo_cache.lock() = Some((password.clone(), Instant::now() + SUDO_CACHE_IDLE));
            }
        }
        let after_ready = stdout_raw.find(READY_MARKER).map(|i| &stdout_raw[i + READY_MARKER.len()..]).unwrap_or("").to_string();
        if let Some(caps) = ERROR_MARKER.captures(&after_ready) {
            let cleaned = clean_remote_message(&stderr_raw, 600);
            return Err(CoreError::new(format!("{}{}", options.error_prefix, &caps[1]), if cleaned.is_empty() { caps[1].to_string() } else { cleaned })
                .with("remoteExitCode", exit.map(Value::from).unwrap_or(Value::Null)));
        }
        let interaction = match interaction_result {
            Some(Ok(value)) => Some(value),
            Some(Err(error)) => return Err(error),
            None => None,
        };
        Ok((RunResult { code: exit, stdout: after_ready, stderr: clean_remote_message(&stderr_raw, 600), privileged: options.privileged }, interaction))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sudo_failure_classification() {
        assert_eq!(classify_sudo_failure("Sorry, try again.", Some(1), None), "SUDO_PASSWORD_INCORRECT");
        assert_eq!(classify_sudo_failure("bob is not in the sudoers file.", Some(1), None), "SUDO_NOT_PERMITTED");
        assert_eq!(classify_sudo_failure("sudo: sorry, you must have a tty to run sudo", Some(1), None), "SUDO_TTY_REQUIRED");
        assert_eq!(classify_sudo_failure("sh: sudo: not found", Some(127), None), "SUDO_UNAVAILABLE");
        assert_eq!(classify_sudo_failure("x", Some(1), Some("SUDO_PASSWORD_REQUIRED")), "SUDO_PASSWORD_REQUIRED");
        assert_eq!(classify_sudo_failure("other", Some(1), None), "SUDO_FAILED");
    }

    #[test]
    fn remote_message_cleanup() {
        assert_eq!(clean_remote_message("[chj-sudo-auth]line1\n\nline2\x07\n", 600), "line1 | line2");
        assert!(normalize_password(Some("a\nb")).is_err());
        assert_eq!(normalize_password(None).unwrap(), "");
    }

    fn interaction() -> (Interaction, mpsc::UnboundedReceiver<InteractionCommand>) {
        let (commands, rx) = mpsc::unbounded_channel();
        (Interaction { commands, shared: Arc::new(InteractionShared {
            stdout: Mutex::new(String::new()), consumed: Mutex::new(0),
            done: Mutex::new(false), notify: Notify::new(),
        }) }, rx)
    }

    #[tokio::test]
    async fn interaction_consumes_matches_once_and_handles_fragmented_output() {
        let (io, _) = interaction();
        let writer = io.clone();
        let pattern = Regex::new("CHJ_REPLY=([0-9]+)\\n").unwrap();
        let (matched, ()) = tokio::join!(io.wait_for(&pattern, Duration::from_secs(1)), async move {
            *writer.shared.stdout.lock() = "banner\nCHJ_RE".into();
            writer.shared.notify.notify_waiters();
            tokio::task::yield_now().await;
            writer.shared.stdout.lock().push_str("PLY=42\n");
            writer.shared.notify.notify_waiters();
        });
        let matched = matched.unwrap();
        assert_eq!(matched.before, "banner\n");
        assert_eq!(matched.groups[1].as_deref(), Some("42"));
        assert_eq!(io.wait_for(&Regex::new("CHJ_REPLY").unwrap(), Duration::ZERO).await.err().unwrap().code, "REMOTE_COMMAND_TIMEOUT");
        io.shared.stdout.lock().push_str("CHJ_REPLY=43\n");
        let next = io.wait_for(&Regex::new("CHJ_REPLY=([0-9]+)").unwrap(), Duration::ZERO).await.unwrap();
        assert_eq!(next.groups[1].as_deref(), Some("43"));
    }

    #[tokio::test]
    async fn interaction_eof_settles_waiter_and_send_end_have_exact_framing() {
        let (io, mut commands) = interaction();
        io.send("COMMIT"); io.end();
        assert!(matches!(commands.recv().await, Some(InteractionCommand::Send(s)) if s == "COMMIT\n"));
        assert!(matches!(commands.recv().await, Some(InteractionCommand::End)));
        let writer = io.clone();
        let pattern = Regex::new("NEVER").unwrap();
        let (result, ()) = tokio::join!(io.wait_for(&pattern, Duration::from_secs(1)), async move {
            *writer.shared.done.lock() = true;
            writer.shared.notify.notify_waiters();
        });
        assert_eq!(result.err().unwrap().code, "PROTOCOL_ERROR");
    }

}