botserver/src/directory/mod.rs

use crate::core::shared::state::AppState;
use axum::{
    extract::{Query, State},
    http::StatusCode,
    response::{IntoResponse, Json},
};
use log::{error, info};
use std::collections::HashMap;
use std::sync::Arc;
use uuid::Uuid;

pub mod auth_routes;
pub mod bootstrap;
pub mod client;
pub mod groups;
pub mod router;
pub mod users;

// Zitadel directory service integration - v4.13.1

pub use client::{ZitadelClient, ZitadelConfig};

pub struct AuthService {
    client: Arc<ZitadelClient>,
}

impl std::fmt::Debug for AuthService {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AuthService")
            .field("client", &"Arc<ZitadelClient>")
            .finish()
    }
}

impl AuthService {
    pub fn new(config: ZitadelConfig) -> anyhow::Result<Self> {
        let client = ZitadelClient::new(config)?;
        Ok(Self {
            client: Arc::new(client),
        })
    }

    pub fn client(&self) -> &ZitadelClient {
        &self.client
    }
}

pub async fn auth_handler(
    State(state): State<Arc<AppState>>,
    Query(params): Query<HashMap<String, String>>,
) -> impl IntoResponse {
    let bot_name = params.get("bot_name").cloned().unwrap_or_default();
    let existing_user_id = params
        .get("user_id")
        .and_then(|s| Uuid::parse_str(s).ok());
    let existing_session_id = params
        .get("session_id")
        .and_then(|s| Uuid::parse_str(s).ok());
    
    info!("Auth handler called: bot_name={}, existing_user_id={:?}, existing_session_id={:?}", 
          bot_name, existing_user_id, existing_session_id);

    let user_id = {
        let mut sm = state.session_manager.lock().await;
        match sm.get_or_create_anonymous_user(existing_user_id) {
            Ok(id) => id,
            Err(e) => {
                error!("Failed to create anonymous user: {}", e);
                return (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(serde_json::json!({ "error": "Failed to create user" })),
                );
            }
        }
    };

    let (bot_id, bot_name) = match tokio::task::spawn_blocking({
        let bot_name = bot_name.clone();
        let conn = state.conn.clone();
        move || {
            let mut db_conn = conn
                .get()
                .map_err(|e| format!("Failed to get database connection: {}", e))?;
            use crate::core::shared::models::schema::bots::dsl::*;
            use diesel::prelude::*;
            match bots
                .filter(name.eq(&bot_name))
                .filter(is_active.eq(true))
                .select((id, name))
                .first::<(Uuid, String)>(&mut db_conn)
                .optional()
            {
                Ok(Some((id_val, name_val))) => Ok((id_val, name_val)),
                Ok(None) => match bots
                    .filter(is_active.eq(true))
                    .select((id, name))
                    .first::<(Uuid, String)>(&mut db_conn)
                    .optional()
                {
                    Ok(Some((id_val, name_val))) => Ok((id_val, name_val)),
                    Ok(None) => Err("No active bots found".to_string()),
                    Err(e) => Err(format!("DB error: {}", e)),
                },
                Err(e) => Err(format!("DB error: {}", e)),
            }
        }
    })
    .await
    {
        Ok(Ok(res)) => res,
        Ok(Err(e)) => {
            error!("{}", e);
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": e })),
            );
        }
        Err(e) => {
            error!("Spawn blocking failed: {}", e);
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": "DB thread error" })),
            );
        }
    };

    let session = {
        let mut sm = state.session_manager.lock().await;
        
        // Try to get existing session by ID first
        if let Some(existing_session_id) = existing_session_id {
            info!("Attempting to get existing session: {}", existing_session_id);
            match sm.get_session_by_id(existing_session_id) {
                Ok(Some(sess)) => {
                    info!("Successfully retrieved existing session: {}", sess.id);
                    sess
                }
                Ok(None) => {
                    // Session not found, create a new one
                    info!("Session {} not found in database, creating new session", existing_session_id);
                    match sm.create_session(user_id, bot_id, "Auth Session") {
                        Ok(sess) => sess,
                        Err(e) => {
                            error!("Failed to create session: {}", e);
                            return (
                                StatusCode::INTERNAL_SERVER_ERROR,
                                Json(serde_json::json!({ "error": e.to_string() })),
                            );
                        }
                    }
                }
                Err(e) => {
                    // Error getting session, create a new one
                    error!("Error getting session {}: {}", existing_session_id, e);
                    match sm.create_session(user_id, bot_id, "Auth Session") {
                        Ok(sess) => sess,
                        Err(e) => {
                            error!("Failed to create session: {}", e);
                            return (
                                StatusCode::INTERNAL_SERVER_ERROR,
                                Json(serde_json::json!({ "error": e.to_string() })),
                            );
                        }
                    }
                }
            }
        } else {
            // No session_id provided, get or create session
            info!("No session_id provided, getting or creating session for user {}", user_id);
            match sm.get_or_create_user_session(user_id, bot_id, "Auth Session") {
                Ok(Some(sess)) => sess,
                Ok(None) => {
                    error!("Failed to create session");
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        Json(serde_json::json!({ "error": "Failed to create session" })),
                    );
                }
                Err(e) => {
                    error!("Failed to create session: {}", e);
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        Json(serde_json::json!({ "error": e.to_string() })),
                    );
                }
            }
        }
    };

    let auth_script_path = format!("./work/{}.gbai/{}.gbdialog/auth.ast", bot_name, bot_name);
    if tokio::fs::metadata(&auth_script_path).await.is_ok() {
        let auth_script = match tokio::fs::read_to_string(&auth_script_path).await {
            Ok(content) => content,
            Err(e) => {
                error!("Failed to read auth script: {}", e);
                return (
                    StatusCode::OK,
                    Json(serde_json::json!({
                        "user_id": session.user_id,
                        "session_id": session.id,
                        "bot_id": bot_id,
                        "bot_name": bot_name,
                        "status": "authenticated"
                    })),
                );
            }
        };

        let state_clone = Arc::clone(&state);
        let session_clone = session.clone();
        let bot_id = session.bot_id;
        match tokio::task::spawn_blocking(move || {
            let mut script_service =
                crate::basic::ScriptService::new(state_clone.clone(), session_clone);

script_service.load_bot_config_params(&state_clone, bot_id);

            match script_service.run(&auth_script) {
                Ok(_) => Ok(()),
                Err(e) => Err(format!("Script execution error: {}", e)),
            }
        })
        .await
        {
            Ok(Ok(())) => {}
            Ok(Err(e)) => {
                error!("Auth script error: {}", e);
            }
            Err(e) => {
                error!("Auth script task error: {}", e);
            }
        }
    }

    (
        StatusCode::OK,
        Json(serde_json::json!({
            "user_id": session.user_id,
            "session_id": session.id,
            "bot_id": bot_id,
            "bot_name": bot_name,
            "status": "authenticated"
        })),
    )
}