use std::arch::asm;
use std::ffi::{CStr, CString};
use std::fs::File;
use std::os::fd::AsRawFd;
use std::os::unix::fs::PermissionsExt;
use nix::libc::{self, setgid, CLONE_NEWCGROUP, MS_NODEV, MS_NOSUID};
use nix::sched::{clone, CloneCb, CloneFlags};
use nix::sys::signal::{kill, Signal};
use nix::sys::wait::{wait, waitpid, waitid, WaitPidFlag};
use nix::unistd::{chdir, chroot, dup2, execv, pivot_root, setuid, sleep, Gid, Pid, Uid, User, setgroups};
use nix::mount::{mount, MntFlags, MsFlags, umount2, umount};
use nix::env::clearenv;
use std::path::{Path, PathBuf};
use clap::Parser;
use error::{Result, RockerError};
mod error;
use uuid;
use std::{io, fs};
use toml;
use serde::{Deserialize, Serialize};
use std::fmt::Display;

static WORKSPACE: &str = "/root/rocker";
static USER_NAME: &str = "rocker";
static INFO_FILE: &str = "info.toml";
static LOCK_FILE: &str = ".lock";
static mut STACK: [u8; 1024*1024*1] = [0; 1024*1024*1];
static CLONE_FLAG: i32 = 0b1101100000000100000000000000000;  // CLONE_NEWNS | CLONE_NEWUTS | CLONE_NEWPID | CLONE_NEWIPC | CLONE_NEWNET;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct RockerArgs {
    // --wait/--log --run /bin/bash --image busybox 
    #[arg(long)]
    run: Option<String>,
    #[arg(long)]
    image: Option<String>,
    #[arg(long)]
    volume: Option<String>,
    #[arg(long)]
    env: Option<String>,

    #[arg(long)]
    log: bool,
    #[arg(long)]
    wait: bool,

    // --logs container_id
    #[arg(long)]
    logs: Option<String>,

    // --ps
    #[arg(long)]
    ps: bool,
    // --psa
    #[arg(long)]
    psa: bool,
    
    // rm container_id
    #[arg(long)]
    rm: Option<String>,

    // stop container_id
    #[arg(long)]
    stop: Option<String>
    
}


/// 从images解压到volumes
fn extend_image(image_name: &String) -> Result<PathBuf> {
    // 源文件
    let image_path = Path::new(WORKSPACE).join("images").join(image_name);
    if image_path.exists() == false {
        return Err(RockerError::from(io::Error::new(io::ErrorKind::NotFound, "未找到镜像")));
    }
    let image_path_str = image_path.to_str().unwrap();  // 安全的unwrap

    // volumes只读层
    let volume_path = Path::new(WORKSPACE).join("volumes").join(image_name);
    if volume_path.exists() {
        return Ok(volume_path);
    } else {
        create_dir(&volume_path, true)?;
    }
    let volume_path_str = volume_path.to_str().unwrap();  // 安全的unwrap
    
    // 解压缩
    let out = std::process::Command::new("tar")
        .arg("-xvf")
        .arg(image_path_str)
        .arg("-C")
        .arg(volume_path_str)
        .output()?;
    let std_out = String::from_utf8_lossy(&out.stdout);
    let std_err = String::from_utf8_lossy(&out.stderr);
    if std_err.len() == 0 {
        println!("解压缩完毕: {std_out:?}");
        Ok(volume_path)
    } else {
        // 删除 volume_path
        std::fs::remove_dir_all(volume_path)?;
        Err(RockerError::from(io::Error::new(io::ErrorKind::Other, format!("解压缩镜像失败: {std_err}"))))
    }
}

fn init_container_lock<P: AsRef<Path>>(container_work_path:P) -> Result<i32> {
    use nix::sys::stat::Mode;
    use nix::fcntl::{OFlag, open};
    let lock_path = container_work_path.as_ref().join(LOCK_FILE);
    let lock_path_str = lock_path.as_os_str();
    let oflag = OFlag::O_RDWR | OFlag::O_CREAT;
    let mode = Mode::empty();
    Ok(open(lock_path_str, oflag, mode)?)
}

fn init_container_overlay<P: AsRef<Path>>(volume_path: P, upper_path: P, merged_path: P) -> Result<()> {
    let lower_dir = volume_path.as_ref().to_string_lossy().to_string();
    let upper_dir = upper_path.as_ref().to_string_lossy().to_string();
    let merged_dir = merged_path.as_ref().to_string_lossy().to_string();

    let dirs = format!("lowerdir={lower_dir},upperdir={upper_dir},workdir={merged_dir}");
    println!("dirs: {dirs:?}");
    let out = std::process::Command::new("mount")
        .arg("-t")
        .arg("overlay")
        .arg("overlay")
        .arg("-o")
        .arg(dirs)
        .arg(merged_dir)
        .output()?;
    
    let std_out = String::from_utf8_lossy(&out.stdout);
    let std_err = String::from_utf8_lossy(&out.stderr);

    if std_err.len() == 0 {
        println!("容器文件系统创建完成: {std_out:?}");
    } else {
        return Err(RockerError::from(io::Error::new(io::ErrorKind::Other, format!("容器文件系统创建失败: {std_err:?}"))));
    }
    Ok(())
}

fn init_container_custom_volume<P: AsRef<Path>>(container_merged_path: P, custom_volume_s: &String) -> Result<()> {
    for custom_volume in custom_volume_s.split(",") {
        let custom_volume_v = custom_volume.split(":").collect::<Vec<&str>>();
        if custom_volume_v.len() < 2 {
            return Err(RockerError::OtherError(format!("volume 参数格式不正确: {custom_volume}")));
        }
        let host_path = custom_volume_v[0];
        let container_path_buf = {
            if custom_volume_v[1].starts_with("/") {
                container_merged_path.as_ref().join(&custom_volume_v[1][1..])
            } else {
                container_merged_path.as_ref().join(&custom_volume_v[1])
            }
        };
        let container_path = container_path_buf.to_string_lossy().to_string();

        // 创建宿主机和容器内的目录
        create_dir(Path::new(host_path), true)?;
        create_dir(&container_path, true)?;

        // 绑定
        let out = std::process::Command::new("mount")
            .arg("-o")
            .arg("bind")
            .arg(host_path)
            .arg(container_path)
            .output()?;
        let std_out = String::from_utf8_lossy(&out.stdout);
        let std_err = String::from_utf8_lossy(&out.stderr);
        if std_err.len() == 0 {
            println!("创建自定义 volume: {custom_volume:?}");
        } else {
            return Err(RockerError::OtherError(format!("创建volume失败: {std_err}")))
        }
    }
    Ok(())
}


fn init_container_pivot<P: AsRef<Path>>(merged_path: P) -> Result<()> {
    // 在我们没有设置 chroot之前, 需要先把所有挂载点的传播类型改为 private, 避免进程中的系统调用污染全局
    mount(None::<&str>, "/", None::<&str>, MsFlags::MS_PRIVATE | MsFlags::MS_REC, None::<&str>)?;

    // 修改overlayfs 为rootfs
    chdir(merged_path.as_ref())?;
    let pwd_path = std::env::current_dir()?;
    let pwd_str = pwd_path.to_string_lossy().to_string();

    // 挂载bind
    mount(Some(pwd_str.as_str()), pwd_str.as_str(), Some("bind"), MsFlags::MS_BIND | MsFlags::MS_REC, Some(""))?;
    
    // 创建 rootfs/.pivot_root 目录用于存储 old_root
    let pivot_root_dir = format!("{pwd_str}/.pivot_root");

    // 将系统rootfs切换到新的rootfs, 并设置权限
    std::fs::create_dir(&pivot_root_dir)?;
    std::fs::set_permissions(&pivot_root_dir, PermissionsExt::from_mode(0o777))?;
    pivot_root(pwd_str.as_str(), pivot_root_dir.as_str())?;

    // 修改当前进程工作目录(注意我们之前已经到rootfs内, 并且把根目录设置完毕了)
    chdir("/")?;

    // 卸载 old_root, 并删除临时文件
    umount2(".pivot_root", MntFlags::MNT_DETACH).unwrap();
    std::fs::remove_dir(".pivot_root").unwrap();
    Ok(())
}


fn init_container_mount() -> Result<()> {
    // 挂载proc
    let mount_flags = MsFlags::MS_NODEV | MsFlags::MS_NOEXEC | MsFlags::MS_NOSUID;
    mount(Some("proc"), "/proc", Some("proc"), mount_flags, Some(""))?;  

    // 挂载dev
    mount(Some("tmpfs"), "/dev", Some("tmpfs"), mount_flags, Some("mode=755"))?;  
    Ok(())
}

fn init_container_log(log: bool) -> Result<()> {
    let log_path = Path::new("logs");
    create_dir(log_path, true)?;
    let log_fd = File::create(log_path.join("log"))?;
    if log {
        unsafe {
            let log_fd_raw = log_fd.as_raw_fd();
            dup2(log_fd_raw, 1)?;
            dup2(log_fd_raw, 2)?;
        }
    }
    Ok(())
}

fn init_container_env() -> Result<()> {
    unsafe {
        clearenv().map_err(|e|RockerError::OtherError(format!("清除env失败: {e:?}")))?;
    }
    Ok(())
}
fn init_container_user(uid: Uid, gid: Gid) -> Result<()>{
    unsafe {
        setgid(gid.as_raw());
        setgroups(&[gid])?;
    }
    setuid(uid)?;
    Ok(())
}

fn create_dir<P: AsRef<Path>>(path: P, is_any:bool) -> Result<()> {
    fs::create_dir_all(&path)?;
    if is_any {
        fs::set_permissions(&path, PermissionsExt::from_mode(0o777))?;
    }
    Ok(())
}


fn parse_cmd(run: &String) -> Vec<CString>{
    let args= run
        .split(" ")
        .filter_map(|s| {
            match CString::new(s) {
                Ok(cs) => {
                    Some(cs)
                }
                Err(e) => {
                    println!("{e:?}");
                    None
                }
            }
        })
        .collect::<Vec<CString>>();
    args
}


fn check_container_is_running(pid: &Pid, main_exe: &Path) -> Result<bool> {
    // 检查pid对应的exe是否和外部传过来的相同
    let child_exe_s= format!("/proc/{pid}/exe");
    let child_exe_path = Path::new(child_exe_s.as_str());
    let target_child_exe_path = fs::read_link(child_exe_path)?;
    if target_child_exe_path != main_exe {
        return Ok(true);
    }
    Ok(false)
}

fn run_container(_container_id: &String, cmd: &String, args: &RockerArgs, volume_path: &PathBuf) -> Result<i32> {
    // 禁止同时wait和log
    if args.wait && args.log {
        return Err(RockerError::OtherError("--wait/--log 禁止同时使用".to_string()));
    }

    // 初始化容器工作目录
    let container_work_path = Path::new(WORKSPACE).join("containers").join(&_container_id);
    let container_upper_path = container_work_path.join("upper");
    let container_merged_path = container_work_path.join("merged");
    create_dir(&container_work_path, true)?;
    create_dir(&container_upper_path, true)?;
    create_dir(&container_merged_path, true)?;

    let rocker_user_info =  User::from_name(USER_NAME)?.ok_or(RockerError::OtherError(format!("没找到 用户: {USER_NAME}")))?;
    let rocker_uid = rocker_user_info.uid;
    let rocker_gid = rocker_user_info.gid;

    let _cb = || {
        init_container_lock(&container_work_path).unwrap();
        init_container_overlay(volume_path, &container_upper_path, &container_merged_path).unwrap();
        if let Some(custom_volume) = &args.volume {
            init_container_custom_volume(&container_merged_path, custom_volume).unwrap();
        }
        init_container_pivot(&container_merged_path).unwrap();
        init_container_mount().unwrap();
        init_container_log(args.log).unwrap();
        init_container_env().unwrap();
        init_container_user(rocker_uid, rocker_gid).unwrap();
        
        let cmd_vec = parse_cmd(cmd);
        match execv(&cmd_vec[0], &cmd_vec) {
            Err(e) => {
                println!("execv {cmd_vec:?}失败: {e:?}");
            }
            _ => {},
        };
        0
    };
    
    let main_exe = std::env::current_exe()?;
    unsafe {
        match clone(Box::new(_cb), STACK.as_mut_slice(), CloneFlags::from_bits_truncate(CLONE_FLAG), None) {
            Ok(child_pid) =>  {
                println!("clone ok: {child_pid:?}");

                // check_container_is_running
                let mut cnt = 0;
                while let Ok(running) = check_container_is_running(&child_pid, &main_exe) {
                    cnt += 1;
                    std::thread::sleep(std::time::Duration::from_millis(10));
                    if running {
                        break;
                    }
                    if cnt > 1000 {
                        println!("{child_pid} 启动 超时");
                        break;
                    }
                }
                
                // wait
                if args.wait {
                    match waitpid(child_pid, Some(WaitPidFlag::WUNTRACED)) {
                        Ok(status) => {
                            println!("{child_pid:?} exit: {status:?}");
                        }
                        Err(e) => {
                            println!("{child_pid} wait err: {e}");
                        }
                    }
                }
                Ok(child_pid.as_raw())
            }
            Err(e) => {
                return Err(RockerError::OtherError(format!("clone err: {e}")));
            }
        }
    }
}

#[derive(Deserialize, Serialize, Debug, PartialEq)]
enum ContainerStatus {
    READY,
    RUNNING,
    STOP,
}

impl Display for ContainerStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::READY => write!(f, "😀"),
            Self::RUNNING => write!(f, "✅"),
            Self::STOP => write!(f, "❌"),
        }
    }
    
}


#[derive(Deserialize, Serialize, Debug)]
struct ContainerInfo {
    id: String,
    pid: i32,
    run: String,  // /bin/bash
    image: String,  // busybox
    volume: String,  // /root/tmp:/root/tmp,/root/tmp1:/root/tmp1
    env: String,  // a=1,b=2,c=3 或者 env文件路径
    status: ContainerStatus,
}

impl Display for ContainerInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let volume: String = self.volume.chars().take(20).collect();
        let env: String = self.env.chars().take(20).collect();
        write!(f, "\x1b[4m{:<10} {:<8} {:<10} {:<20} {:<20} {:<20} {:<10}\x1b[24m", self.id, self.pid, self.image, self.run, volume, env, &self.status)
    }
}

fn save_container_info(args: &RockerArgs, container_id: &String, pid: i32) -> Result<()> {
    let container_info_path = Path::new(WORKSPACE).join("containers").join(container_id).join(INFO_FILE);
    let container_info = ContainerInfo {
        id: container_id.clone(),
        pid: pid,
        run: args.run.as_ref().unwrap().clone(),
        image: args.image.as_ref().unwrap().clone(),
        volume: args.volume.clone().unwrap_or("".to_string()),
        env: args.env.clone().unwrap_or("".to_string()),
        status: ContainerStatus::READY,
    };
    let toml_str = toml::to_string(&container_info)?;
    fs::write(container_info_path, toml_str)?;
    Ok(())
}

fn get_container_info(container_id: &str) -> Result<ContainerInfo> {
    let container_work_path = Path::new(WORKSPACE).join("containers").join(container_id);
    let container_info_path = container_work_path.join(INFO_FILE);
    let lock_path = container_work_path.join(LOCK_FILE);
    let info_str = fs::read_to_string(container_info_path)?;
    let mut container_info: ContainerInfo = toml::from_str(&info_str)?;

    // 判断是否正在运行, 首先得到该容器所有的fd
    let proc_fd_path = Path::new("/proc").join(container_info.pid.to_string()).join("fd");
    let is_running = if let Ok(fd_dir) = fs::read_dir(proc_fd_path) {
        fd_dir.filter_map(|p|p.ok())
            .filter_map(|f| fs::read_link(f.path()).ok())
            .any(|p|p == lock_path)
    } else {
        false
    };
    if is_running {
        container_info.status = ContainerStatus::RUNNING;
    } else {
        container_info.status = ContainerStatus::STOP;
    }
    Ok(container_info)
}

fn get_all_container_info() -> Result<Vec<ContainerInfo>> {
    let containers_path = Path::new(WORKSPACE).join("containers");
    let all_containers_info = fs::read_dir(containers_path)?
        .map(|res| res.map(|e| e.file_name()))
        .filter_map(|p| p.ok())
        .map(|f|f.to_string_lossy().to_string())
        .filter_map(|s|get_container_info(s.as_str()).ok())
        .collect::<Vec<ContainerInfo>>();
    Ok(all_containers_info)
}

/// 读取所有容器的状态
fn show_containers(is_show_all: bool) -> Result<()> {
    println!("{:<10} {:<8} {:<10} {:<20} {:<20} {:<20} {:<10}", "id", "pid", "image", "run", "volume", "env", "status");
    for container_info in get_all_container_info()? {
        if is_show_all{
            println!("{container_info}");
        } else if container_info.status == ContainerStatus::RUNNING {
            println!("{container_info}");
        }
    }
    Ok(())
}

fn stop_container(containers_id: &str, is_remove: bool) -> Result<()> {    
    for container_id in containers_id.split(" ") {
        if let Ok(container_info) = get_container_info(container_id) {
            let container_work_path = Path::new(WORKSPACE).join("containers").join(container_id);
            let container_merged_path = container_work_path.join("merged");
            // 正在运行中的需要 kill
            if container_info.status == ContainerStatus::RUNNING {
                kill(Pid::from_raw(container_info.pid), Signal::SIGTERM)?;
            }
            // 卸载自定义挂载点
            if container_info.volume != "" {
                container_info.volume
                    .split(",")
                    .filter_map(|v| v.split(":").last())
                    .map(|v| container_merged_path.join(v).to_string_lossy().to_string())
                    .for_each(|s| {
                        match umount(s.as_str()) {
                            Ok(s) => println!("卸载自定卷: {s:?}"),
                            Err(e) => panic!("卸载卷失败: {e:?}"),
                        }
                    });
            }
            // 卸载overlayfs
            match umount(container_merged_path.to_str().unwrap()) {
                Ok(s) => println!("卸载overlayfs卷: {s:?}"),
                Err(e) => println!("卸载overlayfs失败: {e:?}"),
            }

            println!("停止容器: {container_id:?}");
            // 删除容器目录
            if is_remove {
                match fs::remove_dir_all(container_work_path) {
                    Ok(s) => println!("删除容器 {container_id} 成功"),
                    Err(e) => println!("删除容器失败: {e:?}"),
                }
            }
        } else {
            println!("容器不存在: {container_id}")
        }
    }

    Ok(())
}


fn main() -> Result<()>{
    let args = RockerArgs::parse();
    if let (Some(cmd), Some(image_name)) = (&args.run, &args.image) {
        // run
        let volume_path = extend_image(image_name)?;
        let container_id = uuid::Uuid::new_v4().to_string()[0..8].to_string();
        let mut pid = -1;
        match run_container(&container_id,&cmd, &args, &volume_path) {
            Ok(child_pid) => {
                pid = child_pid;
            }
            Err(e) => {
                println!("run_container失败: {e}");
            }
        }        
        save_container_info(&args, &container_id, pid)?; // todo 无论出不错, 都要保存一个信息, 后面需要删除用清理
    } else if args.ps || args.psa {
        // --ps
        show_containers(args.psa)?
    } else if let Some(containers_id) = &args.rm {
        // --rm
        stop_container(containers_id, true)?;
    } else if let Some(containers_id) = &args.stop {
        // --stop
        stop_container(containers_id, false)?;
    }
    // } else if let Some(containers_id) = &args.start {
    //     // --start
    //     start_container(containers_id)?;
    // } else if let Some(containers_id) = &args.exec {
    //     // --exec
    //     exec_container(containers_id, &cmd, &args)?;
    // }

    // exec

    // logs


    Ok(())
}