use std::{
    ffi::{CStr, CString},
    fmt::Write as _,
    fs,
    io::{stdin, IsTerminal},
    mem,
    os::fd::AsFd,
    path::Path,
    process::{Command, Stdio},
    ptr, thread,
    time::Duration,
};

use anyhow::{anyhow, bail, ensure, Result};
use bcachefs_kernel::c::{
    bch_key, bch_sb_handle,
    bch2_chacha20, bch_encrypted_key, bch_sb_field_crypt,
};
use bch_bindgen::keyutils::{self, keyctl_search};
use log::info;
use rustix::termios;

use crate::prompt::{Waited, Watch};
use crate::unlock_socket::UnlockSocket;
use uuid::Uuid;
use zeroize::{ZeroizeOnDrop, Zeroizing};

use crate::prompt::Prompt;
use crate::ErrnoError;

/// Check if a superblock has an encrypted passphrase set.
pub fn sb_is_encrypted(sb: &bch_sb_handle) -> bool {
    sb.sb()
        .crypt()
        .map(|c| c.key().is_encrypted())
        .unwrap_or(false)
}

/// Target keyring for key storage.
#[derive(Clone, Copy, Debug, Default, clap::ValueEnum)]
#[clap(rename_all = "snake_case")]
pub enum Keyring {
    Session,
    #[default]
    User,
    UserSession,
}

impl Keyring {
    pub fn id(self) -> i32 {
        match self {
            Keyring::Session => keyutils::KEY_SPEC_SESSION_KEYRING,
            Keyring::User => keyutils::KEY_SPEC_USER_KEYRING,
            Keyring::UserSession => keyutils::KEY_SPEC_USER_SESSION_KEYRING,
        }
    }
}

#[derive(Clone, Debug, Default, clap::ValueEnum, strum::Display)]
pub enum UnlockPolicy {
    /// Don't ask for passphrase, if the key cannot be found in the keyring just
    /// fail
    Fail,
    /// Wait for passphrase to become available before mounting
    Wait,
    /// Interactively prompt the user for a passphrase
    #[default]
    Ask,
    /// Try to read the passphrase from `stdin` without prompting
    Stdin,
}

/// An unlocked filesystem, and whether we are holding the key that unlocked it.
///
/// The difference decides how the kernel gets it. A key we derived ourselves
/// can be handed over directly, which is what the mount helper does; one that
/// was already in a keyring we have never seen the bytes of, so all we can do
/// is leave it there for bch2_request_key() to find.
pub enum Unlocked {
    Key(PassphraseCorrect),
    InKeyring,
}

impl Unlocked {
    /// The key as the kernel's `user_key` mount parameter wants it, if we have
    /// it to give.
    ///
    /// The same bytes add_key() would have put in a keyring: both ends treat a
    /// struct bch_key as an opaque block, so there is no byte order here to get
    /// wrong.
    pub fn hex(&self) -> Option<Zeroizing<String>> {
        let Self::Key(k) = self else { return None };

        // SAFETY: bch_key is four __le64s - no padding, no pointers, and the
        // borrow lives no longer than @k.
        let bytes: &[u8] = unsafe {
            std::slice::from_raw_parts(
                ptr::addr_of!(k.passphrase_key).cast(),
                mem::size_of_val(&k.passphrase_key),
            )
        };

        let mut s = String::with_capacity(bytes.len() * 2);
        for b in bytes {
            let _ = write!(s, "{b:02x}");
        }
        Some(Zeroizing::new(s))
    }

    /// Put it where bch2_request_key() will look, for the mount(2) path - that
    /// has nowhere to carry a parameter, so the keyring is the only channel.
    pub fn to_keyring(&self) -> Result<()> {
        match self {
            Self::Key(k)    => KeyHandle::new(k, Keyring::User).map(|_| ()),
            Self::InKeyring => Ok(()),
        }
    }
}

impl UnlockPolicy {
    pub fn apply(&self, sb: &bch_sb_handle) -> Result<Unlocked> {
        let uuid = sb.sb().uuid();

        info!("Using filesystem unlock policy '{self}' on {uuid}");

        match self {
            // Somebody else's key, in a keyring: we never see the bytes.
            Self::Fail => KeyHandle::new_from_search(&uuid).map(|_| Unlocked::InKeyring),
            Self::Wait => KeyHandle::wait_for_unlock(&uuid).map(|_| Unlocked::InKeyring),
            Self::Ask => {
                let passphrase = Passphrase::ask_in_terminal()?;
                Ok(Unlocked::Key(passphrase
                    .check(sb)
                    .ok_or_else(|| anyhow!("incorrect passphrase"))?))
            }
            Self::Stdin => {
                let passphrase = Passphrase::read_from_stdin()?;
                Ok(Unlocked::Key(passphrase
                    .check(sb)
                    .ok_or_else(|| anyhow!("incorrect passphrase"))?))
            }
        }
    }
}


/// Proof that a bcachefs key has been added to or found in the kernel keyring.
pub struct KeyHandle;

impl KeyHandle {
    pub fn format_key_name(uuid: &Uuid) -> CString {
        CString::new(format!("bcachefs:{uuid}")).unwrap()
    }

    pub fn new(passphrase_correct: &PassphraseCorrect, keyring: Keyring) -> Result<Self> {
        let key_name = Self::format_key_name(&passphrase_correct.fs_uuid);
        let key_name = CStr::as_ptr(&key_name);
        let key_type = c"user";

        let key_id = unsafe {
            keyutils::add_key(
                key_type.as_ptr(),
                key_name,
                ptr::addr_of!(passphrase_correct.passphrase_key).cast(),
                mem::size_of_val(&passphrase_correct.passphrase_key),
                keyring.id(),
            )
        };

        if key_id > 0 {
            info!("Added key to keyring");
            Ok(KeyHandle)
        } else {
            Err(anyhow!("failed to add key to keyring: {}", errno::errno()))
        }
    }

    fn search_keyring(keyring: i32, key_name: &CStr) -> Result<()> {
        let key_name = CStr::as_ptr(key_name);
        let key_type = c"user";

        let key_id = unsafe { keyctl_search(keyring, key_type.as_ptr(), key_name, 0) };

        if key_id > 0 {
            info!("Found key in keyring");
            Ok(())
        } else {
            Err(ErrnoError(errno::errno()).into())
        }
    }

    pub fn new_from_search(uuid: &Uuid) -> Result<Self> {
        let key_name = Self::format_key_name(uuid);

        Self::search_keyring(keyutils::KEY_SPEC_SESSION_KEYRING, &key_name)
            .or_else(|_| Self::search_keyring(keyutils::KEY_SPEC_USER_KEYRING, &key_name))
            .or_else(|_| Self::search_keyring(keyutils::KEY_SPEC_USER_SESSION_KEYRING, &key_name))
            .map(|_| KeyHandle)
    }

    fn wait_for_unlock(uuid: &Uuid) -> Result<Self> {
        loop {
            match Self::new_from_search(uuid) {
                Err(_) => thread::sleep(Duration::from_secs(1)),
                r => break r,
            }
        }
    }
}

#[derive(ZeroizeOnDrop)]
pub struct Passphrase(CString);

impl Passphrase {
    pub(crate) fn get(&self) -> &CStr {
        &self.0
    }

    /// From a line somebody typed or sent. The line ending is theirs, not part
    /// of the passphrase - and it is CR, LF or both, because an initramfs
    /// console that nothing has configured sends CR for enter.
    pub fn from_line(line: &str) -> Result<Self> {
        Ok(Self(CString::new(line.trim_end_matches(['\n', '\r']))?))
    }

    /// Ask, and take the passphrase from whoever supplies one first.
    ///
    /// @socket is a second way to answer for callers that want one - mounting
    /// does, because whoever is booting the machine may not be at it. It has
    /// already checked what it was given, so a passphrase it holds is right.
    pub fn ask_and_check(
        sb: &bch_sb_handle,
        mut socket: Option<&mut UnlockSocket<'_>>,
    ) -> Result<PassphraseCorrect> {
        let typed = match StdinType::detect() {
            // A terminal plymouth is drawing over shows nothing, and this
            // prompt has no timeout - it waits there for a person who cannot
            // see it. Which way a question should go when that happens is
            // Prompt::detect()'s decision, and asking it is not the same as
            // copying it: it answers Agent for a covered terminal only where
            // there is an agent framework to reach plymouth by, and Terminal
            // where there isn't - an initramfs that isn't systemd, which
            // leaves the invisible terminal as the only thing we have.
            StdinType::Terminal if matches!(Prompt::detect(), Some(Prompt::Agent)) =>
                Some(Self::ask_from_systemd_and_check(sb)?),
            StdinType::Terminal => Self::read_from_terminal(
                    "Enter passphrase: ",
                    socket.as_deref_mut().map(|s| s as &mut dyn Watch),
                )?
                .and_then(|p| p.check(sb)),
            StdinType::DevNull => Some(Self::ask_from_systemd_and_check(sb)?),
            StdinType::Other => Self::read_from_stdin()?.check(sb),
        };

        // Nothing typed was right - but the socket answers by ending the
        // prompt, so a passphrase waiting there is why we are here.
        typed
            .or_else(|| socket.and_then(UnlockSocket::take))
            .ok_or_else(|| anyhow!("incorrect passphrase"))
    }

    fn ask_from_systemd_and_check(sb: &bch_sb_handle) -> Result<PassphraseCorrect> {
        let uuid = sb.sb().uuid();
        let label = crate::prompt::fs_name(sb);
        for i in 0..3 {
            let mut command = Command::new("systemd-ask-password");
            command
                // Reached either with no tty, where this is a no-op (measured:
                // identical timeout behaviour with and without), or with one
                // plymouth is covering, where it is what stops us prompting
                // onto the invisible terminal we just declined to use.
                .arg("--no-tty")
                .arg("--icon=drive-harddisk")
                .arg(format!("--id=cryptsetup:UUID={}", uuid.as_hyphenated()))
                .arg("--keyname=cryptsetup")
                .arg("--credential=cryptsetup.passphrase")
                .arg("--timeout=0")
                .arg("--multiple")
                .arg("-n");
            if i == 0 {
                command.arg("--accept-cached");
            }
            let output = command
                .arg(format!("Please enter passphrase for disk {label}:"))
                .stdin(Stdio::inherit())
                .stderr(Stdio::inherit())
                .output()?;
            if !output.status.success() {
                bail!("systemd-ask-password returned an error");
            }
            for passphrase in output.stdout.split(|b| *b == b'\0') {
                let p = Self(
                    CString::new(passphrase)
                        .expect("passphrase should not contain a NUL byte because the output was split on NUL bytes")
                );
                if let Some(passphrase_correct) = p.check(sb) {
                    return Ok(passphrase_correct);
                }
            }
        }
        bail!("incorrect passphrase limit reached");
    }

    /// Prompt for a passphrase with echo disabled.
    fn ask_in_terminal_with_prompt(prompt: &str) -> Result<Self> {
        Ok(Self::read_from_terminal(prompt, None)?
            .expect("no watch was given, so nothing else could have answered"))
    }

    /// `None` when @watch answered before anything was typed.
    fn read_from_terminal(prompt: &str, watch: Option<&mut dyn Watch>) -> Result<Option<Self>> {
        let old = termios::tcgetattr(stdin())?;
        let mut new = old.clone();
        new.local_modes.remove(termios::LocalModes::ECHO);
        /*
         * We may be prompting on an early-boot console (initramfs) that no
         * shell has ever configured: without ICRNL, enter sends '\r', which
         * read_line() doesn't terminate on - keystrokes appear eaten, and
         * the eventually-assembled passphrase has embedded '\r's and is
         * rejected. Ensure line-input sanity rather than inheriting it:
         */
        new.input_modes.insert(termios::InputModes::ICRNL);
        new.local_modes.insert(termios::LocalModes::ICANON);
        termios::tcsetattr(stdin(), termios::OptionalActions::Flush, &new)?;

        eprint!("{}", prompt);

        // ICANON above is what makes this safe to poll: stdin only becomes
        // readable once a whole line has been typed, so the read below has one
        // waiting and does not block either.
        let waited = crate::prompt::wait(stdin().as_fd(), None, watch);

        let mut line = Zeroizing::new(String::new());
        let res = match waited {
            Ok(Waited::Readable) => stdin().read_line(&mut line).map(|_| ()),
            _ => Ok(()),
        };

        termios::tcsetattr(stdin(), termios::OptionalActions::Flush, &old)?;
        eprintln!();
        res?;

        match waited? {
            Waited::Readable => Ok(Some(Self::from_line(&line)?)),
            // Moot and Timeout can't arrive: no timeout is passed, and nothing
            // withdraws a passphrase prompt.
            _ => Ok(None),
        }
    }

    // blocks indefinitely if no input is available on stdin
    pub fn ask_in_terminal() -> Result<Self> {
        Self::ask_in_terminal_with_prompt("Enter passphrase: ")
    }

    /// Prompt for a new passphrase twice and verify they match.
    pub fn ask_for_new_passphrase() -> Result<Self> {
        if !stdin().is_terminal() {
            return Self::read_from_stdin();
        }
        let pass1 = Self::ask_in_terminal_with_prompt("Enter new passphrase: ")?;
        let pass2 = Self::ask_in_terminal_with_prompt("Enter same passphrase again: ")?;
        ensure!(pass1.get().to_bytes() == pass2.get().to_bytes(), "Passphrases do not match");
        Ok(pass1)
    }

    // blocks indefinitely if no input is available on stdin
    pub fn read_from_stdin() -> Result<Self> {
        info!("Trying to read passphrase from stdin...");

        let mut line = Zeroizing::new(String::new());
        stdin().read_line(&mut line)?;

        Ok(Self(CString::new(line.trim_end_matches('\n'))?))
    }

    pub fn read_from_file(passphrase_file: impl AsRef<Path>) -> Result<Self> {
        let passphrase_file = passphrase_file.as_ref();

        info!(
            "Attempting to unlock key with passphrase from file {}",
            passphrase_file.display()
        );

        let passphrase = Zeroizing::new(fs::read_to_string(passphrase_file)?);

        Ok(Self(CString::new(passphrase.trim_end_matches('\n'))?))
    }

    fn derive(&self, crypt: &bch_sb_field_crypt) -> bch_key {
        let crypt_ptr = (crypt as *const bch_sb_field_crypt).cast_mut();

        unsafe { bch_bindgen::c::derive_passphrase(crypt_ptr, self.get().as_ptr()) }
    }

    /// Re-encrypt a filesystem key with this passphrase.
    /// Returns the encrypted key suitable for writing to crypt->key.
    pub fn encrypt_key(
        &self,
        sb: &bch_sb_handle,
        key: bch_key,
    ) -> bch_encrypted_key {
        let crypt = sb.sb().crypt().expect("called on encrypted fs");
        let mut new_key = bch_encrypted_key::new_unencrypted(key);
        
        let mut passphrase_key: bch_key = self.derive(crypt);

        unsafe {
            bch2_chacha20(
                ptr::addr_of_mut!(passphrase_key),
                sb.sb().nonce(),
                ptr::addr_of_mut!(new_key).cast(),
                mem::size_of_val(&new_key),
            )
        };

        new_key
    }

    pub fn check(&self, sb: &bch_sb_handle) -> Option<PassphraseCorrect> {
        let crypt = sb
            .sb()
            .crypt()
            .expect("superblock should have crypt when calling Passphrase::check");
        assert!(
            crypt.key().is_encrypted(),
            "sb_key should be encrypted when calling Passphrase::check",
        );

        let mut passphrase_key: bch_key = self.derive(crypt);

        let mut cleartext_sb_key = crypt.key().clone();
        unsafe {
            bch2_chacha20(
                ptr::addr_of_mut!(passphrase_key),
                sb.sb().nonce(),
                ptr::addr_of_mut!(cleartext_sb_key).cast(),
                mem::size_of_val(&cleartext_sb_key),
            )
        };
        if cleartext_sb_key.is_encrypted() {
            return None;
        }

        Some(PassphraseCorrect {
            fs_uuid: sb.sb().uuid(),
            passphrase_key,
            cleartext_sb_key,
        })
    }
}

pub struct PassphraseCorrect {
    pub fs_uuid:          Uuid,
    pub passphrase_key:   bch_key,
    pub cleartext_sb_key: bch_encrypted_key,
}

/// A passphrase is the one question a pipe can answer, which is why this
/// doesn't go through crate::prompt: Other reads the passphrase off stdin
/// rather than meaning nobody is there.
enum StdinType {
    Terminal,
    DevNull,
    Other,
}

impl StdinType {
    fn detect() -> StdinType {
        if stdin().is_terminal() {
            StdinType::Terminal
        } else if crate::prompt::stdin_is_dev_null() {
            StdinType::DevNull
        } else {
            StdinType::Other
        }
    }
}
