use anyhow::{Context, Result};
use base64::{Engine as _, engine::general_purpose};
use clap::Parser;
use serde_json;
use std::env;
use std::fs;
use std::path::PathBuf;
use std::process::Command;
use std::time::Duration;
use tokio::time::sleep;

const DEFAULT_PROFILE_YAML: &str = r#"name: default

description: create foo user, docker

devices:
  eth0:
    nictype: bridged
    parent: lxdbr0
    type: nic

  root:
    path: /
    pool: default
    size: 50GiB
    type: disk

  # NOTE: Not mounting scripts directory since it won't be reachable when building the image

config:
  security.privileged: 'false'
  security.nesting: 'true'
  security.protection.shift: 'true'

  cloud-init.user-data: |-
    #cloud-config

    # Set the hostname
    hostname: lxd

    # Set the timezone
    timezone: Asia/Tokyo

    # Users configuration
    users:
      - name: foo
        gecos: Foo
        primary_group: foo
        groups: users,sudo
        sudo: ['ALL=(ALL) NOPASSWD:ALL']
        shell: /bin/bash
        lock_passwd: false

    # Package updates and upgrades
    package_update: true
    package_upgrade: true

    # Packages to install
    packages:
      - git
      - curl
      - vim
      - git-lfs
      - nano
      - gnupg2
      - gcc
      - g++
      - glances
      - net-tools
      - sshfs
      - nfs-common
      - sshpass
      - gh
      - python3-pip
      - cargo
      - s3fs
      - pre-commit

    # Custom scripts to run on first boot
    runcmd:
      - |
        # Read all environment variables from systemd (PID 1) and write them to /etc/environment
        # Exclude PATH as it's already set in /etc/environment
        cat /proc/1/environ | tr '\0' '\n' | grep -v '^PATH=' | grep '=' >> /etc/environment
      # Clone myawesomescripts repo and copy scripts from virtualization/lxd/
      - |
        REPO_URL="https://github.com/JeanMaximilienCadic/myawesomescripts.git"
        TEMP_DIR=$(mktemp -d)
        git clone "$REPO_URL" "$TEMP_DIR" 2>/dev/null || true
        if [ -d "$TEMP_DIR/virtualization/lxd" ]; then
          sudo cp "$TEMP_DIR/virtualization/lxd/sbin"/* /usr/local/sbin/ 2>/dev/null || true
          sudo chmod +x /usr/local/sbin/* 2>/dev/null || true
        fi
        rm -rf "$TEMP_DIR"
      # Copy skeleton .bashrc, add PATH and source helper scripts
      - cp /etc/skel/.bashrc /home/foo/.bashrc
      - echo "export PATH=$PATH:/home/foo/.cargo/bin" >> /home/foo/.bashrc
      - echo ". /usr/local/sbin/gh_auth" >> /home/foo/.bashrc
      - echo ". /usr/local/sbin/gh_ssh" >> /home/foo/.bashrc
      - echo ". /usr/local/sbin/mc_alias" >> /home/foo/.bashrc
      # MinIO Client (mc)
      - curl https://dl.min.io/client/mc/release/linux-amd64/mc -o /usr/local/bin/mc
      - chmod +x /usr/local/bin/mc
      # Ensure correct ownership
      - chown foo:foo -R /home/foo
      # Ensure the directories are created
      - mkdir -p /srv/code /srv/lake /FileStore /workspace
      - chown -R foo:foo /srv /FileStore /workspace
      - sed -i 's/#user_allow_other/user_allow_other/' /etc/fuse.conf  # Uncomment user_allow_other
      # Docker
      - curl -fsSL https://get.docker.com | sh
      - usermod -aG docker foo  # Add 'foo' user to the docker group instead of $USER
      # Task
      - curl -fsSL https://taskfile.dev/install.sh | sh
      - sudo mv task /usr/local/bin/  # Move "task" to the correct location
      # UV
      - export HOME=/root && curl -LsSf https://astral.sh/uv/install.sh | bash
      - mv /root/.local/bin/uv /usr/local/bin/

    # Disable SSH key regulation
    ssh_pwauth: True

    # Disable cloud-init logs
    output:
      all: "| tee -a /var/log/cloud-init-output.log"

project: default
"#;

#[derive(Parser, Debug)]
#[command(
    author,
    version,
    about = "LXD Container Management Tool",
    long_about = r#"
LXD Container Management Tool

This tool automates the setup and configuration of LXD containers with:
- Profile creation/update with cloud-init configuration
- Static IP assignment
- Environment variable configuration
- SSH key installation

USAGE EXAMPLES:

  # Show help (default):
  cargo run

  # Display default embedded profile configuration:
  cargo run -- show-config

  # Basic usage with default settings (requires lxd/.env file):
  cargo run -- build --env-file lxd/.env

  # Use existing LXD profile (skips env file validation):
  cargo run -- build --from-config my-existing-profile

  # Custom project and container:
  cargo run -- build --project myproject --container mycontainer --env-file lxd/.env

  # Custom bridge network:
  cargo run -- build --bridge mybridge --env-file lxd/.env

  # Use custom profile YAML file:
  cargo run -- build --profile-yaml /path/to/custom.yaml

  # Full example with all options:
  cargo run -- build \
    --project myproject \
    --container mycontainer \
    --container-ip 10.2.0.100 \
    --bridge lxdbr1 \
    --env-file /path/to/.env \
    --profile-yaml /path/to/profile.yaml

ENVIRONMENT VARIABLES:

  All command-line arguments can be set via environment variables:
  - PROJECT: LXD project name (default: "default")
  - CONTAINER: Container name (default: "default-dev")
  - CONTAINER_IP: Container IP address (default: "10.2.0.2")
  - BRIDGE: Network bridge name (default: "lxdbr0")
  - IMAGE: Container image (default: "ubuntu:24.04")
  - PROFILE_YAML: Path to custom profile YAML file
  - ENV_FILE: Path to environment variables file
  - FROM_CONFIG: Existing LXD profile name to reuse

REQUIREMENTS:

  When using the default embedded profile (--profile-yaml not specified):
  - The --env-file argument is REQUIRED
  - The environment file must exist and contain KEY=VALUE pairs

  When using a custom profile YAML (--profile-yaml specified):
  - The --env-file argument is optional
  - Environment variables can be configured in the YAML file itself
"#,
    subcommand_required = false,
    arg_required_else_help = true
)]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(clap::Subcommand, Debug)]
enum Commands {
    /// Build and configure the LXD container
    Build(BuildArgs),
    /// Display the default embedded profile configuration
    ShowConfig,
    /// Restart LXD network daemon and flush iptables
    RestartNetwork,
}

#[derive(Parser, Debug)]
struct BuildArgs {
    /// LXD project name
    #[arg(long, env = "PROJECT", default_value = "default")]
    project: String,

    /// Container name
    #[arg(long, env = "CONTAINER", default_value = "default-dev")]
    container: String,

    /// Container IP address
    #[arg(long, env = "CONTAINER_IP", default_value = "10.2.0.2")]
    container_ip: String,

    /// Network bridge name
    #[arg(long, env = "BRIDGE", default_value = "lxdbr0")]
    bridge: String,

    /// Container image (e.g., ubuntu:24.04, ubuntu:22.04)
    #[arg(long, env = "IMAGE", default_value = "ubuntu:24.04")]
    image: String,

    /// Path to profile YAML file (if not specified, uses embedded default)
    #[arg(long, env = "PROFILE_YAML")]
    profile_yaml: Option<PathBuf>,

    /// Path to environment file (required when using default config)
    #[arg(long, env = "ENV_FILE")]
    env_file: Option<PathBuf>,

    /// Use existing LXD profile name (skips env file and profile YAML validation)
    #[arg(long, env = "FROM_CONFIG")]
    from_config: Option<String>,
}

struct Config {
    project: String,
    container: String,
    container_ip: String,
    bridge: String,
    image: String,
    project_root: PathBuf,
    profile_yaml: Option<PathBuf>,
    env_file: Option<PathBuf>,
    from_config: Option<String>,
}

impl Config {
    fn new(args: BuildArgs) -> Result<Self> {
        let project_root = env::current_dir()?;

        Ok(Config {
            project: args.project,
            container: args.container,
            container_ip: args.container_ip,
            bridge: args.bridge,
            image: args.image,
            project_root,
            profile_yaml: args.profile_yaml,
            env_file: args.env_file,
            from_config: args.from_config,
        })
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Some(Commands::Build(args)) => {
            let config = Config::new(args)?;

            println!(
                "Using PROJECT={}, CONTAINER={}, CONTAINER_IP={}",
                config.project, config.container, config.container_ip
            );

            // Step 0: Validate env file (skip if using existing config or custom profile YAML)
            validate_env_file(&config)?;

            // Step 1: Delete container if it exists
            delete_container(&config)?;

            // Step 2: Update profile (skip if using existing profile)
            update_profile(&config)?;

            // Step 3: Initialize container
            init_container(&config)?;

            // Step 4: Set static IP
            set_static_ip(&config)?;

            // Step 5: Set credentials
            set_credentials(&config)?;

            // Step 6: Start container
            start_container(&config)?;

            // Step 7: Wait for container to be ready
            wait_container_ready().await?;

            // Step 8: Set SSH keys
            set_ssh_keys(&config)?;

            println!("Container setup completed successfully!");
            Ok(())
        }
        Some(Commands::ShowConfig) => {
            println!("{}", DEFAULT_PROFILE_YAML);
            Ok(())
        }
        Some(Commands::RestartNetwork) => {
            restart_network().await
        }
        None => {
            // This should not happen due to arg_required_else_help, but handle it anyway
            Ok(())
        }
    }
}

fn validate_env_file(config: &Config) -> Result<()> {
    // Skip validation if using existing LXD profile
    if config.from_config.is_some() {
        return Ok(());
    }

    // Skip validation if custom profile YAML is provided
    if config.profile_yaml.is_some() {
        return Ok(());
    }

    // Determine env file path: use specified path or default to lxd/.env
    let env_file = if let Some(ref env_path) = config.env_file {
        env_path.clone()
    } else {
        config.project_root.join("lxd/.env")
    };

    if !env_file.exists() {
        eprintln!(
            "❌ Error: Environment file not found at {:?}",
            env_file
        );
        eprintln!(
            "❌ Error: The default configuration cannot be fully executed without environment variables."
        );
        eprintln!(
            "❌ Error: Please specify the path to the environment file using --env-file or provide a custom YAML file using --profile-yaml."
        );
        anyhow::bail!("Environment file not found at {:?}", env_file);
    }

    Ok(())
}

fn delete_container(config: &Config) -> Result<()> {
    println!("Deleting container '{}' if it exists...", config.container);
    let output = Command::new("lxc")
        .args(&[
            "delete",
            &config.container,
            "--force",
            "--project",
            &config.project,
        ])
        .output();

    // Ignore errors - container might not exist
    match output {
        Ok(_) => println!("Container deletion attempted (may not have existed)"),
        Err(_) => println!("Container deletion skipped (container may not exist)"),
    }
    Ok(())
}

fn update_profile(config: &Config) -> Result<()> {
    // If using existing profile, just verify it exists
    if let Some(ref profile_name) = config.from_config {
        println!("Using existing LXD profile '{}' for project: {}", profile_name, config.project);
        let check_output = Command::new("lxc")
            .args(&["profile", "show", profile_name, "--project", &config.project])
            .output()
            .context("Failed to check if profile exists")?;

        if !check_output.status.success() {
            anyhow::bail!(
                "Profile '{}' does not exist in project '{}'",
                profile_name,
                config.project
            );
        }
        println!(
            "Profile '{}' verified in project '{}'",
            profile_name,
            config.project
        );
        return Ok(());
    }

    // Otherwise, update or create the default profile
    let profile_name = "default";
    println!("Updating default profile for project: {}", config.project);

    // Check if profile exists
    let check_output = Command::new("lxc")
        .args(&["profile", "show", profile_name, "--project", &config.project])
        .output()
        .context("Failed to check if profile exists")?;

    // Get profile content: use file if specified, otherwise use embedded default
    let mut profile_content = if let Some(ref yaml_path) = config.profile_yaml {
        println!("Using profile YAML from: {:?}", yaml_path);
        if !yaml_path.exists() {
            anyhow::bail!("Profile YAML file not found at {:?}", yaml_path);
        }
        fs::read_to_string(yaml_path)
            .with_context(|| format!("Failed to read profile YAML: {:?}", yaml_path))?
    } else {
        println!("Using embedded default profile YAML");
        DEFAULT_PROFILE_YAML.to_string()
    };

    // Replace the project name and bridge in the YAML template
    profile_content = profile_content.replace("project: default", &format!("project: {}", config.project))
        .replace("project: default", &format!("project: {}", config.project))
        .replace("parent: lxdbr0", &format!("parent: {}", config.bridge));

    if check_output.status.success() {
        println!("Profile '{}' exists in project '{}', updating...", profile_name, config.project);
        let mut child = Command::new("lxc")
            .args(&["profile", "edit", profile_name, "--project", &config.project])
            .stdin(std::process::Stdio::piped())
            .spawn()
            .context("Failed to spawn lxc profile edit")?;

        use std::io::Write;
        child
            .stdin
            .as_mut()
            .unwrap()
            .write_all(profile_content.as_bytes())
            .context("Failed to write profile content")?;

        child.wait().context("Failed to wait for lxc profile edit")?;
    } else {
        println!(
            "Profile '{}' does not exist in project '{}', creating...",
            profile_name,
            config.project
        );
        let mut child = Command::new("lxc")
            .args(&["profile", "create", profile_name, "--project", &config.project])
            .stdin(std::process::Stdio::piped())
            .spawn()
            .context("Failed to spawn lxc profile create")?;

        use std::io::Write;
        child
            .stdin
            .as_mut()
            .unwrap()
            .write_all(profile_content.as_bytes())
            .context("Failed to write profile content")?;

        child.wait().context("Failed to wait for lxc profile create")?;
    }

    println!(
        "Profile '{}' updated successfully for project '{}'",
        profile_name,
        config.project
    );
    Ok(())
}

fn init_container(config: &Config) -> Result<()> {
    println!("Initializing container '{}' with image '{}'...", config.container, config.image);
    let profile_name = config.from_config.as_deref().unwrap_or("default");
    Command::new("lxc")
        .args(&[
            "init",
            &config.image,
            &config.container,
            "-p",
            profile_name,
            "--project",
            &config.project,
        ])
        .status()
        .context("Failed to initialize container")?;
    println!("Container initialized successfully");
    Ok(())
}

fn set_static_ip(config: &Config) -> Result<()> {
    println!("Setting static IP address: {}", config.container_ip);
    // Remove existing eth0 device if it exists (ignore errors)
    let _ = Command::new("lxc")
        .args(&[
            "config",
            "device",
            "remove",
            &config.container,
            "eth0",
            "--project",
            &config.project,
        ])
        .output();

    // Add new eth0 device with static IP
    Command::new("lxc")
        .args(&[
            "config",
            "device",
            "add",
            &config.container,
            "eth0",
            "nic",
            "nictype=bridged",
            &format!("parent={}", config.bridge),
            &format!("ipv4.address={}", config.container_ip),
            "--project",
            &config.project,
        ])
        .status()
        .context("Failed to set static IP")?;
    println!("Static IP set successfully");
    Ok(())
}

fn set_credentials(config: &Config) -> Result<()> {
    println!("Setting credentials for container: {}", config.container);
    
    // Skip if using existing config file
    if config.from_config.is_some() {
        println!("Skipping environment variable setup (using existing config)");
        return Ok(());
    }

    // Skip if custom profile YAML is provided (env vars might be in YAML)
    if config.profile_yaml.is_some() {
        println!("Skipping environment variable setup (using custom profile YAML)");
        return Ok(());
    }

    // Determine env file path: use specified path or default to lxd/.env
    let env_file = if let Some(ref env_path) = config.env_file {
        env_path.clone()
    } else {
        config.project_root.join("lxd/.env")
    };

    // This should not happen if validate_env_file was called, but check anyway
    if !env_file.exists() {
        eprintln!(
            "❌ Error: Environment file not found at {:?}",
            env_file
        );
        eprintln!(
            "❌ Error: The default configuration cannot be fully executed without environment variables."
        );
        eprintln!(
            "❌ Error: Please specify the path to the environment file using --env-file or provide a custom YAML file using --profile-yaml."
        );
        anyhow::bail!("Environment file not found at {:?}", env_file);
    }

    // Use the project from config
    let project = &config.project;
    println!("Setting credentials for container '{}' in project '{}'", config.container, project);

    // Read and parse .env file
    let env_content = fs::read_to_string(&env_file)
        .with_context(|| format!("Failed to read .env file: {:?}", env_file))?;

    let mut var_count = 0;
    for line in env_content.lines() {
        // Skip empty lines and comments
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }

        // Skip lines without '='
        if !line.contains('=') {
            continue;
        }

        // Extract variable name and value
        if let Some(equal_pos) = line.find('=') {
            let var_name = line[..equal_pos].trim();
            let var_value = line[equal_pos + 1..].trim();

            if var_name.is_empty() {
                continue;
            }

            println!("Setting {}...", var_name);
            Command::new("lxc")
                .args(&[
                    "config",
                    "set",
                    &config.container,
                    &format!("environment.{}", var_name),
                    var_value,
                    "--project",
                    &project,
                ])
                .status()
                .with_context(|| format!("Failed to set environment variable: {}", var_name))?;
            var_count += 1;
        }
    }

    if var_count == 0 {
        eprintln!(
            "⚠️  Warning: No valid environment variables found in {:?}",
            env_file
        );
        return Ok(());
    }

    println!("Successfully set {} environment variable(s)!", var_count);
    Ok(())
}

fn start_container(config: &Config) -> Result<()> {
    println!("Starting container '{}'...", config.container);
    Command::new("lxc")
        .args(&["start", &config.container, "--project", &config.project])
        .status()
        .context("Failed to start container")?;
    println!("Container started successfully");
    Ok(())
}

async fn wait_container_ready() -> Result<()> {
    println!("Waiting for container to be ready...");
    sleep(Duration::from_secs(5)).await;
    Ok(())
}

fn set_ssh_keys(config: &Config) -> Result<()> {
    println!("Setting SSH keys for container: {}", config.container);
    let auth_keys_file = config.project_root.join("lxd/authorized_keys");
    let ssh_user = "foo";

    if !auth_keys_file.exists() {
        println!(
            "Warning: authorized_keys file not found at {:?}; skipping SSH key installation.",
            auth_keys_file
        );
        return Ok(());
    }

    // Use the project from config
    let project = &config.project;
    println!("Setting SSH keys for container '{}' in project '{}'", config.container, project);

    // Read SSH keys from file
    let keys_content = fs::read_to_string(&auth_keys_file)
        .with_context(|| format!("Failed to read authorized_keys file: {:?}", auth_keys_file))?;

    // Filter out comments and empty lines
    let ssh_keys: Vec<&str> = keys_content
        .lines()
        .map(|l| l.trim())
        .filter(|l| !l.is_empty() && !l.starts_with('#'))
        .collect();

    if ssh_keys.is_empty() {
        println!(
            "Warning: No SSH keys found in {:?}; skipping SSH key installation.",
            auth_keys_file
        );
        return Ok(());
    }

    println!(
        "Ensuring SSH authorized keys for user '{}'...",
        ssh_user
    );
    println!("Found {} key(s) to install", ssh_keys.len());

    // Check if container is running
    let status_output = Command::new("lxc")
        .args(&[
            "list",
            &config.container,
            "--project",
            &project,
            "--format",
            "json",
        ])
        .output()
        .context("Failed to check container status")?;

    let status_json: serde_json::Value = serde_json::from_slice(&status_output.stdout)
        .context("Failed to parse container status JSON")?;

    let container_status = status_json
        .as_array()
        .and_then(|arr| arr.get(0))
        .and_then(|c| c.get("status"))
        .and_then(|s| s.as_str())
        .unwrap_or("");

    if container_status != "Running" {
        anyhow::bail!(
            "Container '{}' is not running. Start the container first.",
            config.container
        );
    }

    // Create .ssh directory
    Command::new("lxc")
        .args(&[
            "exec",
            &config.container,
            "--project",
            &project,
            "--",
            "bash",
            "-c",
            &format!(
                "mkdir -p /home/{}/.ssh && chmod 700 /home/{}/.ssh && chown {}:{} /home/{}/.ssh",
                ssh_user, ssh_user, ssh_user, ssh_user, ssh_user
            ),
        ])
        .status()
        .context("Failed to create .ssh directory")?;

    // Write keys to authorized_keys file
    let keys_content_clean = ssh_keys.join("\n");
    let keys_b64 = general_purpose::STANDARD.encode(keys_content_clean.as_bytes());

    let script = format!(
        r#"
echo '{}' | base64 -d > /tmp/authorized_keys_new
if [ -f /home/{}/.ssh/authorized_keys ]; then
    cat /home/{}/.ssh/authorized_keys /tmp/authorized_keys_new | sort -u > /tmp/authorized_keys_merged
else
    cp /tmp/authorized_keys_new /tmp/authorized_keys_merged
fi
mv /tmp/authorized_keys_merged /home/{}/.ssh/authorized_keys
chmod 600 /home/{}/.ssh/authorized_keys
chown {}:{} /home/{}/.ssh/authorized_keys
rm -f /tmp/authorized_keys_new /tmp/authorized_keys_merged
"#,
        keys_b64, ssh_user, ssh_user, ssh_user, ssh_user, ssh_user, ssh_user, ssh_user
    );

    Command::new("lxc")
        .args(&[
            "exec",
            &config.container,
            "--project",
            &project,
            "--",
            "bash",
            "-c",
            &script,
        ])
        .status()
        .context("Failed to set SSH keys")?;

    println!("SSH authorized keys ensured for {}.", ssh_user);
    Ok(())
}

async fn restart_network() -> Result<()> {
    println!("Restarting LXD network daemon and flushing iptables...");

    // Flush iptables for all variants
    println!("Flushing iptables...");
    let iptables_variants = ["iptables", "iptables-legacy", "ip6tables", "ip6tables-legacy"];
    
    for ipt in &iptables_variants {
        let commands = vec![
            format!("{} --flush", ipt),
            format!("{} --flush -t nat", ipt),
            format!("{} --delete-chain", ipt),
            format!("{} --delete-chain -t nat", ipt),
            format!("{} -P FORWARD ACCEPT", ipt),
            format!("{} -P INPUT ACCEPT", ipt),
            format!("{} -P OUTPUT ACCEPT", ipt),
        ];

        for cmd in commands {
            // Ignore errors for iptables commands (they might fail if rules don't exist)
            let _ = Command::new("bash")
                .args(&["-c", &cmd])
                .output();
        }
    }

    // Reload snap.lxd.daemon
    println!("Reloading snap.lxd.daemon...");
    Command::new("systemctl")
        .args(&["reload", "snap.lxd.daemon"])
        .status()
        .context("Failed to reload snap.lxd.daemon")?;

    // Restart snap lxd.daemon
    println!("Restarting snap lxd.daemon...");
    Command::new("snap")
        .args(&["restart", "lxd.daemon"])
        .status()
        .context("Failed to restart lxd.daemon")?;

    println!("LXD network daemon restarted successfully!");
    Ok(())
}


