use std::ffi::OsString;
use std::io::Write;
use std::path::{Path, PathBuf};
use age::ssh::Identity;
use flate2::Compression;
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use clap::{Parser, Subcommand};
use ssh_key::{Algorithm, LineEnding, PrivateKey};
use crate::error::DrawerError;
use crate::parsing::KeyType::*;

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum OperationType {
    Open,
    Close,
    Key,
    Info,
}

#[derive(Debug, PartialEq)]
pub struct DrawerOperation {
    pub operation_type: OperationType,
    pub drawer_file: PathBuf,
    pub key: KeyType,
    pub target_path: Option<PathBuf>,
    pub force: bool,
}

#[derive(Debug, PartialEq)]
pub enum KeyType {
    EnvVar, Path(PathBuf)
}

#[derive(Parser)]
#[command(name = "drawer", about = "Encrypt and decrypt project directories as drawer files")]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Decrypt and expand a drawer file into a directory
    Open {
        drawer_file: PathBuf,
        target_path: Option<PathBuf>,
        #[arg(short = 'i')]
        key: Option<PathBuf>,
        #[arg(short = 'f')]
        force: bool,
    },
    /// Compress and encrypt a directory into a drawer file
    Close {
        drawer_file: PathBuf,
        target_path: Option<PathBuf>,
        #[arg(short = 'i')]
        key: Option<PathBuf>,
    },
    /// Generate a new SSH key and save it to a file
    Key {
        filename: PathBuf,
    },
    /// Show information about a drawer file
    Info {
        drawer_file: PathBuf,
        #[arg(short = 'i')]
        key: Option<PathBuf>,
    },
}

pub fn parse_args<I, T>(args: I) -> Result<DrawerOperation, DrawerError>
where
    I: IntoIterator<Item = T>,
    T: Into<OsString> + Clone,
{
    let cli = Cli::try_parse_from(args)
        .map_err(|e| DrawerError::ParseError(e.to_string()))?;

    match cli.command {
        Command::Open { drawer_file, target_path, key, force } => {
            let key = resolve_key(key);
            Ok(DrawerOperation {
                operation_type: OperationType::Open,
                drawer_file,
                key,
                target_path,
                force,
            })
        }
        Command::Close { drawer_file, target_path, key } => {
            let key = resolve_key(key);
            Ok(DrawerOperation {
                operation_type: OperationType::Close,
                drawer_file,
                key,
                target_path,
                force: false,
            })
        }
        Command::Key { filename } => {
            Ok(DrawerOperation {
                operation_type: OperationType::Key,
                drawer_file: filename,
                key: KeyType::EnvVar,
                target_path: None,
                force: false,
            })
        }
        Command::Info { drawer_file, key } => {
            let key = resolve_key(key);
            Ok(DrawerOperation {
                operation_type: OperationType::Info,
                drawer_file,
                key,
                target_path: None,
                force: false,
            })
        }
    }
}

fn resolve_key(key: Option<PathBuf>) -> KeyType {
    if let Some(k) = key {
        return Path(k);
    }
    EnvVar
}

enum ArchiveInfo {
    SingleFile { size: u64 },
    Folder { file_count: usize },
}

fn archive_info(data: &[u8]) -> Result<ArchiveInfo, DrawerError> {
    let cursor = std::io::Cursor::new(data);
    let gz = GzDecoder::new(cursor);
    let mut archive = tar::Archive::new(gz);
    let mut entries = archive.entries().map_err(|_| DrawerError::UntarFailed)?;

    let first = entries.next()
        .ok_or(DrawerError::UntarFailed)?
        .map_err(|_| DrawerError::UntarFailed)?;

    if first.header().entry_type().is_file() {
        let size = first.header().size().map_err(|_| DrawerError::UntarFailed)?;
        if entries.next().is_none() {
            return Ok(ArchiveInfo::SingleFile { size });
        }
    }

    let cursor = std::io::Cursor::new(data);
    let gz = GzDecoder::new(cursor);
    let mut archive = tar::Archive::new(gz);
    let file_count = archive.entries()
        .map_err(|_| DrawerError::UntarFailed)?
        .filter_map(|e| e.ok())
        .filter(|e| e.header().entry_type().is_file())
        .count();
    Ok(ArchiveInfo::Folder { file_count })
}

fn single_file_entry(data: &[u8]) -> Option<PathBuf> {
    let cursor = std::io::Cursor::new(data);
    let gz = GzDecoder::new(cursor);
    let mut archive = tar::Archive::new(gz);
    let mut entries = archive.entries().ok()?;
    let entry = entries.next()?.ok()?;
    if !entry.header().entry_type().is_file() {
        return None;
    }
    if entries.next().is_some() {
        return None;
    }
    Some(entry.path().ok()?.into_owned())
}

impl DrawerOperation {
    // Returns the default target when no target_path is specified.
    // For close and for directory drawers on open: the drawer file stem.
    // For single-file drawers on open: the filename stored in the archive.
    fn default_target(&self, single_file: Option<&std::path::Path>) -> PathBuf {
        match single_file {
            Some(p) => PathBuf::from(p.file_name().unwrap_or_default()),
            None => PathBuf::from(self.drawer_file.file_stem().unwrap_or_default()),
        }
    }

    pub fn perform(&self) -> Result<(), DrawerError> {
        self.validate()?;
        match self.operation_type {
            OperationType::Open => self.perform_open(),
            OperationType::Close => self.perform_close(),
            OperationType::Key => self.perform_key(),
            OperationType::Info => self.perform_info(),
        }
    }

    fn perform_open(&self) -> Result<(), DrawerError> {
        let encrypted = std::fs::read(&self.drawer_file)
            .map_err(|_| DrawerError::DrawerFileNotFound(self.drawer_file.clone()))?;

        let decrypted = self.decrypt(encrypted)?;

        let sf = single_file_entry(&decrypted);
        let target = self.target_path.as_ref().cloned()
            .unwrap_or_else(|| self.default_target(sf.as_deref()));

        if self.target_path.is_none() {
            println!("(assuming into {})", target.display());
        }

        if !self.force && target.exists() {
            if target.is_file() {
                return Err(DrawerError::TargetNotEmpty(target.clone()));
            }
            let mut entries = std::fs::read_dir(&target)
                .map_err(|_| DrawerError::TargetInvalid(target.clone()))?;
            if entries.next().is_some() {
                return Err(DrawerError::TargetNotEmpty(target.clone()));
            }
        }

        if sf.is_some() {
            let cursor = std::io::Cursor::new(&decrypted);
            let gz = GzDecoder::new(cursor);
            let mut archive = tar::Archive::new(gz);
            let mut entry = archive.entries()
                .map_err(|_| DrawerError::UntarFailed)?
                .next()
                .ok_or(DrawerError::UntarFailed)?
                .map_err(|_| DrawerError::UntarFailed)?;
            entry.unpack(&target).map_err(|_| DrawerError::UntarFailed)?;
        } else {
            let cursor = std::io::Cursor::new(decrypted);
            let gz = GzDecoder::new(cursor);
            let mut archive = tar::Archive::new(gz);
            archive.unpack(&target).map_err(|_| DrawerError::UntarFailed)?;
        }
        Ok(())
    }

    fn decrypt(&self, encrypted: Vec<u8>) -> Result<Vec<u8>, DrawerError> {
        let (key_path, identity) = self.identity()?;

        let decryptor = age::Decryptor::new(encrypted.as_slice())
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        let mut decrypted: Vec<u8> = Vec::new();
        let mut stream = decryptor.decrypt(std::iter::once(&identity as &dyn age::Identity))
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        std::io::Read::read_to_end(&mut stream, &mut decrypted)
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        Ok(decrypted)
    }

    fn identity(&self) -> Result<(PathBuf, Identity), DrawerError> {
        let key_path = self.key_path()?;
        let file = std::fs::File::open(&key_path)
            .map_err(|_| DrawerError::KeyInvalid(key_path.clone()))?;
        let reader = std::io::BufReader::new(file);
        let identity = age::ssh::Identity::from_buffer(reader, None)
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        Ok((key_path, identity))
    }

    fn perform_close(&self) -> Result<(), DrawerError> {
        let target = self.target_path.as_ref().cloned()
            .unwrap_or_else(|| self.default_target(None));

        let mut buf: Vec<u8> = Vec::new();
        {
            let gz = GzEncoder::new(&mut buf, Compression::default());
            let mut archive = tar::Builder::new(gz);
            if target.is_dir() {
                archive.append_dir_all(".", &target)
                    .map_err(|_| DrawerError::TarFailed)?;
            } else {
                let name = target.file_name().unwrap_or_default();
                archive.append_path_with_name(&target, name)
                    .map_err(|_| DrawerError::TarFailed)?;
            }
            let gz = archive.into_inner()
                .map_err(|_| DrawerError::TarFailed)?;
            gz.finish()
                .map_err(|_| DrawerError::TarFailed)?;
        }

        let encrypted = self.encrypt(&buf)?;

        std::fs::write(&self.drawer_file, &encrypted)
            .map_err(|_| DrawerError::WriteFailed(self.drawer_file.clone()))
    }

    fn encrypt(&self, buf: &[u8]) -> Result<Vec<u8>, DrawerError> {
        let (_, identity) = self.identity()?;
        let recipient = age::ssh::Recipient::try_from(identity)
            .map_err(|_| DrawerError::EncryptFailed)?;

        let mut encrypted: Vec<u8> = Vec::new();
        let encryptor = age::Encryptor::with_recipients(
            std::iter::once(&recipient as &dyn age::Recipient)
        ).map_err(|_| DrawerError::EncryptFailed)?;
        let mut writer = encryptor.wrap_output(&mut encrypted)
            .map_err(|_| DrawerError::EncryptFailed)?;
        writer.write_all(buf).map_err(|_| DrawerError::EncryptFailed)?;
        writer.finish().map_err(|_| DrawerError::EncryptFailed)?;
        Ok(encrypted)
    }

    pub fn validate(&self) -> Result<(), DrawerError> {
        if self.operation_type != OperationType::Key {
            if self.drawer_file.extension().and_then(|e| e.to_str()) != Some("drawer") {
                return Err(DrawerError::DrawerFileInvalidExtension(self.drawer_file.clone()));
            }
            self.validate_key()?;
        }
        match self.operation_type {
            OperationType::Open | OperationType::Info => self.validate_drawer_file(),
            OperationType::Close => self.validate_close(),
            OperationType::Key => self.validate_key_gen(),
        }
    }

    fn validate_key_gen(&self) -> Result<(), DrawerError> {
        self.validate_drawer_file()?;
        let parent = match self.drawer_file.parent() {
            None => Path::new("."),
            Some(p) => {
                if p == Path::new("") {
                    Path::new(".")
                } else {
                    p
                }
            },
        };
        if !parent.exists() {
            return Err(DrawerError::KeyOutputFailed(self.drawer_file.clone()));
        }
        Ok(())
    }

    fn key_path(&self) -> Result<PathBuf, DrawerError> {
        match &self.key {
            EnvVar => {
                let val = std::env::var("DRAWER_KEY").map_err(|_| DrawerError::NoKeyPath)?;
                Ok(PathBuf::from(val))
            }
            Path(p) => Ok(p.clone()),
        }
    }

    fn validate_key(&self) -> Result<(), DrawerError> {
        let key_path = self.key_path()?;
        let file = std::fs::File::open(&key_path)
            .map_err(|_| DrawerError::KeyInvalid(key_path.clone()))?;
        let reader = std::io::BufReader::new(file);
        let identity = age::ssh::Identity::from_buffer(reader, None)
            .map_err(|_| DrawerError::KeyNotValidSsh)?;
        age::ssh::Recipient::try_from(identity)
            .map_err(|_| DrawerError::KeyNotValidSsh)?;
        Ok(())
    }

    // Ensure the drawer file actually exists
    fn validate_drawer_file(&self) -> Result<(), DrawerError> {
        if !self.drawer_file.exists() {
            return Err(DrawerError::DrawerFileNotFound(self.drawer_file.clone()));
        }
        Ok(())
    }

    fn perform_info(&self) -> Result<(), DrawerError> {
        let encrypted = std::fs::read(&self.drawer_file)
            .map_err(|_| DrawerError::DrawerFileNotFound(self.drawer_file.clone()))?;
        let decrypted = self.decrypt(encrypted)?;
        let info = archive_info(&decrypted)?;
        match info {
            ArchiveInfo::SingleFile { size } =>
                println!("{} (file, {} bytes)", self.drawer_file.display(), size),
            ArchiveInfo::Folder { file_count } =>
                println!("{} (folder, {} files)", self.drawer_file.display(), file_count),
        }
        Ok(())
    }

    fn perform_key(&self) -> Result<(), DrawerError> {
        let private_key = PrivateKey::random(&mut rand::rngs::OsRng, Algorithm::Ed25519)
            .map_err(|_| DrawerError::KeyGenerateFailed)?;
        let pem = private_key.to_openssh(LineEnding::LF)
            .map_err(|_| DrawerError::KeyGenerateFailed)?;
        std::fs::write(&self.drawer_file, pem.as_bytes())
            .map_err(|_| DrawerError::KeyOutputFailed(self.drawer_file.clone()))
    }

    fn validate_close(&self) -> Result<(), DrawerError> {
        let target = self.target_path.as_ref().cloned()
            .unwrap_or_else(|| self.default_target(None));
        if self.target_path.is_none() {
            println!("(assuming from {})", target.display());
        }
        let meta = std::fs::metadata(&target)
            .map_err(|_| DrawerError::TargetInvalid(target.clone()))?;
        if !meta.is_dir() && !meta.is_file() {
            return Err(DrawerError::TargetInvalid(target.clone()));
        }
        if meta.is_dir() && std::fs::read_dir(&target).is_err() {
            return Err(DrawerError::TargetInvalid(target.clone()));
        }
        let canon_target = target.canonicalize()
            .unwrap_or_else(|_| target.clone());
        let drawer_parent = match self.drawer_file.parent() {
            None => Path::new("."),
            Some(p) if p == Path::new("") => {
                Path::new(".")
            },
            Some(p) => p,
        };
        let canon_drawer = drawer_parent.canonicalize()
            .unwrap_or_else(|_| drawer_parent.to_path_buf())
            .join(self.drawer_file.file_name().unwrap_or_default());
        if canon_drawer.starts_with(&canon_target) {
            return Err(DrawerError::DrawerInsideTarget);
        }
        Ok(())
    }
}