use actix_web::{App, HttpResponse, HttpServer, web};
use clap::Parser;
use ipnetwork::IpNetwork;
use rlua::{Function, Lua};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::env;
use std::fs;
use std::io::Write;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::process::Command;
use tokio::sync::RwLock;
use tokio::time::sleep;

mod markov;
mod metrics;

use markov::MarkovGenerator;
use metrics::{
    ACTIVE_CONNECTIONS, BLOCKED_IPS, HITS_COUNTER, PATH_HITS, UA_HITS, metrics_handler,
    status_handler,
};

// Command-line arguments using clap
#[derive(Parser, Debug, Clone)]
#[clap(
    author,
    version,
    about = "Markov chain based HTTP tarpit/honeypot that delays and tracks potential attackers"
)]
struct Args {
    #[clap(
        long,
        default_value = "0.0.0.0:8888",
        help = "Address and port to listen for incoming HTTP requests (format: ip:port)"
    )]
    listen_addr: String,

    #[clap(
        long,
        default_value = "0.0.0.0:9100",
        help = "Address and port to expose Prometheus metrics and status endpoint (format: ip:port)"
    )]
    metrics_addr: String,

    #[clap(long, help = "Disable Prometheus metrics server completely")]
    disable_metrics: bool,

    #[clap(
        long,
        default_value = "127.0.0.1:80",
        help = "Backend server address to proxy legitimate requests to (format: ip:port)"
    )]
    backend_addr: String,

    #[clap(
        long,
        default_value = "1000",
        help = "Minimum delay in milliseconds between chunks sent to attacker"
    )]
    min_delay: u64,

    #[clap(
        long,
        default_value = "15000",
        help = "Maximum delay in milliseconds between chunks sent to attacker"
    )]
    max_delay: u64,

    #[clap(
        long,
        default_value = "600",
        help = "Maximum time in seconds to keep an attacker in the tarpit before disconnecting"
    )]
    max_tarpit_time: u64,

    #[clap(
        long,
        default_value = "3",
        help = "Number of hits to honeypot patterns before permanently blocking an IP"
    )]
    block_threshold: u32,

    #[clap(
        long,
        help = "Base directory for all application data (overrides XDG directory structure)"
    )]
    base_dir: Option<PathBuf>,

    #[clap(
        long,
        help = "Path to JSON configuration file (overrides command line options)"
    )]
    config_file: Option<PathBuf>,

    #[clap(
        long,
        default_value = "info",
        help = "Log level: trace, debug, info, warn, error"
    )]
    log_level: String,
}

// Configuration structure
#[derive(Clone, Debug, Deserialize, Serialize)]
struct Config {
    listen_addr: String,
    metrics_addr: String,
    disable_metrics: bool,
    backend_addr: String,
    min_delay: u64,
    max_delay: u64,
    max_tarpit_time: u64,
    block_threshold: u32,
    trap_patterns: Vec<String>,
    whitelist_networks: Vec<String>,
    markov_corpora_dir: String,
    lua_scripts_dir: String,
    data_dir: String,
    config_dir: String,
    cache_dir: String,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8888".to_string(),
            metrics_addr: "0.0.0.0:9100".to_string(),
            disable_metrics: false,
            backend_addr: "127.0.0.1:80".to_string(),
            min_delay: 1000,
            max_delay: 15000,
            max_tarpit_time: 600,
            block_threshold: 3,
            trap_patterns: vec![
                "/vendor/phpunit".to_string(),
                "eval-stdin.php".to_string(),
                "/wp-admin".to_string(),
                "/wp-login.php".to_string(),
                "/xmlrpc.php".to_string(),
                "/phpMyAdmin".to_string(),
                "/solr/".to_string(),
                "/.env".to_string(),
                "/config".to_string(),
                "/api/".to_string(),
                "/actuator/".to_string(),
            ],
            whitelist_networks: vec![
                "192.168.0.0/16".to_string(),
                "10.0.0.0/8".to_string(),
                "172.16.0.0/12".to_string(),
                "127.0.0.0/8".to_string(),
            ],
            markov_corpora_dir: "./corpora".to_string(),
            lua_scripts_dir: "./scripts".to_string(),
            data_dir: "./data".to_string(),
            config_dir: "./conf".to_string(),
            cache_dir: "./cache".to_string(),
        }
    }
}

// Gets standard XDG directory paths for config, data and cache
fn get_xdg_dirs() -> (PathBuf, PathBuf, PathBuf) {
    let config_home = env::var_os("XDG_CONFIG_HOME")
        .map(PathBuf::from)
        .unwrap_or_else(|| {
            let home = env::var_os("HOME").map_or_else(|| PathBuf::from("."), PathBuf::from);
            home.join(".config")
        });

    let data_home = env::var_os("XDG_DATA_HOME")
        .map(PathBuf::from)
        .unwrap_or_else(|| {
            let home = env::var_os("HOME").map_or_else(|| PathBuf::from("."), PathBuf::from);
            home.join(".local").join("share")
        });

    let cache_home = env::var_os("XDG_CACHE_HOME")
        .map(PathBuf::from)
        .unwrap_or_else(|| {
            let home = env::var_os("HOME").map_or_else(|| PathBuf::from("."), PathBuf::from);
            home.join(".cache")
        });

    let config_dir = config_home.join("eris");
    let data_dir = data_home.join("eris");
    let cache_dir = cache_home.join("eris");

    (config_dir, data_dir, cache_dir)
}

impl Config {
    // Create configuration from command-line args
    fn from_args(args: &Args) -> Self {
        let (config_dir, data_dir, cache_dir) = if let Some(base_dir) = &args.base_dir {
            let base_str = base_dir.to_string_lossy().to_string();
            (
                format!("{base_str}/conf"),
                format!("{base_str}/data"),
                format!("{base_str}/cache"),
            )
        } else {
            let (c, d, cache) = get_xdg_dirs();
            (
                c.to_string_lossy().to_string(),
                d.to_string_lossy().to_string(),
                cache.to_string_lossy().to_string(),
            )
        };

        Self {
            listen_addr: args.listen_addr.clone(),
            metrics_addr: args.metrics_addr.clone(),
            disable_metrics: args.disable_metrics,
            backend_addr: args.backend_addr.clone(),
            min_delay: args.min_delay,
            max_delay: args.max_delay,
            max_tarpit_time: args.max_tarpit_time,
            block_threshold: args.block_threshold,
            markov_corpora_dir: format!("{data_dir}/corpora"),
            lua_scripts_dir: format!("{data_dir}/scripts"),
            data_dir,
            config_dir,
            cache_dir,
            ..Default::default()
        }
    }

    // Load configuration from a JSON file
    fn load_from_file(path: &Path) -> std::io::Result<Self> {
        let content = fs::read_to_string(path)?;
        let config = serde_json::from_str(&content)?;
        Ok(config)
    }

    // Save configuration to a JSON file
    fn save_to_file(&self, path: &Path) -> std::io::Result<()> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let content = serde_json::to_string_pretty(self)?;
        fs::write(path, content)?;
        Ok(())
    }

    // Create required directories if they don't exist
    fn ensure_dirs_exist(&self) -> std::io::Result<()> {
        let dirs = [
            &self.markov_corpora_dir,
            &self.lua_scripts_dir,
            &self.data_dir,
            &self.config_dir,
            &self.cache_dir,
        ];

        for dir in dirs {
            fs::create_dir_all(dir)?;
            log::debug!("Created directory: {dir}");
        }

        Ok(())
    }
}

// State of bots/IPs hitting the honeypot
#[derive(Clone, Debug)]
struct BotState {
    hits: HashMap<IpAddr, u32>,
    blocked: HashSet<IpAddr>,
    active_connections: HashSet<IpAddr>,
    data_dir: String,
    cache_dir: String,
}

impl BotState {
    fn new(data_dir: &str, cache_dir: &str) -> Self {
        Self {
            hits: HashMap::new(),
            blocked: HashSet::new(),
            active_connections: HashSet::new(),
            data_dir: data_dir.to_string(),
            cache_dir: cache_dir.to_string(),
        }
    }

    // Load previous state from disk
    fn load_from_disk(data_dir: &str, cache_dir: &str) -> Self {
        let mut state = Self::new(data_dir, cache_dir);
        let blocked_ips_file = format!("{data_dir}/blocked_ips.txt");

        if let Ok(content) = fs::read_to_string(&blocked_ips_file) {
            let mut loaded = 0;
            for line in content.lines() {
                if let Ok(ip) = line.parse::<IpAddr>() {
                    state.blocked.insert(ip);
                    loaded += 1;
                }
            }
            log::info!("Loaded {loaded} blocked IPs from {blocked_ips_file}");
        } else {
            log::info!("No blocked IPs file found at {blocked_ips_file}");
        }

        // Check for temporary hit counter cache
        let hit_cache_file = format!("{cache_dir}/hit_counters.json");
        if let Ok(content) = fs::read_to_string(&hit_cache_file) {
            if let Ok(hit_map) = serde_json::from_str::<HashMap<String, u32>>(&content) {
                for (ip_str, count) in hit_map {
                    if let Ok(ip) = ip_str.parse::<IpAddr>() {
                        state.hits.insert(ip, count);
                    }
                }
                log::info!("Loaded hit counters for {} IPs", state.hits.len());
            }
        }

        BLOCKED_IPS.set(state.blocked.len() as f64);
        state
    }

    // Persist state to disk for later reloading
    fn save_to_disk(&self) {
        // Save blocked IPs
        if let Err(e) = fs::create_dir_all(&self.data_dir) {
            log::error!("Failed to create data directory: {e}");
            return;
        }

        let blocked_ips_file = format!("{}/blocked_ips.txt", self.data_dir);

        match fs::File::create(&blocked_ips_file) {
            Ok(mut file) => {
                let mut count = 0;
                for ip in &self.blocked {
                    if writeln!(file, "{ip}").is_ok() {
                        count += 1;
                    }
                }
                log::info!("Saved {count} blocked IPs to {blocked_ips_file}");
            }
            Err(e) => {
                log::error!("Failed to create blocked IPs file: {e}");
            }
        }

        // Save hit counters to cache
        if let Err(e) = fs::create_dir_all(&self.cache_dir) {
            log::error!("Failed to create cache directory: {e}");
            return;
        }

        let hit_cache_file = format!("{}/hit_counters.json", self.cache_dir);
        let mut hit_map = HashMap::new();
        for (ip, count) in &self.hits {
            hit_map.insert(ip.to_string(), *count);
        }

        match fs::File::create(&hit_cache_file) {
            Ok(file) => {
                if let Err(e) = serde_json::to_writer(file, &hit_map) {
                    log::error!("Failed to write hit counters to cache: {e}");
                } else {
                    log::debug!("Saved hit counters for {} IPs to cache", hit_map.len());
                }
            }
            Err(e) => {
                log::error!("Failed to create hit counter cache file: {e}");
            }
        }
    }
}

// Lua scripts for response generation and customization
struct ScriptManager {
    script_content: String,
    scripts_loaded: bool,
}

impl ScriptManager {
    fn new(scripts_dir: &str) -> Self {
        let mut script_content = String::new();
        let mut scripts_loaded = false;

        // Try to load scripts from directory
        let script_dir = Path::new(scripts_dir);
        if script_dir.exists() {
            log::debug!("Loading Lua scripts from directory: {scripts_dir}");
            if let Ok(entries) = fs::read_dir(script_dir) {
                for entry in entries {
                    if let Ok(entry) = entry {
                        let path = entry.path();
                        if path.extension().and_then(|ext| ext.to_str()) == Some("lua") {
                            if let Ok(content) = fs::read_to_string(&path) {
                                log::debug!("Loaded Lua script: {}", path.display());
                                script_content.push_str(&content);
                                script_content.push('\n');
                                scripts_loaded = true;
                            } else {
                                log::warn!("Failed to read Lua script: {}", path.display());
                            }
                        }
                    }
                }
            }
        } else {
            log::warn!("Lua scripts directory does not exist: {scripts_dir}");
        }

        // If no scripts were loaded, use a default script
        if !scripts_loaded {
            log::info!("No Lua scripts found, loading default scripts");
            script_content = r#"
            function generate_honeytoken(token)
                local token_types = {"API_KEY", "AUTH_TOKEN", "SESSION_ID", "SECRET_KEY"}
                local prefix = token_types[math.random(#token_types)]
                local suffix = string.format("%08x", math.random(0xffffff))
                return prefix .. "_" .. token .. "_" .. suffix
            end

            function enhance_response(text, response_type, path, token)
                local result = text
                local honeytoken = generate_honeytoken(token)

                -- Add some fake sensitive data
                result = result .. "\n<!-- DEBUG: " .. honeytoken .. " -->"
                result = result .. "\n<div style='display:none'>Server ID: " .. token .. "</div>"

                return result
            end
            "#
            .to_string();
            scripts_loaded = true;
        }

        Self {
            script_content,
            scripts_loaded,
        }
    }

    // Lua is a powerful configuration language we can use to expand functionality of
    // Eris, e.g., with fake tokens or honeytrap content.
    fn expand_response(&self, text: &str, response_type: &str, path: &str, token: &str) -> String {
        if !self.scripts_loaded {
            return format!("{text}\n<!-- Token: {token} -->");
        }

        let lua = Lua::new();
        if let Err(e) = lua.load(&self.script_content).exec() {
            log::warn!("Error loading Lua script: {e}");
            return format!("{text}\n<!-- Error: Failed to load Lua script -->");
        }

        let globals = lua.globals();
        match globals.get::<_, Function>("enhance_response") {
            Ok(enhance_func) => {
                match enhance_func.call::<_, String>((text, response_type, path, token)) {
                    Ok(result) => result,
                    Err(e) => {
                        log::warn!("Error calling Lua function enhance_response: {e}");
                        format!("{text}\n<!-- Error calling Lua enhance_response -->")
                    }
                }
            }
            Err(e) => {
                log::warn!("Lua enhance_response function not found: {e}");
                format!("{text}\n<!-- Lua enhance_response function not found -->")
            }
        }
    }
}

// Main connection handler - decides whether to tarpit or proxy
async fn handle_connection(
    mut stream: TcpStream,
    config: Arc<Config>,
    state: Arc<RwLock<BotState>>,
    markov_generator: Arc<MarkovGenerator>,
    script_manager: Arc<ScriptManager>,
) {
    let peer_addr = match stream.peer_addr() {
        Ok(addr) => addr.ip(),
        Err(e) => {
            log::debug!("Failed to get peer address: {e}");
            return;
        }
    };

    log::debug!("New connection from: {peer_addr}");

    // Check if IP is already blocked
    if state.read().await.blocked.contains(&peer_addr) {
        log::debug!("Rejected connection from blocked IP: {peer_addr}");
        let _ = stream.shutdown().await;
        return;
    }

    // Read the HTTP request
    let mut buffer = [0; 8192];
    let mut request_data = Vec::new();

    // Read with timeout to prevent hanging
    let read_fut = async {
        loop {
            match stream.read(&mut buffer).await {
                Ok(0) => break,
                Ok(n) => {
                    request_data.extend_from_slice(&buffer[..n]);
                    // Stop reading at empty line, this is the end of HTTP headers
                    if request_data.len() > 2 && &request_data[request_data.len() - 2..] == b"\r\n"
                    {
                        break;
                    }
                }
                Err(e) => {
                    log::debug!("Error reading from stream: {e}");
                    break;
                }
            }
        }
    };

    let timeout_fut = sleep(Duration::from_secs(5));

    tokio::select! {
        () = read_fut => {},
        () = timeout_fut => {
            log::debug!("Connection timeout from: {peer_addr}");
            let _ = stream.shutdown().await;
            return;
        }
    }

    // Parse the request
    let request_str = String::from_utf8_lossy(&request_data);
    let request_lines: Vec<&str> = request_str.lines().collect();

    if request_lines.is_empty() {
        log::debug!("Empty request from: {peer_addr}");
        let _ = stream.shutdown().await;
        return;
    }

    // Parse request line
    let request_parts: Vec<&str> = request_lines[0].split_whitespace().collect();
    if request_parts.len() < 3 {
        log::debug!("Malformed request from {}: {}", peer_addr, request_lines[0]);
        let _ = stream.shutdown().await;
        return;
    }

    let method = request_parts[0];
    let path = request_parts[1];
    let protocol = request_parts[2];

    log::debug!("Request: {method} {path} {protocol} from {peer_addr}");

    // Parse headers
    let mut headers = HashMap::new();
    for line in &request_lines[1..] {
        if line.is_empty() {
            break;
        }

        if let Some(idx) = line.find(':') {
            let key = line[..idx].trim();
            let value = line[idx + 1..].trim();
            headers.insert(key, value.to_string());
        }
    }

    let user_agent = headers
        .get("user-agent")
        .cloned()
        .unwrap_or_else(|| "unknown".to_string());

    // Check if this request matches our tarpit patterns
    let should_tarpit = should_tarpit(path, &peer_addr, &config).await;

    if should_tarpit {
        log::info!("Tarpit triggered: {method} {path} from {peer_addr} (UA: {user_agent})");

        // Update metrics
        HITS_COUNTER.inc();
        PATH_HITS.with_label_values(&[path]).inc();
        UA_HITS.with_label_values(&[&user_agent]).inc();

        // Update state and check for blocking threshold
        {
            let mut state = state.write().await;
            state.active_connections.insert(peer_addr);
            ACTIVE_CONNECTIONS.set(state.active_connections.len() as f64);

            *state.hits.entry(peer_addr).or_insert(0) += 1;
            let hit_count = state.hits[&peer_addr];
            log::debug!("Hit count for {peer_addr}: {hit_count}");

            // Block IPs that hit tarpits too many times
            if hit_count >= config.block_threshold && !state.blocked.contains(&peer_addr) {
                log::info!("Blocking IP {peer_addr} after {hit_count} hits");
                state.blocked.insert(peer_addr);
                BLOCKED_IPS.set(state.blocked.len() as f64);
                state.save_to_disk();

                // Try to add to firewall
                let peer_addr_str = peer_addr.to_string();
                tokio::spawn(async move {
                    log::debug!("Adding IP {peer_addr_str} to firewall blacklist");
                    match Command::new("nft")
                        .args([
                            "add",
                            "element",
                            "inet",
                            "filter",
                            "eris_blacklist",
                            "{",
                            &peer_addr_str,
                            "}",
                        ])
                        .output()
                        .await
                    {
                        Ok(output) => {
                            if !output.status.success() {
                                log::warn!(
                                    "Failed to add IP {} to firewall: {}",
                                    peer_addr_str,
                                    String::from_utf8_lossy(&output.stderr)
                                );
                            }
                        }
                        Err(e) => {
                            log::warn!("Failed to execute nft command: {e}");
                        }
                    }
                });
            }
        }

        // Generate a deceptive response using Markov chains and Lua
        let response =
            generate_deceptive_response(path, &user_agent, &markov_generator, &script_manager)
                .await;

        // Send the response with the tarpit delay strategy
        tarpit_connection(
            stream,
            response,
            peer_addr,
            state.clone(),
            config.min_delay,
            config.max_delay,
            config.max_tarpit_time,
        )
        .await;
    } else {
        log::debug!("Proxying request: {method} {path} from {peer_addr}");

        // Proxy non-matching requests to the actual backend
        proxy_to_backend(
            stream,
            method,
            path,
            protocol,
            &headers,
            &config.backend_addr,
        )
        .await;
    }
}

// Determine if a request should be tarpitted based on path and IP
async fn should_tarpit(path: &str, ip: &IpAddr, config: &Config) -> bool {
    // Don't tarpit whitelisted IPs (internal networks, etc)
    for network_str in &config.whitelist_networks {
        if let Ok(network) = network_str.parse::<IpNetwork>() {
            if network.contains(*ip) {
                log::debug!("IP {ip} is in whitelist network {network_str}");
                return false;
            }
        }
    }

    // Check if the request path matches any of our trap patterns
    for pattern in &config.trap_patterns {
        if path.contains(pattern) {
            log::debug!("Path '{path}' matches trap pattern '{pattern}'");
            return true;
        }
    }

    // No trap patterns matched
    false
}

// Generate a deceptive HTTP response that appears legitimate
async fn generate_deceptive_response(
    path: &str,
    user_agent: &str,
    markov: &MarkovGenerator,
    script_manager: &ScriptManager,
) -> String {
    // Choose response type based on path to seem more realistic
    let response_type = if path.contains("phpunit") || path.contains("eval") {
        "php_exploit"
    } else if path.contains("wp-") {
        "wordpress"
    } else if path.contains("api") {
        "api"
    } else {
        "generic"
    };

    log::debug!("Generating {response_type} response for path: {path}");

    // Generate tracking token for this interaction
    let tracking_token = format!(
        "BOT_{}_{}",
        user_agent
            .chars()
            .filter(|c| c.is_alphanumeric())
            .collect::<String>(),
        chrono::Utc::now().timestamp()
    );

    // Generate base response using Markov chain text generator
    let markov_text = markov.generate(response_type, 30);

    // Use Lua to enhance with honeytokens and other deceptive content
    let enhanced =
        script_manager.expand_response(&markov_text, response_type, path, &tracking_token);

    // Return full HTTP response with appropriate headers
    format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nX-Powered-By: PHP/7.4.3\r\nConnection: keep-alive\r\n\r\n{enhanced}"
    )
}

// Slowly feed a response to the client with random delays to waste attacker time
async fn tarpit_connection(
    mut stream: TcpStream,
    response: String,
    peer_addr: IpAddr,
    state: Arc<RwLock<BotState>>,
    min_delay: u64,
    max_delay: u64,
    max_tarpit_time: u64,
) {
    let start_time = Instant::now();
    let mut chars = response.chars().collect::<Vec<_>>();

    // Randomize the char order slightly to confuse automated tools
    for i in (0..chars.len()).rev() {
        if i > 0 && rand::random::<f32>() < 0.1 {
            chars.swap(i, i - 1);
        }
    }

    log::debug!(
        "Starting tarpit for {} with {} chars, min_delay={}ms, max_delay={}ms",
        peer_addr,
        chars.len(),
        min_delay,
        max_delay
    );

    let mut position = 0;
    let mut chunks_sent = 0;
    let mut total_delay = 0;

    // Send the response character by character with random delays
    while position < chars.len() {
        // Check if we've exceeded maximum tarpit time
        let elapsed_secs = start_time.elapsed().as_secs();
        if elapsed_secs > max_tarpit_time {
            log::info!("Tarpit maximum time ({max_tarpit_time} sec) reached for {peer_addr}");
            break;
        }

        // Decide how many chars to send in this chunk (usually 1, sometimes more)
        let chunk_size = if rand::random::<f32>() < 0.9 {
            1
        } else {
            (rand::random::<f32>() * 3.0).floor() as usize + 1
        };

        let end = (position + chunk_size).min(chars.len());
        let chunk: String = chars[position..end].iter().collect();

        // Try to write chunk
        if stream.write_all(chunk.as_bytes()).await.is_err() {
            log::debug!("Connection closed by client during tarpit: {peer_addr}");
            break;
        }

        if stream.flush().await.is_err() {
            log::debug!("Failed to flush stream during tarpit: {peer_addr}");
            break;
        }

        position = end;
        chunks_sent += 1;

        // Apply random delay between min and max configured values
        let delay_ms = (rand::random::<f32>() * (max_delay - min_delay) as f32) as u64 + min_delay;
        total_delay += delay_ms;
        sleep(Duration::from_millis(delay_ms)).await;
    }

    log::debug!(
        "Tarpit stats for {}: sent {} chunks, {}% of data, total delay {}ms over {}s",
        peer_addr,
        chunks_sent,
        position * 100 / chars.len(),
        total_delay,
        start_time.elapsed().as_secs()
    );

    // Remove from active connections
    if let Ok(mut state) = state.try_write() {
        state.active_connections.remove(&peer_addr);
        ACTIVE_CONNECTIONS.set(state.active_connections.len() as f64);
    }

    let _ = stream.shutdown().await;
}

// Forward a legitimate request to the real backend server
async fn proxy_to_backend(
    mut client_stream: TcpStream,
    method: &str,
    path: &str,
    protocol: &str,
    headers: &HashMap<&str, String>,
    backend_addr: &str,
) {
    // Connect to backend server
    let server_stream = match TcpStream::connect(backend_addr).await {
        Ok(stream) => stream,
        Err(e) => {
            log::warn!("Failed to connect to backend {backend_addr}: {e}");
            let _ = client_stream.shutdown().await;
            return;
        }
    };

    log::debug!("Connected to backend server at {backend_addr}");

    // Forward the original request
    let mut request = format!("{method} {path} {protocol}\r\n");
    for (key, value) in headers {
        request.push_str(&format!("{key}: {value}\r\n"));
    }
    request.push_str("\r\n");

    let mut server_stream = server_stream;
    if server_stream.write_all(request.as_bytes()).await.is_err() {
        log::debug!("Failed to write request to backend server");
        let _ = client_stream.shutdown().await;
        return;
    }

    // Set up bidirectional forwarding between client and backend
    let (mut client_read, mut client_write) = client_stream.split();
    let (mut server_read, mut server_write) = server_stream.split();

    // Client -> Server
    let client_to_server = async {
        let mut buf = [0; 8192];
        let mut bytes_forwarded = 0;

        loop {
            match client_read.read(&mut buf).await {
                Ok(0) => break,
                Ok(n) => {
                    bytes_forwarded += n;
                    if server_write.write_all(&buf[..n]).await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }

        log::debug!("Client -> Server: forwarded {bytes_forwarded} bytes");
    };

    // Server -> Client
    let server_to_client = async {
        let mut buf = [0; 8192];
        let mut bytes_forwarded = 0;

        loop {
            match server_read.read(&mut buf).await {
                Ok(0) => break,
                Ok(n) => {
                    bytes_forwarded += n;
                    if client_write.write_all(&buf[..n]).await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }

        log::debug!("Server -> Client: forwarded {bytes_forwarded} bytes");
    };

    // Run both directions concurrently
    tokio::select! {
        () = client_to_server => {},
        () = server_to_client => {},
    }

    log::debug!("Proxy connection completed");
}

// Set up nftables firewall rules for IP blocking
async fn setup_firewall() -> Result<(), String> {
    log::info!("Setting up firewall rules");

    // Check if nft command exists
    let nft_exists = Command::new("which")
        .arg("nft")
        .output()
        .await
        .map(|output| output.status.success())
        .unwrap_or(false);

    if !nft_exists {
        log::warn!("nft command not found. Firewall rules will not be set up.");
        return Ok(());
    }

    // Create table if it doesn't exist
    let output = Command::new("nft")
        .args(["list", "table", "inet", "filter"])
        .output()
        .await;

    match output {
        Ok(output) => {
            if !output.status.success() {
                log::info!("Creating nftables table");
                let result = Command::new("nft")
                    .args(["create", "table", "inet", "filter"])
                    .output()
                    .await;

                if let Err(e) = result {
                    return Err(format!("Failed to create nftables table: {e}"));
                }
            }
        }
        Err(e) => {
            log::warn!("Failed to check if nftables table exists: {e}");
            log::info!("Will try to create it anyway");
            let result = Command::new("nft")
                .args(["create", "table", "inet", "filter"])
                .output()
                .await;

            if let Err(e) = result {
                return Err(format!("Failed to create nftables table: {e}"));
            }
        }
    }

    // Create blacklist set if it doesn't exist
    let output = Command::new("nft")
        .args(["list", "set", "inet", "filter", "eris_blacklist"])
        .output()
        .await;

    match output {
        Ok(output) => {
            if !output.status.success() {
                log::info!("Creating eris_blacklist set");
                let result = Command::new("nft")
                    .args([
                        "create",
                        "set",
                        "inet",
                        "filter",
                        "eris_blacklist",
                        "{ type ipv4_addr; flags interval; }",
                    ])
                    .output()
                    .await;

                if let Err(e) = result {
                    return Err(format!("Failed to create blacklist set: {e}"));
                }
            }
        }
        Err(e) => {
            log::warn!("Failed to check if blacklist set exists: {e}");
            return Err(format!("Failed to check if blacklist set exists: {e}"));
        }
    }

    // Add rule to drop traffic from blacklisted IPs
    let output = Command::new("nft")
        .args(["list", "chain", "inet", "filter", "input"])
        .output()
        .await;

    // Check if our rule already exists
    match output {
        Ok(output) => {
            let rule_exists = String::from_utf8_lossy(&output.stdout)
                .contains("ip saddr @eris_blacklist counter drop");

            if !rule_exists {
                log::info!("Adding drop rule for blacklisted IPs");
                let result = Command::new("nft")
                    .args([
                        "add",
                        "rule",
                        "inet",
                        "filter",
                        "input",
                        "ip saddr @eris_blacklist",
                        "counter",
                        "drop",
                    ])
                    .output()
                    .await;

                if let Err(e) = result {
                    return Err(format!("Failed to add firewall rule: {e}"));
                }
            }
        }
        Err(e) => {
            log::warn!("Failed to check if firewall rule exists: {e}");
            return Err(format!("Failed to check if firewall rule exists: {e}"));
        }
    }

    log::info!("Firewall setup complete");
    Ok(())
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    // Parse command line arguments
    let args = Args::parse();

    // Initialize the logger
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or(&args.log_level))
        .format_timestamp_millis()
        .init();

    log::info!("Starting eris tarpit system");

    // Load configuration
    let config = if let Some(config_path) = &args.config_file {
        log::info!("Loading configuration from {config_path:?}");
        match Config::load_from_file(config_path) {
            Ok(cfg) => {
                log::info!("Configuration loaded successfully");
                cfg
            }
            Err(e) => {
                log::warn!("Failed to load configuration file: {e}");
                log::info!("Using configuration from command-line arguments");
                Config::from_args(&args)
            }
        }
    } else {
        log::info!("Using configuration from command-line arguments");
        Config::from_args(&args)
    };

    // Ensure required directories exist
    match config.ensure_dirs_exist() {
        Ok(()) => log::info!("Directory setup completed"),
        Err(e) => {
            log::error!("Failed to create required directories: {e}");
            log::info!("Will continue with default in-memory configuration");
        }
    }

    // Save config for reference if it was loaded from command line
    if args.config_file.is_none() {
        if let Err(e) = fs::create_dir_all(&config.config_dir) {
            log::warn!("Failed to create config directory: {e}");
        } else {
            let config_path = Path::new(&config.config_dir).join("config.json");
            if !config_path.exists() {
                if let Err(e) = config.save_to_file(&config_path) {
                    log::warn!("Failed to save default configuration: {e}");
                } else {
                    log::info!("Saved default configuration to {config_path:?}");
                }
            }
        }
    }

    log::info!("Using directories:");
    log::info!("  Config:  {}", config.config_dir);
    log::info!("  Corpora: {}", config.markov_corpora_dir);
    log::info!("  Scripts: {}", config.lua_scripts_dir);
    log::info!("  Data:    {}", config.data_dir);
    log::info!("  Cache:   {}", config.cache_dir);

    let config = Arc::new(config);

    // Setup firewall rules for IP blocking
    match setup_firewall().await {
        Ok(()) => {}
        Err(e) => {
            log::warn!("Failed to set up firewall rules: {e}");
            log::info!("IP blocking will be managed in memory only");
        }
    }

    // Initialize bot state for both servers
    let tarpit_state = Arc::new(RwLock::new(BotState::load_from_disk(
        &config.data_dir,
        &config.cache_dir,
    )));
    let metrics_state = tarpit_state.clone();

    // Initialize Markov chain text generator
    log::info!(
        "Initializing Markov chain generator from {}",
        config.markov_corpora_dir
    );
    let markov_generator = Arc::new(MarkovGenerator::new(&config.markov_corpora_dir));

    // Initialize Lua script manager
    log::info!("Loading Lua scripts from {}", config.lua_scripts_dir);
    let script_manager = Arc::new(ScriptManager::new(&config.lua_scripts_dir));

    // Clone config for metrics server
    let metrics_config = config.clone();

    // Start the main tarpit server
    let tarpit_server = tokio::spawn(async move {
        log::info!("Starting tarpit server on {}", config.listen_addr);

        let listener = match TcpListener::bind(&config.listen_addr).await {
            Ok(l) => l,
            Err(e) => {
                return Err(format!("Failed to bind to {}: {}", config.listen_addr, e));
            }
        };

        log::info!("Tarpit server listening on {}", config.listen_addr);

        loop {
            match listener.accept().await {
                Ok((stream, addr)) => {
                    log::debug!("Accepted connection from {addr}");

                    let state_clone = tarpit_state.clone();
                    let markov_clone = markov_generator.clone();
                    let script_manager_clone = script_manager.clone();
                    let config_clone = config.clone();

                    tokio::spawn(async move {
                        handle_connection(
                            stream,
                            config_clone,
                            state_clone,
                            markov_clone,
                            script_manager_clone,
                        )
                        .await;
                    });
                }
                Err(e) => {
                    log::error!("Error accepting connection: {e}");
                }
            }
        }

        #[allow(unreachable_code)]
        Ok::<(), String>(())
    });

    // Start the metrics server with actix_web only if metrics are not disabled
    let metrics_server = if metrics_config.disable_metrics {
        log::info!("Metrics server disabled via configuration");
        None
    } else {
        log::info!("Starting metrics server on {}", metrics_config.metrics_addr);

        let server = HttpServer::new(move || {
            App::new()
                .app_data(web::Data::new(metrics_state.clone()))
                .route("/metrics", web::get().to(metrics_handler))
                .route("/status", web::get().to(status_handler))
                .route("/", web::get().to(|| async {
                    HttpResponse::Ok().body("Botpot Server is running. Visit /metrics for metrics or /status for status.")
                }))
        })
        .bind(&metrics_config.metrics_addr);

        match server {
            Ok(server) => Some(server.run()),
            Err(e) => {
                log::error!(
                    "Failed to bind metrics server to {}: {}",
                    metrics_config.metrics_addr,
                    e
                );
                None
            }
        }
    };

    // Run both servers concurrently if metrics server is enabled
    if let Some(metrics_server) = metrics_server {
        tokio::select! {
            result = tarpit_server => match result {
                Ok(Ok(())) => Ok(()),
                Ok(Err(e)) => {
                    log::error!("Tarpit server error: {e}");
                    Err(std::io::Error::new(std::io::ErrorKind::Other, e))
                },
                Err(e) => {
                    log::error!("Tarpit server task error: {e}");
                    Err(std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))
                },
            },
            result = metrics_server => {
                if let Err(ref e) = result {
                    log::error!("Metrics server error: {e}");
                }
                result
            },
        }
    } else {
        // Just run the tarpit server if metrics are disabled
        match tarpit_server.await {
            Ok(Ok(())) => Ok(()),
            Ok(Err(e)) => {
                log::error!("Tarpit server error: {e}");
                Err(std::io::Error::new(std::io::ErrorKind::Other, e))
            }
            Err(e) => {
                log::error!("Tarpit server task error: {e}");
                Err(std::io::Error::new(
                    std::io::ErrorKind::Other,
                    e.to_string(),
                ))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};
    use tokio::sync::RwLock;

    #[test]
    fn test_config_from_args() {
        let args = Args {
            listen_addr: "127.0.0.1:8080".to_string(),
            metrics_addr: "127.0.0.1:9000".to_string(),
            disable_metrics: true,
            backend_addr: "127.0.0.1:8081".to_string(),
            min_delay: 500,
            max_delay: 10000,
            max_tarpit_time: 300,
            block_threshold: 5,
            base_dir: Some(PathBuf::from("/tmp/eris")),
            config_file: None,
            log_level: "debug".to_string(),
        };

        let config = Config::from_args(&args);
        assert_eq!(config.listen_addr, "127.0.0.1:8080");
        assert_eq!(config.metrics_addr, "127.0.0.1:9000");
        assert!(config.disable_metrics);
        assert_eq!(config.backend_addr, "127.0.0.1:8081");
        assert_eq!(config.min_delay, 500);
        assert_eq!(config.max_delay, 10000);
        assert_eq!(config.max_tarpit_time, 300);
        assert_eq!(config.block_threshold, 5);
        assert_eq!(config.markov_corpora_dir, "/tmp/eris/data/corpora");
        assert_eq!(config.lua_scripts_dir, "/tmp/eris/data/scripts");
        assert_eq!(config.data_dir, "/tmp/eris/data");
        assert_eq!(config.config_dir, "/tmp/eris/conf");
        assert_eq!(config.cache_dir, "/tmp/eris/cache");
    }

    #[tokio::test]
    async fn test_should_tarpit() {
        let config = Config::default();

        // Test trap patterns
        assert!(
            should_tarpit(
                "/vendor/phpunit/whatever",
                &IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)),
                &config
            )
            .await
        );
        assert!(
            should_tarpit(
                "/wp-admin/login.php",
                &IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)),
                &config
            )
            .await
        );
        assert!(should_tarpit("/.env", &IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), &config).await);

        // Test whitelist networks
        assert!(
            !should_tarpit(
                "/wp-admin/login.php",
                &IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                &config
            )
            .await
        );
        assert!(
            !should_tarpit(
                "/vendor/phpunit/whatever",
                &IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)),
                &config
            )
            .await
        );

        // Test legitimate paths
        assert!(
            !should_tarpit(
                "/index.html",
                &IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)),
                &config
            )
            .await
        );
        assert!(
            !should_tarpit(
                "/images/logo.png",
                &IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)),
                &config
            )
            .await
        );
    }

    #[test]
    fn test_script_manager_default_script() {
        let script_manager = ScriptManager::new("/nonexistent_directory");
        assert!(script_manager.scripts_loaded);
        assert!(
            script_manager
                .script_content
                .contains("generate_honeytoken")
        );
        assert!(script_manager.script_content.contains("enhance_response"));
    }

    #[tokio::test]
    async fn test_bot_state() {
        let state = BotState::new("/tmp/eris_test", "/tmp/eris_test_cache");
        let ip1 = IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4));
        let ip2 = IpAddr::V4(Ipv4Addr::new(5, 6, 7, 8));

        let state = Arc::new(RwLock::new(state));

        // Test hit counter
        {
            let mut state = state.write().await;
            *state.hits.entry(ip1).or_insert(0) += 1;
            *state.hits.entry(ip1).or_insert(0) += 1;
            *state.hits.entry(ip2).or_insert(0) += 1;

            assert_eq!(*state.hits.get(&ip1).unwrap(), 2);
            assert_eq!(*state.hits.get(&ip2).unwrap(), 1);
        }

        // Test blocking
        {
            let mut state = state.write().await;
            state.blocked.insert(ip1);
            assert!(state.blocked.contains(&ip1));
            assert!(!state.blocked.contains(&ip2));
        }

        // Test active connections
        {
            let mut state = state.write().await;
            state.active_connections.insert(ip1);
            state.active_connections.insert(ip2);
            assert_eq!(state.active_connections.len(), 2);

            state.active_connections.remove(&ip1);
            assert_eq!(state.active_connections.len(), 1);
            assert!(!state.active_connections.contains(&ip1));
            assert!(state.active_connections.contains(&ip2));
        }
    }

    #[tokio::test]
    async fn test_generate_deceptive_response() {
        // Create a simple markov generator for testing
        let markov = MarkovGenerator::new("/nonexistent/path");
        let script_manager = ScriptManager::new("/nonexistent/path");

        // Test different path types
        let resp1 = generate_deceptive_response(
            "/vendor/phpunit/exec",
            "TestBot/1.0",
            &markov,
            &script_manager,
        )
        .await;
        assert!(resp1.contains("HTTP/1.1 200 OK"));
        assert!(resp1.contains("X-Powered-By: PHP"));

        let resp2 =
            generate_deceptive_response("/wp-admin/", "TestBot/1.0", &markov, &script_manager)
                .await;
        assert!(resp2.contains("HTTP/1.1 200 OK"));

        let resp3 =
            generate_deceptive_response("/api/users", "TestBot/1.0", &markov, &script_manager)
                .await;
        assert!(resp3.contains("HTTP/1.1 200 OK"));

        // Verify tracking token is included
        assert!(resp1.contains("BOT_TestBot"));
    }
}