porta-rs 0.1.1

use anyhow::Result;
use clap::Args;
use std::sync::Arc;
use tokio::sync::Notify;
use tokio::time::{interval, Duration};

use crate::config::{load_config, Config, Role};
use crate::control::client::ControlClient;
use crate::control::protocol::PortEntry;
use crate::forwarder::client_listener::ClientListener;
use crate::forwarder::server_listener::PortListener;
use crate::ufw::rules::UfwManager;
use crate::wireguard::tunnel::Tunnel;

#[derive(Args)]
pub struct RunArgs {
    /// Role to run as (client or server)
    #[arg(short, long)]
    pub role: Option<String>,
}

pub async fn execute(_args: &RunArgs) -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "info".into()),
        )
        .init();

    tracing::info!("Porta v{} starting", env!("CARGO_PKG_VERSION"));

    let config = load_config()?;
    let shutdown = Arc::new(Notify::new());
    let shutdown_clone = shutdown.clone();

    tokio::spawn(async move {
        tokio::signal::ctrl_c().await.ok();
        tracing::info!("Shutdown signal received");
        shutdown_clone.notify_waiters();
    });

    match config.role() {
        Role::Client => run_client(config, shutdown).await,
        Role::Server => run_server(config, shutdown).await,
    }
}

/// Server: listens on public IPs, forwards via WireGuard, handles control commands
async fn run_server(config: Config, shutdown: Arc<Notify>) -> Result<()> {
    tracing::info!("Running as SERVER (machine with public IP)");

    let server_config = match &config {
        Config::Server(s) => s.clone(),
        _ => return Err(anyhow::anyhow!("Expected server config")),
    };

    // Start WireGuard tunnel
    let tunnel = Tunnel::new(config.clone())?;
    tunnel.start().await?;
    tracing::info!("WireGuard tunnel started on port {}", server_config.listen_port);

    // Setup UFW
    let ufw = UfwManager::new();
    if ufw.is_active() {
        let _ = ufw.allow_port(server_config.listen_port, &crate::config::Protocol::Tcp);
        let _ = ufw.allow_port(server_config.listen_port + 1, &crate::config::Protocol::Tcp);
        tracing::info!("UFW rules configured");
    } else {
        tracing::warn!("UFW is not active");
    }

    // Start control server
    let control = crate::control::server::ControlServer::new(&server_config).await?;

    // Get shared client state from control server
    let clients_state = control.clients.clone();

    tokio::spawn(async move {
        if let Err(e) = control.run().await {
            tracing::error!("Control server error: {}", e);
        }
    });

    tracing::info!("Control server started on port {}", server_config.listen_port + 1);

    // Port listener (will be started when client registers ports)
    let port_listener = Arc::new(PortListener::new());
    let client_vpn_ip = "10.0.0.2".to_string();

    // Start listening on configured ports from config
    let config_ports = get_static_ports(&server_config);
    for mapping in &config_ports {
        if let Err(e) = port_listener.start(mapping, &client_vpn_ip).await {
            tracing::error!("Failed to start listener: {}", e);
        }
    }

    // Background task: sync ports from control plane
    let listener_clone = port_listener.clone();
    let vpn_ip = client_vpn_ip.clone();
    tokio::spawn(async move {
        let mut ticker = interval(Duration::from_secs(5));
        loop {
            ticker.tick().await;
            let clients = clients_state.read().clone();
            for (_id, state) in &clients {
                for port in &state.ports {
                    let mapping = crate::config::PortMapping {
                        id: format!("dyn-{}-{}", state.client_id, port.remote_port),
                        local_port: port.local_port,
                        remote_port: port.remote_port,
                        protocol: port.protocol.parse().unwrap_or(crate::config::Protocol::Tcp),
                        description: port.description.clone(),
                        enabled: true,
                    };
                    if let Err(_e) = listener_clone.start(&mapping, &vpn_ip).await {
                        tracing::debug!("Port {} already listening", port.remote_port);
                    }
                }
            }
        }
    });

    tracing::info!(
        "Server ready with {} static port(s). Waiting for client connections...",
        config_ports.len()
    );

    shutdown.notified().await;
    tracing::info!("Shutting down...");
    tunnel.stop().await?;
    tracing::info!("Porta server stopped");
    Ok(())
}

/// Client: receives traffic from WireGuard, forwards to local services
async fn run_client(config: Config, shutdown: Arc<Notify>) -> Result<()> {
    tracing::info!("Running as CLIENT (machine behind CGNAT)");

    let client_config = match &config {
        Config::Client(c) => c.clone(),
        _ => return Err(anyhow::anyhow!("Expected client config")),
    };

    // Start WireGuard tunnel
    let tunnel = Tunnel::new(config.clone())?;
    tunnel.start().await?;
    tracing::info!(
        "WireGuard tunnel started, server: {}:{}",
        client_config.server_address,
        client_config.server_port
    );

    // Start client listeners for each port
    let client_listener = ClientListener::new();
    let local_bind_ip = &client_config.local_ip.split('/').next().unwrap_or("10.0.0.2");

    for mapping in config.ports() {
        if mapping.enabled {
            if let Err(e) = client_listener.start(mapping, local_bind_ip).await {
                tracing::error!("Failed to start listener for port {}: {}", mapping.remote_port, e);
            }
        }
    }

    tracing::info!(
        "Client listening on {} port(s)",
        config.ports().iter().filter(|p| p.enabled).count()
    );

    // Connect to control server and register ports
    let control_client = ControlClient::new(&client_config)?;

    // Register ports with server
    let port_entries: Vec<PortEntry> = config.ports()
        .iter()
        .filter(|p| p.enabled)
        .map(|p| PortEntry {
            remote_port: p.remote_port,
            local_port: p.local_port,
            protocol: p.protocol.to_string(),
            description: p.description.clone(),
        })
        .collect();

    if !port_entries.is_empty() {
        match control_client.register_ports(&port_entries).await {
            Ok(response) => {
                tracing::info!("Server response: {:?}", response);
            }
            Err(e) => {
                tracing::warn!("Failed to register ports with server: {} (server may not be ready yet)", e);
            }
        }
    }

    // Heartbeat loop
    let heartbeat_client = ControlClient::new(&client_config)?;
    tokio::spawn(async move {
        let mut ticker = interval(Duration::from_secs(30));
        loop {
            ticker.tick().await;
            match heartbeat_client.heartbeat().await {
                Ok(_) => tracing::debug!("Heartbeat sent"),
                Err(e) => tracing::warn!("Heartbeat failed: {}", e),
            }
        }
    });

    tracing::info!("Client ready. Press Ctrl+C to stop.");

    shutdown.notified().await;
    tracing::info!("Shutting down...");
    tunnel.stop().await?;
    tracing::info!("Porta client stopped");
    Ok(())
}

/// Get static ports from server config (for ports defined in allowed_clients)
fn get_static_ports(config: &crate::config::ServerConfig) -> Vec<crate::config::PortMapping> {
    let mut ports = Vec::new();

    // If there are allowed clients, we forward some default ports
    // In production, this would come from the control plane
    if !config.allowed_clients.is_empty() {
        ports.push(crate::config::PortMapping {
            id: "http".to_string(),
            local_port: 80,
            remote_port: 80,
            protocol: crate::config::Protocol::Tcp,
            description: "HTTP".to_string(),
            enabled: true,
        });
        ports.push(crate::config::PortMapping {
            id: "https".to_string(),
            local_port: 443,
            remote_port: 443,
            protocol: crate::config::Protocol::Tcp,
            description: "HTTPS".to_string(),
            enabled: true,
        });
    }

    ports
}