porta-rs 0.1.1

use anyhow::{Context, Result};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use parking_lot::RwLock;
use std::collections::HashMap;

use crate::config::ServerConfig;
use super::protocol::{Command, Response, PortEntry};

/// Registered client state
#[derive(Debug, Clone)]
pub struct ClientState {
    pub client_id: String,
    pub ports: Vec<PortEntry>,
    pub last_heartbeat: i64,
}

/// Control server - receives commands from clients via TCP
pub struct ControlServer {
    listener: TcpListener,
    pub clients: Arc<RwLock<HashMap<String, ClientState>>>,
}

impl ControlServer {
    /// Create a new control server
    pub async fn new(config: &ServerConfig) -> Result<Self> {
        // Listen on a port near the WireGuard port for control
        let control_port = config.listen_port + 1;
        let addr = format!("0.0.0.0:{}", control_port);

        let listener = TcpListener::bind(&addr)
            .await
            .context(format!("Failed to bind control server on {}", addr))?;

        tracing::info!("Control server listening on {}", addr);

        Ok(Self {
            listener,
            clients: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Run the control server (blocking loop)
    pub async fn run(&self) -> Result<()> {
        tracing::info!("Control server running, waiting for connections...");

        loop {
            match self.listener.accept().await {
                Ok((stream, addr)) => {
                    let clients = self.clients.clone();
                    tokio::spawn(async move {
                        if let Err(e) = Self::handle_client(stream, addr, clients).await {
                            tracing::debug!("Client {} error: {}", addr, e);
                        }
                    });
                }
                Err(e) => {
                    tracing::error!("Accept error: {}", e);
                    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
                }
            }
        }
    }

    /// Handle a single client connection
    async fn handle_client(
        mut stream: TcpStream,
        addr: SocketAddr,
        clients: Arc<RwLock<HashMap<String, ClientState>>>,
    ) -> Result<()> {
        tracing::debug!("New control connection from {}", addr);

        loop {
            // Read command length
            let mut len_buf = [0u8; 4];
            match stream.read_exact(&mut len_buf).await {
                Ok(_) => {}
                Err(e) => {
                    if e.kind() == std::io::ErrorKind::UnexpectedEof {
                        tracing::debug!("Client {} disconnected", addr);
                    }
                    return Ok(());
                }
            }

            let len = u32::from_le_bytes(len_buf) as usize;
            if len > 1024 * 1024 {
                tracing::warn!("Command too large from {}: {} bytes", addr, len);
                return Ok(());
            }

            // Read command data
            let mut data = vec![0u8; len];
            stream.read_exact(&mut data).await?;

            // Parse command
            let command = Command::from_bytes(&data)?;

            // Handle command
            let response = Self::handle_command(&command, &clients);

            // Send response
            let resp_data = response.to_bytes()?;
            let resp_len = resp_data.len() as u32;
            stream.write_all(&resp_len.to_le_bytes()).await?;
            stream.write_all(&resp_data).await?;
        }
    }

    /// Handle a command and return a response
    fn handle_command(
        command: &Command,
        clients: &Arc<RwLock<HashMap<String, ClientState>>>,
    ) -> Response {
        match command {
            Command::RegisterPorts {
                client_id,
                ports,
            } => {
                tracing::info!(
                    "Client '{}' registered {} port(s)",
                    client_id,
                    ports.len()
                );

                for port in ports {
                    tracing::info!(
                        "  Port {}:{} ({}) - {}",
                        port.remote_port,
                        port.local_port,
                        port.protocol,
                        port.description
                    );
                }

                let mut clients = clients.write();
                clients.insert(
                    client_id.clone(),
                    ClientState {
                        client_id: client_id.clone(),
                        ports: ports.clone(),
                        last_heartbeat: chrono::Utc::now().timestamp(),
                    },
                );

                Response::Ok {
                    message: format!("Registered {} port(s)", ports.len()),
                }
            }

            Command::UnregisterPorts { client_id } => {
                tracing::info!("Client '{}' unregistered all ports", client_id);

                let mut clients = clients.write();
                clients.remove(client_id);

                Response::Ok {
                    message: "Ports unregistered".to_string(),
                }
            }

            Command::Heartbeat {
                client_id,
                timestamp,
            } => {
                tracing::debug!("Heartbeat from '{}' at {}", client_id, timestamp);

                let mut clients = clients.write();
                if let Some(state) = clients.get_mut(client_id) {
                    state.last_heartbeat = *timestamp;
                }

                Response::HeartbeatAck {
                    server_time: chrono::Utc::now().timestamp(),
                }
            }
        }
    }

    /// Get all registered clients
    pub fn get_clients(&self) -> HashMap<String, ClientState> {
        self.clients.read().clone()
    }
}