use crate::{
    load_keyring, log, log_cv_in, log_cv_out, log_err, log_in,
    server::short_link::add_short_link,
    sql::{
        connection_status::UserStatus,
        sql::{self, get_by_user_id, get_by_username, get_iota_by_id},
        user_online_tracker::{self},
    },
    transport::omikron_manager,
    util::{file_util::load_file_vec, logger::PrintType},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use dashmap::DashMap;
use mtp::host::{AuthenticationPolicy, Receiver, Sender};
use mtp::{
    codec::{CommunicationType, CommunicationValue, DataType, DataValue},
    host::{Host, HostConfig, Policy, SendMode},
};
use mtp_crypto::PublicKeyBundle;
use std::net::{IpAddr, Ipv4Addr};
use std::{
    sync::Arc,
    time::{Duration, Instant},
};
use tokio::{
    sync::{Mutex, RwLock},
    time::interval,
};
// ============================================================================
// Configuration
// ============================================================================

const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
const MAX_WAITING_AGE: Duration = Duration::from_secs(60);

// ============================================================================
// Error Types
// ============================================================================

#[derive(Debug, thiserror::Error)]
pub enum OmikronError {
    #[error("Not connected")]
    NotConnected,
    #[error("Not authenticated")]
    NotAuthenticated,
    #[error("Invalid response")]
    InvalidResponse,
    #[error("Authentication failed")]
    AuthenticationFailed,
    #[error("SQL error: {0}")]
    Sql(String),
    #[error("Send error: {0}")]
    Send(String),
}

pub type OmikronResult<T> = Result<T, OmikronError>;

// ============================================================================
// Waiting Task System (Preserved from original)
// ============================================================================

pub struct WaitingTask {
    pub task: Box<dyn Fn(Arc<OmikronConnection>, CommunicationValue) -> bool + Send + Sync>,
    pub inserted_at: Instant,
}

// ============================================================================
// Omikron Connection (mtp/QUIC-based)
// ============================================================================

pub struct OmikronConnection {
    id: u64,
    sender: Mutex<Option<Sender>>,
    pub ping: RwLock<i64>,
    waiting_tasks: DashMap<u32, WaitingTask>,
    cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
}

impl Drop for OmikronConnection {
    fn drop(&mut self) {
        if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
            handle.abort();
        }
    }
}

impl OmikronConnection {
    // -------------------------------------------------------------------------
    // Construction
    // -------------------------------------------------------------------------

    pub fn new(sender: Sender, id: u64) -> Arc<Self> {
        let conn = Arc::new(Self {
            id,
            sender: Mutex::new(Some(sender)),
            ping: RwLock::new(-1),
            waiting_tasks: DashMap::new(),
            cleanup_handle: std::sync::Mutex::new(None),
        });

        conn
    }

    // -------------------------------------------------------------------------
    // Main Handler Loop
    // -------------------------------------------------------------------------

    pub async fn handle(self: Arc<Self>, receiver: &mut Receiver) {
        log_in!(
            self.id as i64,
            PrintType::Omega,
            "Omikron connection started"
        );

        // Start cleanup task
        let cleanup_conn = self.clone();
        let cleanup_handle = tokio::spawn(async move {
            let mut ticker = interval(CLEANUP_INTERVAL);
            loop {
                ticker.tick().await;
                cleanup_conn
                    .waiting_tasks
                    .retain(|_, v| v.inserted_at.elapsed() < MAX_WAITING_AGE);
            }
        });
        *self.cleanup_handle.lock().unwrap() = Some(cleanup_handle);

        while let Ok(cv) = receiver.receive().await {
            if let Err(e) = self.clone().process_message(cv).await {
                log_err!(0, PrintType::Omega, "Error processing message: {}", e);
                if matches!(e, OmikronError::NotConnected) {
                    break;
                }
            }
        }

        self.clone().cleanup().await;
        log_in!(
            self.id as i64,
            PrintType::Omega,
            "Omikron connection closed"
        );
    }

    // -------------------------------------------------------------------------
    // Message Processing
    // -------------------------------------------------------------------------

    async fn process_message(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) {
            log_cv_in!(PrintType::Omikron, &cv);
        }

        let msg_id = cv.get_id();

        // Check waiting tasks first (response to previous request)
        if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) {
            let _ = (task.task)(self.clone(), cv);
            return Ok(());
        }

        // Handle ping regardless of message type
        if cv.is_type(CommunicationType::Ping) {
            return self.handle_ping(cv).await;
        }

        // Authentication is completed by the mtp host before this connection exists.
        let omikron_id = self.id as i64;
        self.clone().handle_authenticated(cv, omikron_id).await
    }

    async fn handle_authenticated(
        self: Arc<Self>,
        cv: CommunicationValue,
        omikron_id: i64,
    ) -> OmikronResult<()> {
        let comm_type = cv.get_comm_type_enum();
        match comm_type {
            // Link shortening
            Some(CommunicationType::ShortenLink) => self.handle_shorten_link(cv).await,

            // Online status tracking
            Some(CommunicationType::UserConnected) => {
                self.handle_user_connected(cv, omikron_id).await;
                Ok(())
            }
            Some(CommunicationType::UserDisconnected) => {
                self.handle_user_disconnected(cv, omikron_id).await;
                Ok(())
            }
            Some(CommunicationType::IotaConnected) => {
                self.handle_iota_connected(cv, omikron_id).await;
                Ok(())
            }
            Some(CommunicationType::IotaDisconnected) => {
                self.handle_iota_disconnected(cv, omikron_id).await;
                Ok(())
            }
            Some(CommunicationType::SyncClientIotaStatus) => {
                self.handle_sync_status(cv, omikron_id).await;
                Ok(())
            }

            Some(CommunicationType::GetUserData) => self.handle_get_user_data(cv).await,
            Some(CommunicationType::GetIotaData) => self.handle_get_iota_data(cv).await,

            Some(CommunicationType::GetRegister) => self.handle_get_register(cv).await,
            Some(CommunicationType::CompleteRegisterIota) => {
                self.handle_complete_register_iota(cv).await
            }
            Some(CommunicationType::CompleteRegisterUser) => {
                self.handle_complete_register_user(cv).await
            }

            Some(CommunicationType::ChangeUserData) => self.handle_change_user_data(cv).await,
            Some(CommunicationType::ChangeIotaData) => self.handle_change_iota_data(cv).await,
            Some(CommunicationType::DeleteUser) => self.handle_delete_user(cv).await,
            Some(CommunicationType::DeleteIota) => self.handle_delete_iota(cv).await,

            Some(CommunicationType::GetNotifications) => self.handle_get_notifications(cv).await,
            Some(CommunicationType::ReadNotification) => self.handle_read_notification(cv).await,
            Some(CommunicationType::PushNotification) => self.handle_push_notification(cv).await,
            Some(CommunicationType::GetStates) => self.handle_get_states(cv).await,

            _ => {
                log_err!(
                    0,
                    PrintType::Omega,
                    "Unknown message type: {:?}",
                    cv.get_type()
                );
                Ok(())
            }
        }
    }

    // -------------------------------------------------------------------------
    // Specific Handlers (ported from original WebSocket implementation)
    // -------------------------------------------------------------------------

    async fn handle_shorten_link(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let link = cv
            .get_data(DataType::Link)
            .as_str()
            .ok_or(OmikronError::InvalidResponse)?;

        let short = add_short_link(link)
            .await
            .map_err(|_| OmikronError::Sql("Shortend link Error".to_string()))?;

        let response = CommunicationValue::new(CommunicationType::ShortenLink)
            .with_id(cv.get_id())
            .add_typed_default(DataType::Link, DataValue::Str(short));

        self.send(&response).await
    }

    async fn handle_user_connected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
        log_in!(PrintType::Omega, "User connected");
        if let Some(user_id) = cv.get_data(DataType::UserId).as_number() {
            let status = cv
                .get_data(DataType::UserState)
                .as_str()
                .and_then(|s| UserStatus::from_str(s))
                .unwrap_or(UserStatus::user_online);
            user_online_tracker::track_user_status(user_id.try_into().unwrap(), status, omikron_id);
        }
    }

    async fn handle_user_disconnected(self: Arc<Self>, cv: CommunicationValue, _omikron_id: i64) {
        log_in!(PrintType::Omega, "User disconnected");
        if let Some(user_id) = cv.get_data(DataType::UserId).as_number() {
            if let Some(status) = user_online_tracker::get_user_status(user_id as i64) {
                user_online_tracker::track_user_status(
                    user_id as i64,
                    UserStatus::user_offline,
                    status.omikron_id,
                );
            }
        }
    }

    async fn handle_iota_connected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
        log_in!(PrintType::Omega, "IOTA connected");
        let iota_id = match cv.get_data(DataType::IotaId).as_number() {
            Some(id) => id as i64,
            None => return,
        };
        user_online_tracker::track_iota_connection(iota_id, omikron_id, true);

        let mut user_ids = Vec::new();
        if let Ok(users) = sql::get_users_by_iota_id(iota_id.try_into().unwrap()).await {
            for (user_id, _, _, _, _, _, _, _, _, _, _, _) in users {
                user_ids.push(DataValue::SignedNumber(user_id.try_into().unwrap()));
                user_online_tracker::track_user_status(
                    user_id.try_into().unwrap(),
                    UserStatus::user_offline,
                    omikron_id,
                );
            }
        } else {
            log_in!(PrintType::General, "SQL error loading users for IOTA");
        }

        let response = CommunicationValue::new(CommunicationType::IotaUserData)
            .with_id(cv.get_id())
            .add_typed_default(DataType::UserIds, DataValue::Array(user_ids));

        let _ = self.send(&response).await;
    }

    async fn handle_iota_disconnected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
        log_in!(PrintType::Omega, "IOTA disconnected");
        let iota_id = match cv.get_data(DataType::IotaId).as_number() {
            Some(id) => id as i64,
            None => return,
        };
        let iota_offline = user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
        if iota_offline {
            if let Ok(users) = sql::get_users_by_iota_id(iota_id.try_into().unwrap()).await {
                let user_ids: Vec<i64> = users.iter().map(|u| u.0.try_into().unwrap()).collect();
                user_online_tracker::untrack_many_users(&user_ids);
            }
        }
    }

    async fn handle_sync_status(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
        if let DataValue::Array(user_ids) = cv.get_data(DataType::UserIds) {
            for user_id_val in user_ids {
                if let DataValue::SignedNumber(user_id) = user_id_val {
                    user_online_tracker::track_user_status(
                        *user_id as i64,
                        UserStatus::user_offline,
                        omikron_id,
                    );
                }
            }
        }

        if let DataValue::Array(iota_ids) = cv.get_data(DataType::IotaIds) {
            for iota_id_val in iota_ids {
                if let DataValue::SignedNumber(iota_id) = iota_id_val {
                    user_online_tracker::track_iota_connection(*iota_id as i64, omikron_id, true);
                }
            }
        }
    }

    async fn handle_get_user_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        // Try by user_id first
        if let Some(user_id) = cv.get_data(DataType::UserId).as_number() {
            if let Ok(user_data) = get_by_user_id(user_id as i64).await {
                let response = self
                    .clone()
                    .build_user_data_response(cv.get_id(), user_data)
                    .await;
                return self.send(&response).await;
            }
        }

        // Try by username
        if let Some(username) = cv.get_data(DataType::Username).as_str() {
            if let Ok(user_data) = get_by_username(username).await {
                let response = self
                    .clone()
                    .build_user_data_response(cv.get_id(), user_data)
                    .await;
                return self.send(&response).await;
            }
        }

        // Not found
        let response =
            CommunicationValue::new(CommunicationType::ErrorNotFound).with_id(cv.get_id());
        self.send(&response).await
    }

    async fn build_user_data_response(
        self: Arc<Self>,
        msg_id: u32,
        user: (
            i64,
            i64,
            String,
            Option<String>,
            Option<String>,
            Option<String>,
            Option<Vec<u8>>,
            i32,
            i64,
            PublicKeyBundle,
            String,
            String,
        ),
    ) -> CommunicationValue {
        let (
            id,
            iota_id,
            username,
            display,
            status,
            about,
            avatar,
            sub_level,
            sub_end,
            public_key,
            _,
            _,
        ) = user;

        let mut response = CommunicationValue::new(CommunicationType::GetUserData)
            .with_id(msg_id)
            .add_typed_default(DataType::Username, DataValue::Str(username.clone()))
            .add_typed_default(DataType::PublicKey, DataValue::Str(public_key.to_base64()))
            .add_typed_default(DataType::UserId, DataValue::SignedNumber(id.into()))
            .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()))
            .add_typed_default(
                DataType::SubLevel,
                DataValue::SignedNumber(sub_level as i128),
            )
            .add_typed_default(DataType::SubEnd, DataValue::SignedNumber(sub_end.into()));

        // Display name (fallback to username)
        let display_name = display.filter(|d| !d.is_empty()).unwrap_or(username);
        response = response.add_typed_default(DataType::Display, DataValue::Str(display_name));

        // Optional fields
        if let Some(s) = status.filter(|s| !s.is_empty()) {
            response = response.add_typed_default(DataType::Status, DataValue::Str(s));
        }
        if let Some(a) = about.filter(|a| !a.is_empty()) {
            response = response.add_typed_default(DataType::About, DataValue::Str(a));
        }
        if let Some(av) = avatar {
            response =
                response.add_typed_default(DataType::Avatar, DataValue::Str(STANDARD.encode(av)));
        }

        // Online status
        let user_status = user_online_tracker::get_user_status(id);
        let iota_connections =
            user_online_tracker::get_iota_omikron_connections(iota_id).unwrap_or_default();

        if let Some(us) = user_status {
            let display_status = if us.connection_type == UserStatus::user_invisible {
                UserStatus::user_offline
            } else {
                us.connection_type.clone()
            };
            response = response.add_typed_default(
                DataType::OnlineStatus,
                DataValue::Str(display_status.to_string()),
            );
            response = response.add_typed_default(
                DataType::OmikronId,
                DataValue::SignedNumber(us.omikron_id.into()),
            );
        } else {
            response = response.add_typed_default(
                DataType::OnlineStatus,
                DataValue::Str(UserStatus::iota_offline.to_string()),
            );
        }

        response = response.add_typed_default(
            DataType::OmikronConnections,
            DataValue::Array(
                iota_connections
                    .into_iter()
                    .map(|id| DataValue::SignedNumber(id as i128))
                    .collect(),
            ),
        );

        response
    }

    async fn handle_get_iota_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        // Try by iota_id
        if let Some(iota_id) = cv.get_data(DataType::IotaId).as_number() {
            if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id as i64).await {
                let response = self
                    .clone()
                    .build_iota_data_response(cv.get_id(), iota_id, public_key, None, None)
                    .await;
                return self.send(&response).await;
            }
        }

        // Try by user_id
        if let Some(user_id) = cv.get_data(DataType::UserId).as_number() {
            if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
                get_by_user_id(user_id as i64).await
            {
                if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
                    let response = self
                        .clone()
                        .build_iota_data_response(
                            cv.get_id(),
                            iota_id,
                            public_key,
                            Some(user_id as i64),
                            None,
                        )
                        .await;
                    return self.send(&response).await;
                }
            }
        }

        // Try by username
        if let Some(username) = cv.get_data(DataType::Username).as_str() {
            if let Ok((user_id, iota_id, _, _, _, _, _, _, _, _, _, _)) =
                get_by_username(username).await
            {
                if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
                    let response = self
                        .clone()
                        .build_iota_data_response(
                            cv.get_id(),
                            iota_id,
                            public_key,
                            Some(user_id),
                            Some(username.to_string()),
                        )
                        .await;
                    return self.send(&response).await;
                }
            }
        }

        let response =
            CommunicationValue::new(CommunicationType::ErrorNotFound).with_id(cv.get_id());
        self.send(&response).await
    }

    async fn build_iota_data_response(
        self: Arc<Self>,
        msg_id: u32,
        iota_id: i64,
        public_key: PublicKeyBundle,
        user_id: Option<i64>,
        username: Option<String>,
    ) -> CommunicationValue {
        let mut response = CommunicationValue::new(CommunicationType::GetIotaData)
            .with_id(msg_id)
            .add_typed_default(DataType::PublicKey, DataValue::Str(public_key.to_base64()))
            .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()));

        if let Some(uid) = user_id {
            response =
                response.add_typed_default(DataType::UserId, DataValue::SignedNumber(uid.into()));
        }
        if let Some(uname) = username {
            response = response.add_typed_default(DataType::Username, DataValue::Str(uname));
        }

        let iota_connections =
            user_online_tracker::get_iota_omikron_connections(iota_id).unwrap_or_default();

        response.add_typed_default(
            DataType::OmikronConnections,
            DataValue::Array(
                iota_connections
                    .into_iter()
                    .map(|id| DataValue::SignedNumber(id as i128))
                    .collect(),
            ),
        )
    }

    async fn handle_get_register(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let register_id = sql::get_register_id().await;
        let response = CommunicationValue::new(CommunicationType::GetRegister)
            .with_id(cv.get_id())
            .add_typed_default(
                DataType::UserId,
                DataValue::SignedNumber(register_id as i128),
            );
        self.send(&response).await
    }

    async fn handle_complete_register_iota(
        self: Arc<Self>,
        cv: CommunicationValue,
    ) -> OmikronResult<()> {
        let iota_id_opt = cv.get_data(DataType::IotaId).as_number().map(|n| n as i64);

        let public_key = cv
            .get_data(DataType::PublicKey)
            .as_str()
            .and_then(|s| PublicKeyBundle::from_base64(s).ok());

        if let Some(public_key) = public_key {
            if let Some(iota_id) = iota_id_opt {
                // Register existing IOTA
                match sql::register_complete_iota(iota_id, public_key).await {
                    Ok(_) => {
                        let response = CommunicationValue::new(CommunicationType::Success)
                            .with_id(cv.get_id());
                        self.send(&response).await
                    }
                    Err(e) => {
                        let response = CommunicationValue::new(CommunicationType::ErrorInternal)
                            .with_id(cv.get_id())
                            .add_typed_default(DataType::ErrorType, DataValue::Str(e.to_string()));
                        self.send(&response).await
                    }
                }
            } else {
                // Create new IOTA
                match sql::create_new_iota(public_key).await {
                    Ok(new_iota_id) => {
                        let response =
                            CommunicationValue::new(CommunicationType::CompleteRegisterIota)
                                .with_id(cv.get_id())
                                .add_typed_default(
                                    DataType::IotaId,
                                    DataValue::SignedNumber(new_iota_id.into()),
                                );
                        self.send(&response).await
                    }
                    Err(e) => {
                        let response = CommunicationValue::new(CommunicationType::ErrorInternal)
                            .with_id(cv.get_id())
                            .add_typed_default(DataType::ErrorType, DataValue::Str(e.to_string()));
                        self.send(&response).await
                    }
                }
            }
        } else {
            self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
                .await
        }
    }

    async fn handle_complete_register_user(
        self: Arc<Self>,
        cv: CommunicationValue,
    ) -> OmikronResult<()> {
        let user_id = cv.get_data(DataType::UserId).as_number().map(|n| n as i64);
        let username = cv
            .get_data(DataType::Username)
            .as_str()
            .map(|s| s.to_string());
        let public_key = cv
            .get_data(DataType::PublicKey)
            .as_str()
            .and_then(|s| PublicKeyBundle::from_base64(s).ok());
        let iota_id = cv.get_sender();
        let reset_token = cv
            .get_data(DataType::ResetToken)
            .as_str()
            .map(|s| s.to_string());

        if let (Some(uid), Some(uname), Some(pk), Some(rt)) =
            (user_id, username, public_key, reset_token)
        {
            match sql::register_complete_user(uid, uname, pk, iota_id as i64, rt).await {
                Ok(_) => {
                    let response =
                        CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
                    self.send(&response).await
                }
                Err(e) => {
                    let response = CommunicationValue::new(CommunicationType::ErrorInternal)
                        .with_id(cv.get_id())
                        .add_typed_default(DataType::ErrorType, DataValue::Str(e.to_string()));
                    self.send(&response).await
                }
            }
        } else {
            self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
                .await
        }
    }

    async fn handle_change_user_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let user_id = cv.get_sender() as i64;
        let mut success = true;
        let mut error_message = String::new();

        // Process each field
        if let Some(username) = cv.get_data(DataType::Username).as_str() {
            if let Err(e) = sql::change_username(user_id, username.to_string()).await {
                success = false;
                error_message = e.to_string();
            }
        }
        if let Some(display) = cv.get_data(DataType::Display).as_str() {
            if let Err(e) = sql::change_display_name(user_id, display.to_string()).await {
                success = false;
                error_message = e.to_string();
            }
        }
        if let Some(avatar) = cv.get_data(DataType::Avatar).as_str() {
            if let Err(e) = sql::change_avatar(user_id, avatar.to_string()).await {
                success = false;
                error_message = e.to_string();
            }
        }
        if let Some(about) = cv.get_data(DataType::About).as_str() {
            if let Err(e) = sql::change_about(user_id, about.to_string()).await {
                success = false;
                error_message = e.to_string();
            }
        }
        if let Some(status) = cv.get_data(DataType::Status).as_str() {
            if let Err(e) = sql::change_status(user_id, status.to_string()).await {
                success = false;
                error_message = e.to_string();
            }
        }
        if let (Some(public_key), Some(private_key_hash)) = (
            cv.get_data(DataType::PublicKey)
                .as_str()
                .and_then(|s| PublicKeyBundle::from_base64(s).ok()),
            cv.get_data(DataType::PrivateKeyHash).as_str(),
        ) {
            if let Err(e) =
                sql::change_keys(user_id, public_key, private_key_hash.to_string()).await
            {
                success = false;
                error_message = e.to_string();
            }
        }

        if success {
            let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
            self.send(&response).await
        } else {
            let response = CommunicationValue::new(CommunicationType::ErrorInternal)
                .with_id(cv.get_id())
                .add_typed_default(DataType::ErrorType, DataValue::Str(error_message));
            self.send(&response).await
        }
    }

    async fn handle_change_iota_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let user_id = cv.get_sender() as i64;

        if let (iota_id, Some(reset_token), Some(new_token)) = (
            cv.get_sender(),
            cv.get_data(DataType::ResetToken).as_str(),
            cv.get_data(DataType::NewToken).as_str(),
        ) {
            match sql::get_by_user_id(user_id).await {
                Ok(user) => {
                    let current_token = user.11; // reset_token field
                    if current_token == reset_token {
                        let mut success = true;
                        let mut error_message = String::new();

                        if let Err(e) = sql::change_iota_id(user_id, iota_id as i64).await {
                            success = false;
                            error_message = e.to_string();
                        }
                        if success {
                            if let Err(e) = sql::change_token(user_id, new_token.to_string()).await
                            {
                                success = false;
                                error_message = e.to_string();
                            }
                        }

                        if success {
                            let response = CommunicationValue::new(CommunicationType::Success)
                                .with_id(cv.get_id());
                            self.send(&response).await
                        } else {
                            let response =
                                CommunicationValue::new(CommunicationType::ErrorInternal)
                                    .with_id(cv.get_id())
                                    .add_typed_default(
                                        DataType::ErrorType,
                                        DataValue::Str(error_message),
                                    );
                            self.send(&response).await
                        }
                    } else {
                        self.send_error_response(
                            cv.get_id(),
                            CommunicationType::ErrorInvalidChallenge,
                        )
                        .await
                    }
                }
                Err(_) => {
                    self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound)
                        .await
                }
            }
        } else {
            self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
                .await
        }
    }

    async fn handle_delete_user(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let user_id = cv.get_sender() as i64;
        match sql::delete_user(user_id).await {
            Ok(_) => {
                let response =
                    CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
                self.send(&response).await
            }
            Err(e) => {
                let response = CommunicationValue::new(CommunicationType::ErrorInternal)
                    .with_id(cv.get_id())
                    .add_typed_default(DataType::ErrorType, DataValue::Str(e.to_string()));
                self.send(&response).await
            }
        }
    }

    async fn handle_delete_iota(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let iota_id = cv.get_sender();
        match sql::delete_iota(iota_id as i64).await {
            Ok(_) => {
                let response =
                    CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
                self.send(&response).await
            }
            Err(e) => {
                let response = CommunicationValue::new(CommunicationType::ErrorInternal)
                    .with_id(cv.get_id())
                    .add_typed_default(DataType::ErrorType, DataValue::Str(e.to_string()));
                self.send(&response).await
            }
        }
    }

    async fn handle_get_notifications(
        self: Arc<Self>,
        cv: CommunicationValue,
    ) -> OmikronResult<()> {
        let user_id = cv.get_sender() as i64;
        let response_array = match sql::get_notifications(user_id).await {
            Ok(notifications) => notifications
                .into_iter()
                .map(|(sender, amount)| {
                    let tm = mtp::type_map::TypeMap::latest();
                    DataValue::Container(vec![
                        (
                            DataType::SenderId.to_id(&tm),
                            DataValue::SignedNumber(sender.into()),
                        ),
                        (
                            DataType::Amount.to_id(&tm),
                            DataValue::SignedNumber(amount.into()),
                        ),
                    ])
                })
                .collect(),
            Err(e) => {
                log!(PrintType::General, "SQL get_notifications error: {}", e);
                vec![]
            }
        };

        let response = CommunicationValue::new(CommunicationType::GetNotifications)
            .with_id(cv.get_id())
            .add_typed_default(DataType::Notifications, DataValue::Array(response_array));
        self.send(&response).await
    }

    async fn handle_read_notification(
        self: Arc<Self>,
        cv: CommunicationValue,
    ) -> OmikronResult<()> {
        let receiver_id = match cv.get_sender() {
            s if s > 0 => s as i64,
            _ => match cv.get_data(DataType::ReceiverId).as_number() {
                Some(id) => id as i64,
                None => return Ok(()),
            },
        };

        if let Some(other_id) = cv
            .get_data(DataType::SenderId)
            .as_number()
            .map(|n| n as i64)
        {
            if let Err(e) = sql::read_notification(receiver_id, other_id).await {
                log!(PrintType::General, "SQL read_notification error: {}", e);
            } else {
                let response = CommunicationValue::new(CommunicationType::ReadNotification)
                    .with_id(cv.get_id());
                let _ = self.send(&response).await;

                // Sync with other Omikron clients
                let sync_cv = CommunicationValue::new(CommunicationType::ReadNotification)
                    .with_receiver(receiver_id as u64)
                    .add_typed_default(
                        DataType::SenderId,
                        DataValue::SignedNumber(other_id.into()),
                    );
                crate::transport::omikron_manager::send_to_user(receiver_id, &sync_cv).await;
            }
        }
        Ok(())
    }

    async fn handle_push_notification(
        self: Arc<Self>,
        cv: CommunicationValue,
    ) -> OmikronResult<()> {
        let receiver_id = match cv.get_receiver() {
            r if r > 0 => r as i64,
            _ => match cv.get_data(DataType::ReceiverId).as_number() {
                Some(id) => id as i64,
                None => return Ok(()),
            },
        };

        let sender_id = match cv.get_data(DataType::SenderId).as_number() {
            Some(id) => id as i64,
            None => cv.get_sender() as i64,
        };

        if let Err(e) = sql::add_notification(receiver_id, sender_id).await {
            log!(PrintType::General, "SQL add_notification error: {}", e);
        } else {
            let response =
                CommunicationValue::new(CommunicationType::PushNotification).with_id(cv.get_id());
            let _ = self.send(&response).await;

            // Sync with other Omikron clients
            let push_cv = CommunicationValue::new(CommunicationType::PushNotification)
                .with_receiver(receiver_id as u64)
                .add_typed_default(
                    DataType::SenderId,
                    DataValue::SignedNumber(sender_id.into()),
                );
            crate::transport::omikron_manager::send_to_user(receiver_id, &push_cv).await;
        }
        Ok(())
    }

    async fn handle_get_states(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        let user_ids = match cv.get_data(DataType::UserIds) {
            DataValue::Array(ids) => ids,
            _ => return Ok(()),
        };

        let mut states = Vec::new();
        for id_val in user_ids {
            if let DataValue::SignedNumber(user_id) = id_val {
                let user_id = *user_id as i64;
                let status = user_online_tracker::get_user_status(user_id);
                let status_str = match status {
                    Some(ref us) => {
                        if us.connection_type == UserStatus::user_invisible {
                            "user_offline".to_string()
                        } else {
                            us.connection_type.to_string()
                        }
                    }
                    None => UserStatus::iota_offline.to_string(),
                };
                let tm = mtp::type_map::TypeMap::latest();
                let mut map = Vec::new();
                map.push((
                    DataType::UserId.to_id(&tm),
                    DataValue::SignedNumber(user_id.into()),
                ));
                map.push((DataType::UserState.to_id(&tm), DataValue::Str(status_str)));
                states.push(DataValue::Container(map));
            }
        }

        let response = CommunicationValue::new(CommunicationType::GetStates)
            .with_id(cv.get_id())
            .add_typed_default(DataType::UserStates, DataValue::Array(states));
        self.send(&response).await
    }

    async fn handle_ping(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
        if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) {
            *self.ping.write().await = *last_ping as i64;
        }

        let response = CommunicationValue::new(CommunicationType::Pong).with_id(cv.get_id());
        self.send(&response).await
    }

    // -------------------------------------------------------------------------
    // Utilities
    // -------------------------------------------------------------------------

    async fn send(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
        if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) {
            log_cv_out!(PrintType::Omikron, cv);
        }

        let guard = self.sender.lock().await;
        let sender = guard.as_ref().ok_or(OmikronError::NotConnected)?;

        sender
            .send(cv)
            .await
            .map_err(|e| OmikronError::Send(e.to_string()))
    }

    async fn send_error_response(
        self: Arc<Self>,
        message_id: u32,
        error_type: CommunicationType,
    ) -> OmikronResult<()> {
        let error = CommunicationValue::new(error_type).with_id(message_id);
        self.send(&error).await
    }

    pub async fn close(self: Arc<Self>) {
        log_in!(
            self.get_omikron_id().await.unwrap_or(0),
            PrintType::Omega,
            "Omikron connection Closed"
        );
    }

    async fn cleanup(self: Arc<Self>) {
        if self.id != 0 {
            log_in!(self.id as i64, PrintType::Omega, "Omikron disconnected");
            omikron_manager::remove_omikron(self.id as i64).await;
            user_online_tracker::untrack_omikron(self.id as i64).await;
        }

        if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
            handle.abort();
        }
    }

    pub async fn get_omikron_id(self: Arc<Self>) -> Option<i64> {
        Some(self.id as i64)
    }

    pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
        self.send(cv).await
    }
}

// ============================================================================
// Server Startup
// ============================================================================

pub async fn get_by_omikron_id(
    omikron_id: u64,
    _description: Option<String>,
) -> Option<PublicKeyBundle> {
    sql::get_omikron_by_id(omikron_id as i64)
        .await
        .ok()
        .map(|(bundle, _ip_address)| bundle)
}
pub async fn complete_register(_pub_key: PublicKeyBundle, _description: Option<String>) -> u64 {
    0
}

pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
    let cert_pem = load_file_vec("certs", "transport_cert.pem").expect("Error loading Pemfile");

    let key_pem = load_file_vec("certs", "transport_key.pem").expect("Error loading Keyfile");

    let host_config = HostConfig::new(
        IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)),
        port,
        cert_pem,
        key_pem,
    )
    .with_policy(Policy {
        send_mode: SendMode::SingleStreamPerMessage,
        max_message_size: 1_000_000_000,
        close_frame_len: u32::MAX,
        application_close_code: 0,
        open_stream_timeout: Duration::from_millis(2_000),
        write_timeout: Duration::from_millis(2_000),
        accept_stream_timeout: Duration::from_millis(10_000),
        read_timeout: Duration::from_millis(30_000),
        keep_alive_interval: Some(Duration::from_secs(6)),
        max_idle_timeout: Some(Duration::from_secs(30)),
        force_close_delay: Duration::from_millis(300),
        max_transient_recv_errors: 20,
        transient_recv_backoff: Duration::from_millis(100),
        receiver_queue_capacity: 1000,
    })
    .with_authentication(
        load_keyring(),
        Box::new(|user_id, description| Box::pin(get_by_omikron_id(user_id, description))),
        Box::new(|pub_key, description| Box::pin(complete_register(pub_key, description))),
    )
    .with_authentication_policy(AuthenticationPolicy::ForceAuthentication);

    let mut host: Host = Host::new(host_config).await?;
    log!("OmikronServer listening on port {}", port);

    while let Ok(Some(mut connection)) = host.accept().await {
        tokio::spawn(async move {
            let conn = OmikronConnection::new(connection.sender, connection.client_id);
            omikron_manager::add_omikron(conn.clone()).await;
            conn.handle(&mut connection.receiver).await;
        });
    }

    Ok(())
}