//! Registration Packet Processor for Citadel Protocol
//!
//! This module handles the registration process for new clients in the Citadel Protocol
//! network. It implements a secure multi-stage handshake that establishes client
//! identity and sets up initial cryptographic parameters.
//!
//! # Features
//!
//! - Multi-stage registration handshake
//! - Secure key exchange
//! - Passwordless registration support
//! - Session state management
//! - Cryptographic parameter negotiation
//! - Registration failure handling
//!
//! # Important Notes
//!
//! - Requires specific session states (NeedsRegister, SocketJustOpened, NeedsConnect)
//! - Supports both password-based and passwordless registration
//! - Implements post-quantum cryptography
//! - Manages registration state transitions
//! - Validates registration parameters
//!
//! # Related Components
//!
//! - `StateContainer`: Manages registration state
//! - `AccountManager`: Handles account creation
//! - `StackedRatchet`: Provides cryptographic primitives
//! - `SessionManager`: Tracks registration process

use super::includes::*;
use crate::error::NetworkError;
use crate::prelude::Ticket;
use crate::proto::misc::platform_ops::PlatformOps;
use crate::proto::node_result::{RegisterFailure, RegisterOkay};
use citadel_crypt::endpoint_crypto_container::{
    AssociatedCryptoParams, AssociatedSecurityLevel, EndpointRatchetConstructor, PeerSessionCrypto,
};
use citadel_crypt::prelude::{ConstructorOpts, Toolset};
use citadel_crypt::ratchets::Ratchet;
use citadel_io::{error, ErrorCode};
use citadel_user::serialization::SyncIO;

/// This will handle a registration packet
#[cfg_attr(feature = "localhost-testing", tracing::instrument(
    level = "trace",
    target = "citadel",
    skip_all,
    ret,
    err,
    fields(is_server = session_ref.is_server, src = packet.parse().unwrap().0.session_cid.get(), target = packet.parse().unwrap().0.target_cid.get()
    )
))]
pub async fn process_register<R: Ratchet, T: PlatformOps>(
    session_ref: &CitadelSession<R, T>,
    packet: HdpPacket,
    remote_addr: SocketAddr,
) -> Result<PrimaryProcessorResult, NetworkError> {
    let session = session_ref.clone();
    let state = session.state.get();

    if state != SessionState::NeedsRegister
        && state != SessionState::SocketJustOpened
        && state != SessionState::NeedsConnect
    {
        log::error!(target: "citadel", "Register packet received, but the system's state is not NeedsRegister. Dropping packet");
        return Ok(PrimaryProcessorResult::Void);
    }

    let task = async move {
        let (header, payload, _, _) = packet.decompose();
        let header = return_if_none!(Ref::new(&header[..]), "Unable to parse header")
            as Ref<&[u8], HdpHeader>;
        debug_assert_eq!(packet_flags::cmd::primary::DO_REGISTER, header.cmd_primary);
        let security_level = header.security_level.into();
        let ticket = Ticket::from(header.context_info.get());

        match header.cmd_aux {
            packet_flags::cmd::aux::do_register::STAGE0 => {
                log::trace!(target: "citadel", "STAGE 0 REGISTER PACKET");
                let algorithm = header.algorithm;
                // Validate incoming transfer first (no state lock required)
                match validation::do_register::validate_stage0::<R>(&payload) {
                    Some((transfer, transient_mode_requested)) => {
                        let timestamp = session.time_tracker.get_global_time_ns();
                        // Check server config for transient mode
                        if transient_mode_requested
                            && !session
                                .account_manager
                                .get_misc_settings()
                                .allow_transient_connections
                        {
                            const REASON: &str =
                                "Transient connections are not allowed on this node";
                            let err = packet_crafter::do_register::craft_failure(
                                algorithm,
                                timestamp,
                                REASON,
                                header.session_cid.get(),
                                ticket,
                            );
                            return Ok(PrimaryProcessorResult::EndSessionAndReplyToSender(
                                err, REASON,
                            ));
                        }

                        let session_password = session.session_password.clone();
                        let cid = header.session_cid.get();

                        // Heavy: new_bob + stage0_bob + finish (off-thread)
                        let (transfer_out, finished_ratchet) =
                            citadel_io::spawn_blocking(move || {
                                let mut bob_constructor =
                                    <R::Constructor as EndpointRatchetConstructor<R>>::new_bob(
                                        cid,
                                        ConstructorOpts::new_vec_init(
                                            Some(transfer.crypto_params()),
                                            transfer.security_level(),
                                        ),
                                        transfer,
                                        session_password.as_ref(),
                                    )
                                    .ok_or(error!(ErrorCode::RegisterBadBobTransfer))?;
                                let transfer_out = bob_constructor
                                    .stage0_bob()
                                    .ok_or(error!(ErrorCode::RegisterStage0BobFailed))?;
                                let finished_ratchet = bob_constructor
                                    .finish()
                                    .ok_or(error!(ErrorCode::RegisterBobConstructorFinishFailed))?;
                                Ok::<_, NetworkError>((transfer_out, finished_ratchet))
                            })
                            .await
                            .map_err(|err| error!(ErrorCode::TaskJoinFailed, err.to_string()))??;

                        // Build response packet before touching state
                        let stage1_packet = packet_crafter::do_register::craft_stage1::<R>(
                            algorithm,
                            timestamp,
                            transfer_out,
                            header.session_cid.get(),
                            ticket,
                        );

                        // Single lock for commit and state-machine validation
                        let ok = {
                            let mut state_container = inner_mut_state!(session.state_container);
                            if state_container.register_state.last_stage
                                == packet_flags::cmd::aux::do_register::STAGE0
                            {
                                state_container.register_state.transient_mode =
                                    Some(transient_mode_requested);
                                state_container.register_state.created_ratchet =
                                    Some(finished_ratchet);
                                state_container.register_state.last_stage =
                                    packet_flags::cmd::aux::do_register::STAGE1;
                                state_container.register_state.on_register_packet_received();
                                true
                            } else {
                                false
                            }
                        };

                        if !ok {
                            warn!(target: "citadel", "Inconsistency between the session's stage and the packet's state. Dropping");
                            return Ok(PrimaryProcessorResult::Void);
                        }

                        Ok(PrimaryProcessorResult::ReplyToSender(stage1_packet))
                    }

                    _ => {
                        log::error!(target: "citadel", "Unable to validate STAGE0_REGISTER packet");
                        let mut state_container = inner_mut_state!(session.state_container);
                        state_container.register_state.on_fail();
                        state_container.register_state.on_register_packet_received();
                        drop(state_container);
                        session.state.set(SessionState::NeedsRegister);
                        Ok(PrimaryProcessorResult::EndSession(
                            "Unable to validate STAGE0_REGISTER packet",
                        ))
                    }
                }
            }

            packet_flags::cmd::aux::do_register::STAGE1 => {
                log::trace!(target: "citadel", "STAGE 1 REGISTER PACKET");
                // Node is Alice. This packet will contain Bob's ciphertext; Alice will now be able to create the shared private key
                // Release state_container before heavy crypto
                let (algorithm, alice_constructor_opt, proposed_credentials_opt) = {
                    let algorithm = header.algorithm;
                    let mut state_container = inner_mut_state!(session.state_container);
                    let alice_constructor = state_container.register_state.constructor.take();
                    let proposed_credentials =
                        state_container.connect_state.proposed_credentials.clone();
                    (algorithm, alice_constructor, proposed_credentials)
                };

                if let Some(mut alice_constructor) = alice_constructor_opt {
                    let transfer: <R::Constructor as EndpointRatchetConstructor<R>>::BobToAliceWireTransfer = return_if_none!(
                        SyncIO::deserialize_from_vector(&payload[..]).ok(),
                        "Unable to deserialize BobToAliceTransfer"
                    );
                    let security_level = transfer.security_level();

                    // Offload stage1_alice + finish (avoid capturing session in closure)
                    let psk = session.session_password.clone();
                    let new_ratchet = citadel_io::spawn_blocking(move || {
                        alice_constructor
                            .stage1_alice(transfer, psk.as_ref())
                            .map_err(|err| NetworkError::generic(err.to_string()))?;
                        alice_constructor
                            .finish()
                            .ok_or(error!(ErrorCode::RegisterAliceConstructorFinishFailed))
                    })
                    .await
                    .map_err(|err| error!(ErrorCode::TaskJoinFailed, err.to_string()))??;

                    let timestamp = session.time_tracker.get_global_time_ns();
                    let proposed_credentials = return_if_none!(
                        proposed_credentials_opt.as_ref(),
                        "Unable to load proposed credentials"
                    );

                    let stage2_packet = packet_crafter::do_register::craft_stage2(
                        &new_ratchet,
                        algorithm,
                        timestamp,
                        proposed_credentials,
                        security_level,
                        ticket,
                    );

                    let mut state_container = inner_mut_state!(session.state_container);
                    state_container.register_state.created_ratchet = Some(new_ratchet);
                    state_container.register_state.last_stage =
                        packet_flags::cmd::aux::do_register::STAGE2;
                    state_container.register_state.on_register_packet_received();

                    Ok(PrimaryProcessorResult::ReplyToSender(stage2_packet))
                } else {
                    log::error!(target: "citadel", "Register stage is one, yet, no PQC is present. Aborting.");
                    Ok(PrimaryProcessorResult::Void)
                }
            }

            packet_flags::cmd::aux::do_register::STAGE2 => {
                log::trace!(target: "citadel", "STAGE 2 REGISTER PACKET");
                // Bob receives this packet. It contains the proposed credentials. We need to register and we're good to go

                let task = {
                    let state_container = inner_state!(session.state_container);
                    if state_container.register_state.last_stage
                        == packet_flags::cmd::aux::do_register::STAGE1
                    {
                        let algorithm = header.algorithm;
                        let ratchet = return_if_none!(
                            state_container.register_state.created_ratchet.clone(),
                            "Unable to load created hyper ratchet"
                        );
                        if let Some((stage2_packet, conn_info)) =
                            validation::do_register::validate_stage2(
                                &ratchet,
                                &header,
                                payload,
                                remote_addr,
                            )
                        {
                            let creds = stage2_packet.credentials;
                            let timestamp = session.time_tracker.get_global_time_ns();
                            let account_manager = session.account_manager.clone();
                            std::mem::drop(state_container);
                            let session_crypto_state = initialize_peer_session_crypto(
                                ratchet.get_cid(),
                                ratchet.clone(),
                                true,
                            );
                            // we must now create the CNAC
                            async move {
                                match account_manager
                                    .register_impersonal_hyperlan_client_network_account(
                                        conn_info,
                                        creds,
                                        session_crypto_state,
                                    )
                                    .await
                                {
                                    Ok(peer_cnac) => {
                                        log::trace!(target: "citadel", "Server successfully created a CNAC during the DO_REGISTER process! CID: {}", peer_cnac.get_cid());
                                        let success_message =
                                            session.create_register_success_message();

                                        let packet = packet_crafter::do_register::craft_success(
                                            &ratchet,
                                            algorithm,
                                            timestamp,
                                            success_message,
                                            security_level,
                                            ticket,
                                        );
                                        // Do not shutdown the session here. It is up to the client to decide
                                        // how to shutdown, or continue (in the case of passwordless mode), the session
                                        Ok(PrimaryProcessorResult::ReplyToSender(packet))
                                    }

                                    Err(err) => {
                                        let err = err.into_string();
                                        const REASON: &str = "Server unable to create client account during the DO_REGISTER process";
                                        log::error!(target: "citadel", "{REASON}: {err}");
                                        let packet = packet_crafter::do_register::craft_failure(
                                            algorithm,
                                            timestamp,
                                            err,
                                            header.session_cid.get(),
                                            ticket,
                                        );

                                        session.session_manager.clear_provisional_session(
                                            &remote_addr,
                                            session.init_time,
                                        );

                                        Ok(PrimaryProcessorResult::ReplyToSender(packet))
                                    }
                                }
                            }
                        } else {
                            log::error!(target: "citadel", "Unable to validate stage2 packet. Aborting");
                            return Ok(PrimaryProcessorResult::Void);
                        }
                    } else {
                        warn!(target: "citadel", "Inconsistency between the session's stage and the packet's state. Dropping");
                        return Ok(PrimaryProcessorResult::Void);
                    }
                };

                task.await
            }

            packet_flags::cmd::aux::do_register::SUCCESS => {
                log::trace!(target: "citadel", "STAGE SUCCESS REGISTER PACKET");
                // This will follow stage 4 in the case of a successful registration. The packet's payload contains the CNAC bytes, encrypted using AES-GCM.
                // The CNAC does not have the credentials (Serde skips the serialization thereof)

                let task = {
                    let state_container = inner_state!(session.state_container);
                    if state_container.register_state.last_stage
                        == packet_flags::cmd::aux::do_register::STAGE2
                    {
                        let ratchet = return_if_none!(
                            state_container.register_state.created_ratchet.clone(),
                            "Unable to load created hyper ratchet"
                        );

                        if let Some((success_message, conn_info)) =
                            validation::do_register::validate_success(
                                &ratchet,
                                &header,
                                payload,
                                remote_addr,
                            )
                        {
                            // Now, register the CNAC locally
                            let credentials = return_if_none!(
                                state_container.connect_state.proposed_credentials.clone(),
                                "Unable to take proposed credentials"
                            );

                            let passwordless = return_if_none!(
                                state_container.register_state.transient_mode,
                                "Passwordless unset (reg)"
                            );

                            drop(state_container);

                            let reg_ticket = session.kernel_ticket.clone();
                            let account_manager = session.account_manager.clone();
                            let kernel_tx = session.kernel_tx.clone();

                            let session_crypto_state =
                                initialize_peer_session_crypto(ratchet.get_cid(), ratchet, false);

                            async move {
                                match account_manager
                                    .register_personal_hyperlan_server(
                                        session_crypto_state,
                                        credentials,
                                        conn_info,
                                    )
                                    .await
                                {
                                    Ok(new_cnac) => {
                                        if passwordless {
                                            CitadelSession::begin_connect(&session, &new_cnac)?;
                                            inner_mut_state!(session.state_container).cnac =
                                                Some(new_cnac);
                                            // begin_connect will handle the connection process from here on out
                                            Ok(PrimaryProcessorResult::Void)
                                        } else {
                                            // Finally, alert the higher-level kernel about the success
                                            session.session_manager.clear_provisional_session(
                                                &remote_addr,
                                                session.init_time,
                                            );
                                            kernel_tx.unbounded_send(NodeResult::RegisterOkay(
                                                RegisterOkay {
                                                    ticket: reg_ticket.get(),
                                                    cid: new_cnac.get_cid(),
                                                    welcome_message: success_message,
                                                },
                                            ))?;
                                            session.shutdown();
                                            Ok(PrimaryProcessorResult::Void)
                                        }
                                    }

                                    Err(err) => {
                                        kernel_tx.unbounded_send(NodeResult::RegisterFailure(
                                            RegisterFailure {
                                                ticket: reg_ticket.get(),
                                                error_message: err.into_string(),
                                            },
                                        ))?;
                                        Ok(PrimaryProcessorResult::EndSession(
                                            "Registration subroutine ended (STATUS: ERR)",
                                        ))
                                    }
                                }
                            }
                        } else {
                            log::error!(target: "citadel", "Unable to validate SUCCESS packet");
                            return Ok(PrimaryProcessorResult::Void);
                        }
                    } else {
                        warn!(target: "citadel", "Inconsistency between the session's stage and the packet's state. Dropping");
                        return Ok(PrimaryProcessorResult::Void);
                    }
                };

                task.await
            }

            packet_flags::cmd::aux::do_register::FAILURE => {
                log::trace!(target: "citadel", "STAGE FAILURE REGISTER PACKET");
                // This node is again Bob. Alice received Bob's stage1 packet, but was unable to connect
                // A failure can be sent at any stage greater than the zeroth
                if inner_state!(session.state_container)
                    .register_state
                    .last_stage
                    > packet_flags::cmd::aux::do_register::STAGE0
                {
                    if let Some(error_message) =
                        validation::do_register::validate_failure(&header, &payload[..])
                    {
                        session.send_to_kernel(NodeResult::RegisterFailure(RegisterFailure {
                            ticket: session.kernel_ticket.get(),
                            error_message: String::from_utf8(error_message)
                                .unwrap_or_else(|_| "Non-UTF8 error message".to_string()),
                        }))?;
                        session.shutdown();
                    } else {
                        log::error!(target: "citadel", "Error validating FAILURE packet");
                        return Ok(PrimaryProcessorResult::Void);
                    }

                    Ok(PrimaryProcessorResult::EndSession(
                        "Registration subroutine ended (Status: FAIL)",
                    ))
                } else {
                    log::warn!(target: "citadel", "A failure packet was received, but the program's registration did not advance past stage 0. Dropping");
                    Ok(PrimaryProcessorResult::Void)
                }
            }

            _ => {
                warn!(target: "citadel", "Invalid auxiliary command. Dropping packet");
                Ok(PrimaryProcessorResult::Void)
            }
        }
    };

    to_concurrent_processor!(task)
}

// Only for registration; does not start the messenger/ratchet manager
fn initialize_peer_session_crypto<R: Ratchet>(
    cid: u64,
    initial_ratchet: R,
    is_server: bool,
) -> PeerSessionCrypto<R> {
    PeerSessionCrypto::new(Toolset::new(cid, initial_ratchet), !is_server)
}
