//! # Connection Packet Processing
//!
//! This module implements the connection establishment protocol for Citadel.
//! It handles the secure handshake process between nodes, including authentication,
//! capability negotiation, and secure channel establishment.
//!
//! # Protocol Flow
//!
//! 1. **Stage 0**: Initial connection request
//!    - Client sends encrypted credentials
//!    - Server validates and processes request
//!    - Capability negotiation (filesystem, UDP support)
//!
//! 2. **Success Stage**:
//!    - Secure channel establishment
//!    - Virtual connection initialization
//!    - Session security settings configuration
//!    - UDP channel setup (if supported)
//!
//! # Security Features
//!
//! - Post-quantum cryptographic handshake
//! - Secure credential transmission
//! - Session-specific security settings
//! - State validation at each stage
//!
//! # Related Components
//!
//! - `PreConnectPacket`: Handles pre-connection setup
//! - `StateContainer`: Manages connection state
//! - `VirtualConnection`: Manages established connections
//! - `SessionSecuritySettings`: Configures connection security
//!

use super::includes::*;
use crate::error::NetworkError;
use crate::prelude::Ticket;
use crate::proto::misc::platform_ops::PlatformOps;
use crate::proto::node_result::{ConnectFail, ConnectSuccess, MailboxDelivery};
use crate::proto::packet_crafter::peer_cmd::C2S_IDENTITY_CID;
use crate::proto::packet_processor::primary_group_packet::get_orientation_safe_ratchet;
use citadel_crypt::ratchets::Ratchet;
use citadel_io::{error, ErrorCode};
use citadel_types::proto::ConnectMode;
use citadel_user::external_services::ServicesObject;

/// This will optionally return an HdpPacket as a response if deemed necessary
#[cfg_attr(feature = "localhost-testing", tracing::instrument(
    level = "trace",
    target = "citadel",
    skip_all,
    ret,
    err,
    fields(is_server = sess_ref.is_server, src = packet.parse().unwrap().0.session_cid.get(), target = packet.parse().unwrap().0.target_cid.get()
    )
))]
pub async fn process_connect<R: Ratchet, T: PlatformOps>(
    sess_ref: &CitadelSession<R, T>,
    packet: HdpPacket,
    header_entropy_bank_vers: u32,
) -> Result<PrimaryProcessorResult, NetworkError> {
    let session = sess_ref.clone();

    let (hr, cnac) = {
        let state_container = inner_state!(session.state_container);
        if !session.is_provisional()
            && state_container.connect_state.last_stage
                != packet_flags::cmd::aux::do_connect::SUCCESS
        {
            log::error!(target: "citadel", "Connect packet received, but the system is not in a provisional state. Dropping");
            return Ok(PrimaryProcessorResult::Void);
        }

        if !state_container.pre_connect_state.success {
            log::error!(target: "citadel", "Connect packet received, but the system has not yet completed the pre-connect stage. Dropping");
            return Ok(PrimaryProcessorResult::Void);
        }

        let hr = return_if_none!(
            get_orientation_safe_ratchet(header_entropy_bank_vers, &state_container, None),
            "Could not get proper HR [connect]"
        );
        let cnac = return_if_none!(state_container.cnac.clone(), "CNAC missing");
        (hr, cnac)
    };

    let (header, payload, _, _) = packet.decompose();

    let (header, payload, ratchet) = return_if_none!(
        validation::aead::validate(hr, &header, payload),
        "Unable to validate connect packet"
    );
    let header = header.clone();
    let security_level = header.security_level.into();
    let ticket = Ticket::from(header.context_info.get());
    let time_tracker = session.time_tracker;

    let task = async move {
        let session = &session;
        match header.cmd_aux {
            // Node is Bob. Bob gets the encrypted username and password (separately encrypted)
            packet_flags::cmd::aux::do_connect::STAGE0 => {
                log::trace!(target: "citadel", "STAGE 0 CONNECT PACKET");
                // Idempotency / replay guard: STAGE0 drives the one-time connection-upgrade path
                // (init_new_c2s_virtual_connection -> upgrade_connection -> ConnectSuccess). The
                // top-level guard above intentionally lets packets through once last_stage == SUCCESS
                // (so the client can process the final status packet), but a replayed/duplicated
                // STAGE0 reaching an already-SUCCESS session would re-run that path and churn/evict
                // the live connection. Drop it here.
                {
                    let state_container = inner_state!(session.state_container);
                    if state_container.connect_state.last_stage
                        == packet_flags::cmd::aux::do_connect::SUCCESS
                    {
                        log::warn!(target: "citadel", "Dropping duplicate/replayed STAGE0 connect packet; session already reached SUCCESS");
                        return Ok(PrimaryProcessorResult::Void);
                    }
                }
                let task = {
                    match validation::do_connect::validate_stage0_packet(&cnac, &payload).await {
                        Ok(stage0_packet) => {
                            // Compute inexpensive values and perform checks before taking the lock
                            let local_uses_file_system = session
                                .account_manager
                                .get_backend_type()
                                .is_filesystem_backend();
                            session
                                .file_transfer_compatible
                                .set_once(local_uses_file_system && stage0_packet.uses_filesystem);
                            let cid = ratchet.get_cid();
                            let success_time = session.time_tracker.get_global_time_ns();
                            let addr = session.remote_peer;
                            let is_personal = !session.is_server;
                            let kernel_ticket = session.kernel_ticket.get();

                            if kernel_ticket != ticket {
                                const REASON: &str = "Ticket mismatch";
                                return Ok(PrimaryProcessorResult::EndSession(REASON));
                            }

                            // Acquire lock only for short, critical state updates; release before awaits
                            let (udp_channel_rx, channel, session_security_settings) = {
                                let mut state_container = inner_mut_state!(session.state_container);
                                state_container.connect_state.last_stage =
                                    packet_flags::cmd::aux::do_connect::SUCCESS;
                                state_container.connect_state.fail_time = None;
                                state_container.connect_state.on_connect_packet_received();
                                let udp_channel_rx = state_container
                                    .pre_connect_state
                                    .udp_channel_oneshot_tx
                                    .rx
                                    .take();
                                let channel = state_container.init_new_c2s_virtual_connection(
                                    &cnac,
                                    kernel_ticket,
                                    header.session_cid.get(),
                                    session,
                                );
                                let session_security_settings = state_container
                                    .session_security_settings
                                    .expect("Should be set");
                                (udp_channel_rx, channel, session_security_settings)
                            };

                            // Upgrade the connect BEFORE updating the CNAC
                            if !session.session_manager.upgrade_connection(addr, cid) {
                                return Ok(PrimaryProcessorResult::EndSession("Unable to upgrade from a provisional to a protected connection (Server)"));
                            }

                            // register w/ peer layer, get mail in the process
                            let account_manager = session.account_manager.clone();

                            async move {
                                let mailbox_items = session
                                    .session_manager
                                    .register_session_with_peer_layer(cid)
                                    .await?;
                                let peers = account_manager
                                    .get_persistence_handler()
                                    .get_hyperlan_peer_list_as_server(cid)
                                    .await?
                                    .unwrap_or_default();

                                let post_login_object = account_manager
                                    .services_handler()
                                    .on_post_login_serverside(cid)
                                    .await?;

                                let success_packet =
                                    packet_crafter::do_connect::craft_final_status_packet(
                                        &ratchet,
                                        true,
                                        mailbox_items,
                                        post_login_object.clone(),
                                        session.create_welcome_message(cid),
                                        peers,
                                        success_time,
                                        security_level,
                                        session.account_manager.get_backend_type(),
                                        kernel_ticket,
                                    );

                                session.session_cid.set(Some(cid));
                                session.state.set(SessionState::Connected);

                                let cxn_type =
                                    VirtualConnectionType::LocalGroupServer { session_cid: cid };
                                let channel_signal = NodeResult::ConnectSuccess(ConnectSuccess {
                                    ticket: kernel_ticket,
                                    session_cid: cid,
                                    remote_addr: addr,
                                    is_personal,
                                    v_conn_type: cxn_type,
                                    services: post_login_object,
                                    welcome_message: format!("Client {cid} successfully established a connection to the local HyperNode"),
                                    channel: channel.into(),
                                    udp_rx_opt: udp_channel_rx,
                                    session_security_settings,
                                });
                                // Store the signal (post-await), using a short-lived lock
                                inner_mut_state!(session.state_container)
                                    .get_endpoint_container_mut(C2S_IDENTITY_CID)?
                                    .channel_signal = Some(channel_signal);
                                Ok(PrimaryProcessorResult::ReplyToSender(success_packet))
                            }
                        }

                        Err(err) => {
                            log::error!(target: "citadel", "Error validating stage2 packet. Reason: {err}");
                            let fail_time = time_tracker.get_global_time_ns();

                            let packet = packet_crafter::do_connect::craft_final_status_packet(
                                &ratchet,
                                false,
                                None,
                                ServicesObject::default(),
                                err.to_string(),
                                Vec::new(),
                                fail_time,
                                security_level,
                                session.account_manager.get_backend_type(),
                                ticket,
                            );
                            return Ok(PrimaryProcessorResult::ReplyToSender(packet));
                        }
                    }
                };

                task.await
            }

            packet_flags::cmd::aux::do_connect::FAILURE => {
                log::trace!(target: "citadel", "STAGE FAILURE CONNECT PACKET");
                let kernel_ticket = session.kernel_ticket.get();

                let mut state_container = inner_mut_state!(session.state_container);
                if let Some(payload) =
                    validation::do_connect::validate_final_status_packet(&payload)
                {
                    let message = String::from_utf8(payload.message.to_vec())
                        .unwrap_or_else(|_| "Invalid UTF-8 message".to_string());
                    log::error!(target: "citadel", "The server refused to login the user. Reason: {}", &message);
                    let cid = ratchet.get_cid();
                    state_container.connect_state.on_fail();
                    drop(state_container);

                    session.session_cid.set(None);
                    session.state.set(SessionState::NeedsConnect);
                    session.disable_dc_signal();

                    session.send_to_kernel(NodeResult::ConnectFail(ConnectFail {
                        ticket: kernel_ticket,
                        cid_opt: Some(cid),
                        error_message: message,
                    }))?;
                    Ok(PrimaryProcessorResult::EndSession(
                        "Failed connecting. Try again",
                    ))
                } else {
                    trace!(target: "citadel", "An invalid FAILURE packet was received; dropping due to invalid signature");
                    Ok(PrimaryProcessorResult::Void)
                }
            }

            // Node is finally Alice. The login either failed or succeeded
            packet_flags::cmd::aux::do_connect::SUCCESS => {
                log::trace!(target: "citadel", "STAGE SUCCESS CONNECT PACKET");

                let task = {
                    // Compute compatibility outside of lock
                    let remote_uses_filesystem = header.group.get() != 0;
                    let local_uses_file_system = session
                        .account_manager
                        .get_backend_type()
                        .is_filesystem_backend();
                    session
                        .file_transfer_compatible
                        .set_once(local_uses_file_system && remote_uses_filesystem);

                    let mut state_container = inner_mut_state!(session.state_container);
                    let last_stage = state_container.connect_state.last_stage;

                    if last_stage == packet_flags::cmd::aux::do_connect::STAGE1 {
                        if let Some(payload) =
                            validation::do_connect::validate_final_status_packet(&payload)
                        {
                            let cnac = cnac.clone();
                            let message = String::from_utf8(payload.message.to_vec())
                                .unwrap_or_else(|_| String::from("Invalid message"));
                            let kernel_ticket = session.kernel_ticket.get();
                            let cid = ratchet.get_cid();

                            state_container.connect_state.on_success();
                            state_container.connect_state.on_connect_packet_received();

                            let use_ka = state_container.keep_alive_timeout_ns != 0;
                            let connect_mode = return_if_none!(
                                state_container.connect_state.connect_mode,
                                "Unable to load connect mode"
                            );
                            let udp_channel_rx = state_container
                                .pre_connect_state
                                .udp_channel_oneshot_tx
                                .rx
                                .take();

                            let channel = state_container.init_new_c2s_virtual_connection(
                                &cnac,
                                kernel_ticket,
                                header.session_cid.get(),
                                session,
                            );
                            let session_security_settings = state_container
                                .session_security_settings
                                .expect("Should be set");
                            std::mem::drop(state_container);

                            session.session_cid.set(Some(cid)); // This makes is_provisional equal to false

                            let addr = session.remote_peer;
                            let is_personal = !session.is_server;

                            // Upgrade the connect BEFORE updating the CNAC
                            if !session
                                .session_manager
                                .upgrade_connection(session.remote_peer, cid)
                            {
                                return Ok(PrimaryProcessorResult::EndSession("Unable to upgrade from a provisional to a protected connection (Client)"));
                            }

                            log::trace!(target: "citadel", "The login to the server was a success. Welcome Message: {}", &message);

                            let post_login_object = payload.post_login_object.clone();
                            //session.post_quantum = pqc;
                            let cxn_type =
                                VirtualConnectionType::LocalGroupServer { session_cid: cid };
                            let peers = payload.peers;

                            let timestamp = session.time_tracker.get_global_time_ns();

                            let persistence_handler =
                                session.account_manager.get_persistence_handler().clone();
                            //session.session_manager.clear_provisional_tracker(session.kernel_ticket);
                            session.state.set(SessionState::Connected);

                            let success_ack = packet_crafter::do_connect::craft_success_ack(
                                &ratchet,
                                timestamp,
                                security_level,
                                ticket,
                            );
                            session.send_to_primary_stream(None, success_ack)?;

                            session.send_to_kernel(NodeResult::ConnectSuccess(ConnectSuccess {
                                ticket: kernel_ticket,
                                session_cid: cid,
                                remote_addr: addr,
                                is_personal,
                                v_conn_type: cxn_type,
                                services: payload.post_login_object,
                                welcome_message: message,
                                channel: channel.into(),
                                udp_rx_opt: udp_channel_rx,
                                session_security_settings,
                            }))?;
                            //finally, if there are any mailbox items, send them to the kernel for processing
                            if let Some(mailbox_delivery) = payload.mailbox {
                                session.send_to_kernel(NodeResult::MailboxDelivery(
                                    MailboxDelivery {
                                        session_cid: cid,
                                        ticket_opt: None,
                                        items: mailbox_delivery,
                                    },
                                ))?;
                            }
                            // TODO: Clean this up to prevent multiple saves
                            async move {
                                persistence_handler
                                    .synchronize_hyperlan_peer_list_as_client(&cnac, peers)
                                    .await?;
                                post_login_object
                                    .setup_client_rtdb(&cnac)
                                    .await
                                    .map_err(|err| NetworkError::generic(err.to_string()))?;

                                if let ConnectMode::Fetch { .. } = connect_mode {
                                    log::trace!(target: "citadel", "[FETCH] complete ...");
                                    // we can end the session now. The ratchets packets have already been sent alongside the connect signal above
                                    return Ok(PrimaryProcessorResult::EndSession(
                                        "Fetch succeeded",
                                    ));
                                }

                                if use_ka {
                                    let ka = packet_crafter::keep_alive::craft_keep_alive_packet(
                                        &ratchet,
                                        timestamp,
                                        security_level,
                                    );
                                    Ok(PrimaryProcessorResult::ReplyToSender(ka))
                                } else {
                                    log::warn!(target: "citadel", "Keep-alive subsystem will not be used for this session as requested");
                                    Ok(PrimaryProcessorResult::Void)
                                }
                            }
                        } else {
                            log::error!(target: "citadel", "An invalid SUCCESS packet was received; dropping due to invalid deserialization");
                            return Ok(PrimaryProcessorResult::Void);
                        }
                    } else {
                        log::error!(target: "citadel", "An invalid SUCCESS packet was received; dropping since the last local stage was not stage 1");
                        return Ok(PrimaryProcessorResult::Void);
                    }
                };

                task.await
            }

            packet_flags::cmd::aux::do_connect::SUCCESS_ACK => {
                log::trace!(target: "citadel", "RECV SUCCESS_ACK");
                if session.is_server {
                    let signal = inner_mut_state!(session.state_container)
                        .get_endpoint_container_mut(C2S_IDENTITY_CID)?
                        .channel_signal
                        .take()
                        .ok_or(error!(ErrorCode::ConnectChannelSignalMissing))?;
                    session.send_to_kernel(signal)?;
                    Ok(PrimaryProcessorResult::Void)
                } else {
                    Err(error!(ErrorCode::ConnectSuccessAckAsClient))
                }
            }

            n => {
                log::error!(target: "citadel", "Invalid auxiliary command: {n}");
                Ok(PrimaryProcessorResult::Void)
            }
        }
    };

    to_concurrent_processor!(task)
}
