#[cfg(any(feature = "research", feature = "llm"))]
pub mod kb_context;
#[cfg(any(feature = "research", feature = "llm"))]
use kb_context::inject_kb_context;
pub mod tool_context;
use tool_context::get_session_tools;
pub mod tool_executor;
use tool_executor::ToolExecutor;
#[cfg(feature = "llm")]
use crate::core::config::ConfigManager;

#[cfg(feature = "drive")]
use crate::drive::drive_monitor::DriveMonitor;
#[cfg(feature = "llm")]
use crate::llm::llm_models;
#[cfg(feature = "llm")]
use crate::llm::OpenAIClient;
#[cfg(feature = "nvidia")]
use crate::nvidia::get_system_metrics;
use crate::core::shared::message_types::MessageType;
use crate::core::shared::models::{BotResponse, UserMessage, UserSession};
use crate::core::shared::state::AppState;
#[cfg(feature = "chat")]
use crate::basic::keywords::add_suggestion::get_suggestions;
use axum::extract::ws::{Message, WebSocket};
use axum::{
    extract::{ws::WebSocketUpgrade, Extension, Path, Query, State},
    http::StatusCode,
    response::{IntoResponse, Json},
};
use diesel::ExpressionMethods;
use diesel::PgConnection;
use diesel::QueryDsl;
use diesel::RunQueryDsl;
use futures::{sink::SinkExt, stream::StreamExt};
#[cfg(feature = "llm")]
use log::trace;
use log::{error, info, warn};
use serde_json;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::sync::Mutex as AsyncMutex;
use uuid::Uuid;
use serde::{Deserialize, Serialize};

pub mod channels;
pub mod multimedia;

pub fn get_default_bot(conn: &mut PgConnection) -> (Uuid, String) {
    use crate::core::shared::models::schema::bots::dsl::*;
    use diesel::prelude::*;

    // First try to get the bot named "default"
    match bots
        .filter(name.eq("default"))
        .filter(is_active.eq(true))
        .select((id, name))
        .first::<(Uuid, String)>(conn)
        .optional()
    {
        Ok(Some((bot_id, bot_name))) => (bot_id, bot_name),
        Ok(None) => {
            warn!("Bot named 'default' not found, falling back to first active bot");
            // Fall back to first active bot
            match bots
                .filter(is_active.eq(true))
                .select((id, name))
                .first::<(Uuid, String)>(conn)
                .optional()
            {
                Ok(Some((bot_id, bot_name))) => (bot_id, bot_name),
                Ok(None) => {
                    warn!("No active bots found, using nil UUID");
                    (Uuid::nil(), "default".to_string())
                }
                Err(e) => {
                    error!("Failed to query fallback bot: {}", e);
                    (Uuid::nil(), "default".to_string())
                }
            }
        }
        Err(e) => {
            error!("Failed to query default bot: {}", e);
            (Uuid::nil(), "default".to_string())
        }
    }
}

/// Get bot ID by name from database
pub fn get_bot_id_by_name(conn: &mut PgConnection, bot_name: &str) -> Result<Uuid, String> {
    use crate::core::shared::models::schema::bots::dsl::*;
    use diesel::prelude::*;

    bots
        .filter(name.eq(bot_name))
        .select(id)
        .first::<Uuid>(conn)
        .map_err(|e| format!("Bot '{}' not found: {}", bot_name, e))
}

#[derive(Debug)]
pub struct BotOrchestrator {
    pub state: Arc<AppState>,
    pub mounted_bots: Arc<AsyncMutex<HashMap<String, Arc<DriveMonitor>>>>,
}

#[derive(Debug, Deserialize)]
pub struct BotConfigQuery {
    pub bot_name: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct BotConfigResponse {
    pub public: bool,
    pub theme_color1: Option<String>,
    pub theme_color2: Option<String>,
    pub theme_title: Option<String>,
    pub theme_logo: Option<String>,
    pub theme_logo_text: Option<String>,
}

/// Get bot configuration endpoint
/// Returns bot's public setting and other configuration values
pub async fn get_bot_config(
    Query(params): Query<BotConfigQuery>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<BotConfigResponse>, StatusCode> {
    let bot_name = params.bot_name.unwrap_or_else(|| "default".to_string());

    let mut conn = match state.conn.get() {
        Ok(c) => c,
        Err(e) => {
            error!("Failed to get database connection: {}", e);
            return Err(StatusCode::INTERNAL_SERVER_ERROR);
        }
    };

    // Query bot_configuration table for this bot's configuration
    use crate::core::shared::models::schema::bot_configuration::dsl::*;

    let mut is_public = false;
    let mut theme_color1: Option<String> = None;
    let mut theme_color2: Option<String> = None;
    let mut theme_title: Option<String> = None;
    let mut theme_logo: Option<String> = None;
    let mut theme_logo_text: Option<String> = None;

    let target_bot_id = match get_bot_id_by_name(&mut conn, &bot_name) {
        Ok(found_id) => found_id,
        Err(e) => {
            warn!("Failed to find bot ID for name '{}': {}", bot_name, e);
            return Ok(Json(BotConfigResponse {
                public: false,
                theme_color1: None,
                theme_color2: None,
                theme_title: None,
                theme_logo: None,
                theme_logo_text: None,
            }));
        }
    };

    // Query all config values for this specific bot
    match bot_configuration
        .filter(bot_id.eq(target_bot_id))
        .select((config_key, config_value))
        .load::<(String, String)>(&mut conn)
    {
        Ok(configs) => {
            info!("Config query returned {} entries for bot '{}'", configs.len(), bot_name);
            for (key, value) in configs {
                // Try to strip bot_name prefix, use original if no prefix
                let clean_key = key.strip_prefix(&format!("{}.", bot_name))
                    .or_else(|| key.strip_prefix(&format!("{}_", bot_name)))
                    .unwrap_or(&key);

                // Check if key is for this bot (either prefixed or not)
                let key_for_bot = clean_key == key || key.starts_with(&format!("{}.", bot_name)) || key.starts_with(&format!("{}_", bot_name));

                info!("Key '{}' -> clean_key '{}' -> key_for_bot: {}", key, clean_key, key_for_bot);

                if !key_for_bot {
                    info!("Skipping key '{}' - not for bot '{}'", key, bot_name);
                    continue;
                }

                match clean_key.to_lowercase().as_str() {
                    "public" => {
                        is_public = value.eq_ignore_ascii_case("true") || value == "1";
                    }
                    "theme-color1" => {
                        theme_color1 = Some(value);
                    }
                    "theme-color2" => {
                        theme_color2 = Some(value);
                    }
                    "theme-title" => {
                        theme_title = Some(value);
                    }
                    "theme-logo" => {
                        theme_logo = Some(value);
                    }
                    "theme-logo-text" => {
                        theme_logo_text = Some(value);
                    }
                    _ => {}
                }
            }
            info!("Retrieved config for bot '{}': public={}, theme_color1={:?}, theme_color2={:?}, theme_title={:?}",
                bot_name, is_public, theme_color1, theme_color2, theme_title);
        }
        Err(e) => {
            warn!("Failed to load config for bot '{}': {}", bot_name, e);
            // Return defaults (not public, no theme)
        }
    }

    let config_response = BotConfigResponse {
        public: is_public,
        theme_color1,
        theme_color2,
        theme_title,
        theme_logo,
        theme_logo_text,
    };

    Ok(Json(config_response))
}

impl BotOrchestrator {
    pub fn new(state: Arc<AppState>) -> Self {
        Self {
            state,
            mounted_bots: Arc::new(AsyncMutex::new(HashMap::new())),
        }
    }

    pub fn mount_all_bots(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        info!("Scanning drive for .gbai files to mount bots...");

        let mut bots_mounted = 0;
        let mut bots_created = 0;

        let data_dir = "/opt/gbo/data";

        let directories_to_scan: Vec<std::path::PathBuf> = vec![
            self.state
                .config
                .as_ref()
                .map(|c| c.site_path.clone())
                .unwrap_or_else(|| format!("{}/sites", crate::core::shared::utils::get_stack_path()))
                .into(),
            "./templates".into(),
            "../bottemplates".into(),
            data_dir.into(),
        ];

        for dir_path in directories_to_scan {
            info!("Checking directory for bots: {}", dir_path.display());

            if !dir_path.exists() {
                info!("Directory does not exist, skipping: {}", dir_path.display());
                continue;
            }

            match self.scan_directory(&dir_path, &mut bots_mounted, &mut bots_created) {
                Ok(()) => {}
                Err(e) => {
                    error!("Failed to scan directory {}: {}", dir_path.display(), e);
                }
            }
        }

        info!(
            "BotServer ready - {} bots loaded",
            bots_mounted
        );
        log::debug!(
            "Bot mounting details: {} created, {} already existed",
            bots_created,
            bots_mounted - bots_created
        );

        Ok(())
    }

    fn scan_directory(
        &self,
        dir_path: &std::path::Path,
        bots_mounted: &mut i32,
        bots_created: &mut i32,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let entries =
            std::fs::read_dir(dir_path).map_err(|e| format!("Failed to read directory: {}", e))?;

        for entry in entries.flatten() {
            let name = entry.file_name();

            let bot_name = match name.to_str() {
                Some(n) if n.ends_with(".gbai") => n.trim_end_matches(".gbai"),
                _ => continue,
            };

            info!("Found .gbai file: {}", bot_name);

            match self.ensure_bot_exists(bot_name) {
                Ok(true) => {
                    info!("Bot '{}' already exists in database, mounting", bot_name);
                    *bots_mounted += 1;
                }
                Ok(false) => {
                    // Auto-create bots found in /opt/gbo/data
                    if dir_path.to_string_lossy().contains("/data") {
                        info!("Auto-creating bot '{}' from /opt/gbo/data", bot_name);
                        match self.create_bot_simple(bot_name) {
                            Ok(_) => {
                                info!("Bot '{}' created successfully", bot_name);
                                *bots_created += 1;
                                *bots_mounted += 1;
                            }
                            Err(e) => {
                                error!("Failed to create bot '{}': {}", bot_name, e);
                            }
                        }
                    } else {
                        info!(
                            "Bot '{}' does not exist in database, skipping (run import to create)",
                            bot_name
                        );
                    }
                }
                Err(e) => {
                    error!("Failed to check if bot '{}' exists: {}", bot_name, e);
                }
            }
        }

        Ok(())
    }

    fn ensure_bot_exists(
        &self,
        bot_name: &str,
    ) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
        use diesel::sql_query;

        let mut conn = self
            .state
            .conn
            .get()
            .map_err(|e| format!("Failed to get database connection: {e}"))?;

        #[derive(diesel::QueryableByName)]
        #[diesel(check_for_backend(diesel::pg::Pg))]
        struct BotExistsResult {
            #[diesel(sql_type = diesel::sql_types::Bool)]
            exists: bool,
        }

        let exists: BotExistsResult = sql_query(
            "SELECT EXISTS(SELECT 1 FROM bots WHERE name = $1 AND is_active = true) as exists",
        )
        .bind::<diesel::sql_types::Text, _>(bot_name)
        .get_result(&mut conn)
        .map_err(|e| format!("Failed to check if bot exists: {e}"))?;

        Ok(exists.exists)
    }

    fn create_bot_simple(&self, bot_name: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        use diesel::sql_query;
        use uuid::Uuid;

        let mut conn = self
            .state
            .conn
            .get()
            .map_err(|e| format!("Failed to get database connection: {e}"))?;

        // Check if bot already exists
        let exists = self.ensure_bot_exists(bot_name)?;
        if exists {
            info!("Bot '{}' already exists, skipping creation", bot_name);
            return Ok(());
        }

        // Ensure default tenant exists
        sql_query(
            "INSERT INTO tenants (id, name, slug, created_at) \
             VALUES ('00000000-0000-0000-0000-000000000001', 'Default Tenant', 'default', NOW()) \
             ON CONFLICT (slug) DO NOTHING"
        )
        .execute(&mut conn)
        .ok();

        // Ensure default organization exists (with default slug or use existing)
        sql_query(
            "INSERT INTO organizations (org_id, tenant_id, name, slug, created_at) \
             VALUES ('00000000-0000-0000-0000-000000000001', '00000000-0000-0000-0000-000000000001', 'Default Organization', 'default', NOW()) \
             ON CONFLICT (slug) DO NOTHING"
        )
        .execute(&mut conn)
        .ok();

        // Get default organization by slug
        #[derive(diesel::QueryableByName)]
        #[diesel(check_for_backend(diesel::pg::Pg))]
        struct OrgResult {
            #[diesel(sql_type = diesel::sql_types::Uuid)]
            org_id: uuid::Uuid,
        }

        let org_result: OrgResult = sql_query("SELECT org_id FROM organizations WHERE slug = 'default' LIMIT 1")
            .get_result(&mut conn)
            .map_err(|e| format!("Failed to get default organization: {e}"))?;
        let org_id = org_result.org_id.to_string();

        let bot_id = Uuid::new_v4();

        sql_query(
            "INSERT INTO bots (id, org_id, name, llm_provider, context_provider, is_active, created_at, updated_at)
             VALUES ($1, $2::uuid, $3, 'openai', 'website', true, NOW(), NOW())"
        )
        .bind::<diesel::sql_types::Uuid, _>(bot_id)
        .bind::<diesel::sql_types::Text, _>(org_id.clone())
        .bind::<diesel::sql_types::Text, _>(bot_name)
        .execute(&mut conn)
        .map_err(|e| format!("Failed to create bot: {e}"))?;

        info!("User system created resource: bot {} with org_id {}", bot_id, org_id);
        Ok(())
    }

    #[cfg(feature = "llm")]
    pub async fn stream_response(
        &self,
        message: UserMessage,
        response_tx: mpsc::Sender<BotResponse>,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        trace!(
            "Streaming response for user: {}, session: {}",
            message.user_id,
            message.session_id
        );

        let user_id = Uuid::parse_str(&message.user_id)?;
        let session_id = Uuid::parse_str(&message.session_id)?;
        let message_content = message.content.clone();

        // If a HEAR is blocking the script thread for this session, deliver the input
        // directly and return — the script continues from where it paused.
        if crate::basic::keywords::hearing::deliver_hear_input(
            &self.state,
            session_id,
            message_content.clone(),
        ) {
            trace!("HEAR: delivered input to blocking script for session {session_id}");
            return Ok(());
        }

        let (session, context_data, history, model, key, system_prompt, bot_llm_url, explicit_llm_provider) = {
            let state_clone = self.state.clone();
            tokio::task::spawn_blocking(
                move || -> Result<_, Box<dyn std::error::Error + Send + Sync>> {
                    let mut session = {
                        let mut sm = state_clone.session_manager.blocking_lock();
                        sm.get_session_by_id(session_id)?
                    }
                    .ok_or("Session not found")?;

                    // Store WebSocket session_id in context for TALK routing
                    if let serde_json::Value::Object(ref mut map) = session.context_data {
                        map.insert("websocket_session_id".to_string(), serde_json::Value::String(session_id.to_string()));
                    } else {
                        let mut map = serde_json::Map::new();
                        map.insert("websocket_session_id".to_string(), serde_json::Value::String(session_id.to_string()));
                        session.context_data = serde_json::Value::Object(map);
                    }

                    if !message.content.trim().is_empty() {
                        let mut sm = state_clone.session_manager.blocking_lock();
                        sm.save_message(session.id, user_id, 1, &message.content, 1)?;
                    }

                    let context_data = {
                        let sm = state_clone.session_manager.blocking_lock();
                        sm.get_session_context_data(&session.id, &session.user_id)?
                    };

                    let history = {
                        let mut sm = state_clone.session_manager.blocking_lock();
                        sm.get_conversation_history(session.id, user_id)?
                    };

                    let config_manager = ConfigManager::new(state_clone.conn.clone());

                    // For local LLM server, use the actual model name
                    // Default to DeepSeek model if not configured
                    let model = config_manager
                        .get_config(&session.bot_id, "llm-model", Some("DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf"))
                        .unwrap_or_else(|_| "DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf".to_string());

                    let key = config_manager
                        .get_config(&session.bot_id, "llm-key", Some(""))
                        .unwrap_or_default();

                    // Load bot-specific llm-url (may differ from global default)
                    let bot_llm_url = config_manager
                        .get_bot_config_value(&session.bot_id, "llm-url")
                        .ok();

                    // Load explicit llm-provider from config.csv (e.g., "openai", "bedrock", "claude")
                    // This allows overriding auto-detection from URL
                    let explicit_llm_provider = config_manager
                        .get_bot_config_value(&session.bot_id, "llm-provider")
                        .ok();

                    // Load system-prompt from config.csv, fallback to default
                    let system_prompt = config_manager
                        .get_config(&session.bot_id, "system-prompt", Some("You are a helpful assistant with access to tools that can help you complete tasks. When a user's request matches one of your available tools, use the appropriate tool instead of providing a generic response."))
                        .unwrap_or_else(|_| "You are a helpful General Bots assistant.".to_string());

                    info!("Loaded system-prompt for bot {}: {}", session.bot_id, &system_prompt[..system_prompt.len().min(500)]);

                    Ok((session, context_data, history, model, key, system_prompt, bot_llm_url, explicit_llm_provider))
                },
            )
            .await??
        };

        let mut messages = OpenAIClient::build_messages(&system_prompt, &context_data, &history);

        trace!("Built messages array with {} items, first message role: {:?}",
            messages.as_array().map(|a| a.len()).unwrap_or(0),
            messages.as_array().and_then(|a| a.first()).and_then(|m| m.get("role")));

        // Get bot name for KB and tool injection
        let bot_name_for_context = {
            let conn = self.state.conn.get().ok();
            if let Some(mut db_conn) = conn {
                use crate::core::shared::models::schema::bots::dsl::*;
                bots.filter(id.eq(session.bot_id))
                    .select(name)
                    .first::<String>(&mut db_conn)
                    .unwrap_or_else(|_| "default".to_string())
            } else {
                "default".to_string()
            }
        };

        #[cfg(any(feature = "research", feature = "llm"))]
        {
            // Execute start.bas on first message - ONLY run once per session to load suggestions
            let actual_session_id = session.id.to_string();

            // Check if start.bas has already been executed for this session
            let start_bas_key = format!("start_bas_executed:{}", actual_session_id);
            let should_execute_start_bas = if let Some(cache) = &self.state.cache {
                if let Ok(mut conn) = cache.get_multiplexed_async_connection().await {
                    let executed: Result<Option<String>, redis::RedisError> = redis::cmd("GET")
                        .arg(&start_bas_key)
                        .query_async(&mut conn)
                        .await;
                    matches!(executed, Ok(None))
                } else {
                    true // If cache fails, try to execute
                }
            } else {
                true // If no cache, try to execute
            };

            if should_execute_start_bas {
                // Always execute start.bas for this session (blocking - wait for completion)
                let data_dir = "/opt/gbo/data";
                let start_script_path = format!("{}/{}.gbai/{}.gbdialog/start.bas", data_dir, bot_name_for_context, bot_name_for_context);

                trace!("Executing start.bas for session {} at: {}", actual_session_id, start_script_path);

                if let Ok(metadata) = tokio::fs::metadata(&start_script_path).await {
                    if metadata.is_file() {
                        if let Ok(start_script) = tokio::fs::read_to_string(&start_script_path).await {
                            let state_clone = self.state.clone();
                            let actual_session_id_for_task = session.id;
                            let bot_id_clone = session.bot_id;

                            // Execute start.bas synchronously (blocking)
                            let result = tokio::task::spawn_blocking(move || {
                                let session_result = {
                                    let mut sm = state_clone.session_manager.blocking_lock();
                                    sm.get_session_by_id(actual_session_id_for_task)
                                };

                                let sess = match session_result {
                                    Ok(Some(s)) => s,
                                    Ok(None) => {
                                        return Err(format!("Session {} not found during start.bas execution", actual_session_id_for_task));
                                    }
                                    Err(e) => return Err(format!("Failed to get session: {}", e)),
                                };

                                let mut script_service = crate::basic::ScriptService::new(
                                    state_clone.clone(),
                                    sess
                                );
                                script_service.load_bot_config_params(&state_clone, bot_id_clone);

                                match script_service.compile(&start_script) {
                                    Ok(ast) => match script_service.run(&ast) {
                                        Ok(_) => Ok(()),
                                        Err(e) => Err(format!("Script execution error: {}", e)),
                                    },
                                    Err(e) => Err(format!("Script compilation error: {}", e)),
                                }
                            }).await;

                            match result {
                                Ok(Ok(())) => {
                                    trace!("start.bas completed successfully for session {}", actual_session_id);

                                    // Mark start.bas as executed for this session to prevent re-running
                                    if let Some(cache) = &self.state.cache {
                                        if let Ok(mut conn) = cache.get_multiplexed_async_connection().await {
                                            let _: Result<(), redis::RedisError> = redis::cmd("SET")
                                                .arg(&start_bas_key)
                                                .arg("1")
                                                .arg("EX")
                                                .arg("86400") // Expire after 24 hours
                                                .query_async(&mut conn)
                                                .await;
                                        }
                                    }
                                }
                                Ok(Err(e)) => {
                                    error!("start.bas error for session {}: {}", actual_session_id, e);
                                }
                                Err(e) => {
                                    error!("start.bas task error for session {}: {}", actual_session_id, e);
                                }
                            }
                        }
                    }
                }
            } // End of if should_execute_start_bas

            // If message content is empty, we stop here after potentially running start.bas.
            // This happens when the bot is activated by its name in WhatsApp, where an empty string is sent as a signal.
            if message_content.trim().is_empty() {
                let bot_id_str = message.bot_id.clone();
                let session_id_str = message.session_id.clone();
                
                #[cfg(feature = "chat")]
                let suggestions = get_suggestions(self.state.cache.as_ref(), &bot_id_str, &session_id_str);
                #[cfg(not(feature = "chat"))]
                let suggestions: Vec<crate::core::shared::models::Suggestion> = Vec::new();

                let final_response = BotResponse {
                    bot_id: message.bot_id,
                    user_id: message.user_id,
                    session_id: message.session_id,
                    channel: message.channel,
                    content: String::new(),
                    message_type: MessageType::BOT_RESPONSE,
                    stream_token: None,
                    is_complete: true,
                    suggestions,
                    context_name: None,
                    context_length: 0,
                    context_max_length: 0,
                };

                if let Err(e) = response_tx.send(final_response).await {
                    warn!("Failed to send final response for empty content: {}", e);
                }
                return Ok(());
            }

            if let Some(kb_manager) = self.state.kb_manager.as_ref() {
                let context = crate::core::bot::kb_context::KbInjectionContext {
                    session_id,
                    bot_id: session.bot_id,
                    bot_name: &bot_name_for_context,
                    user_query: &message_content,
                    messages: &mut messages,
                    max_context_tokens: 8000,
                };
                if let Err(e) = inject_kb_context(
                    kb_manager.clone(),
                    self.state.conn.clone(),
                    context,
                )
                .await
                {
                    error!("Failed to inject KB context: {}", e);
                }
            }
        }

        // Sanitize user message to remove any UTF-16 surrogate characters
        let sanitized_message_content = message_content
            .chars()
            .filter(|c| {
                let cp = *c as u32;
                !(0xD800..=0xDBFF).contains(&cp) && !(0xDC00..=0xDFFF).contains(&cp)
            })
            .collect::<String>();

        // Add the current user message to the messages array
        if let Some(msgs_array) = messages.as_array_mut() {
            msgs_array.push(serde_json::json!({
                "role": "user",
                "content": sanitized_message_content
            }));
        }

        let (stream_tx, mut stream_rx) = mpsc::channel::<String>(100);

        // Use bot-specific LLM provider if the bot has its own llm-url configured
        let llm: std::sync::Arc<dyn crate::llm::LLMProvider> = if let Some(ref url) = bot_llm_url {
            info!("Bot has custom llm-url: {}, creating per-bot LLM provider", url);
            // Parse explicit provider type if configured (e.g., "openai", "bedrock", "claude")
            let explicit_type = explicit_llm_provider.as_ref().map(|p| {
                let parsed: crate::llm::LLMProviderType = p.as_str().into();
                info!("Using explicit llm-provider config: {:?} for bot {}", parsed, session.bot_id);
                parsed
            });
            crate::llm::create_llm_provider_from_url(url, Some(model.clone()), None, explicit_type)
        } else {
            self.state.llm_provider.clone()
        };

        let model_clone = model.clone();
        let key_clone = key.clone();

        // Retrieve session tools for tool calling (use actual session.id after potential creation)
        let session_tools = get_session_tools(&self.state.conn, &bot_name_for_context, &session.id);
        let tools_for_llm = match session_tools {
            Ok(tools) => {
                if !tools.is_empty() {
                    Some(tools)
                } else {
                    None
                }
            }
            Err(e) => {
                warn!("Failed to load session tools: {}", e);
                None
            }
        };

        // Clone messages for the async task
        let messages_clone = messages.clone();

        tokio::spawn(async move {
            if let Err(e) = llm
                .generate_stream("", &messages_clone, stream_tx, &model_clone, &key_clone, tools_for_llm.as_ref())
                .await
            {
                error!("LLM streaming error: {}", e);
            }
        });

        let mut full_response = String::new();
        let mut analysis_buffer = String::new();
        let mut in_analysis = false;
        let mut tool_call_buffer = String::new(); // Accumulate potential tool call JSON chunks
        let mut accumulating_tool_call = false; // Track if we're currently accumulating a tool call
        let handler = llm_models::get_handler(&model);

        trace!("Using model handler for {}", model);
        trace!("Receiving LLM stream chunks...");

        #[cfg(feature = "nvidia")]
        {
            let initial_tokens = crate::core::shared::utils::estimate_token_count(&context_data);
            let config_manager = ConfigManager::new(self.state.conn.clone());
            let max_context_size = config_manager
                .get_config(&session.bot_id, "llm-server-ctx-size", None)
                .unwrap_or_default()
                .parse::<usize>()
                .unwrap_or(0);

            if let Ok(metrics) = get_system_metrics() {
                eprintln!(
                    "\nNVIDIA: {:.1}% | CPU: {:.1}% | Tokens: {}/{}",
                    metrics.gpu_usage.unwrap_or(0.0),
                    metrics.cpu_usage,
                    initial_tokens,
                    max_context_size
                );
            }
        }

        while let Some(chunk) = stream_rx.recv().await {

            // ===== GENERIC TOOL EXECUTION =====
            // Add chunk to tool_call_buffer and try to parse
            // Tool calls arrive as JSON that can span multiple chunks

            // Check if this chunk contains JSON (either starts with {/[ or contains {/[)
            let chunk_contains_json = chunk.trim().starts_with('{') || chunk.trim().starts_with('[') ||
                                        chunk.contains('{') || chunk.contains('[');

            let chunk_in_tool_buffer = if accumulating_tool_call {
                // Already accumulating - add entire chunk to buffer
                tool_call_buffer.push_str(&chunk);
                true
            } else if chunk_contains_json {
                // Check if { appears in the middle of the chunk (mixed text + JSON)
                let json_start = chunk.find('{').or_else(|| chunk.find('['));

                if let Some(pos) = json_start {
                    if pos > 0 {
                        // Send the part before { as regular content
                        let regular_part = &chunk[..pos];
                        if !regular_part.trim().is_empty() {
                            full_response.push_str(regular_part);

                            let response = BotResponse {
                                bot_id: message.bot_id.clone(),
                                user_id: message.user_id.clone(),
                                session_id: message.session_id.clone(),
                                channel: message.channel.clone(),
                                content: regular_part.to_string(),
                                message_type: MessageType::BOT_RESPONSE,
                                stream_token: None,
                                is_complete: false,
                                suggestions: Vec::new(),
                                context_name: None,
                                context_length: 0,
                                context_max_length: 0,
                            };

                            if response_tx.send(response).await.is_err() {
                                warn!("Response channel closed");
                                break;
                            }
                        }

                        // Start accumulating from { onwards
                        accumulating_tool_call = true;
                        tool_call_buffer.push_str(&chunk[pos..]);
                        true
                    } else {
                        // Chunk starts with { or [
                        accumulating_tool_call = true;
                        tool_call_buffer.push_str(&chunk);
                        true
                    }
                } else {
                    // Contains {/[ but find() failed - shouldn't happen, but send as regular content
                    false
                }
            } else {
                false
            };

            // Try to parse tool call from accumulated buffer
            let tool_call = if chunk_in_tool_buffer {
                ToolExecutor::parse_tool_call(&tool_call_buffer)
            } else {
                None
            };

            if let Some(tc) = tool_call {
                let execution_result = ToolExecutor::execute_tool_call(
                    &self.state,
                    &bot_name_for_context,
                    &tc,
                    &session_id,
                    &user_id,
                )
                .await;

                if execution_result.success {
                    info!(
                        "[TOOL_EXEC] Tool '{}' executed successfully: {}",
                        tc.tool_name, execution_result.result
                    );

                    // Send tool execution result to user
                    let response = BotResponse {
                        bot_id: message.bot_id.clone(),
                        user_id: message.user_id.clone(),
                        session_id: message.session_id.clone(),
                        channel: message.channel.clone(),
                        content: execution_result.result,
                        message_type: MessageType::BOT_RESPONSE,
                        stream_token: None,
                        is_complete: false,
                        suggestions: Vec::new(),
                        context_name: None,
                        context_length: 0,
                        context_max_length: 0,
                    };

                    if response_tx.send(response).await.is_err() {
                        warn!("Response channel closed during tool execution");
                        break;
                    }
                } else {
                    error!(
                        "[TOOL_EXEC] Tool '{}' execution failed: {:?}",
                        tc.tool_name, execution_result.error
                    );

                    // Send error to user
                    let error_msg = format!(
                        "Erro ao executar ferramenta '{}': {:?}",
                        tc.tool_name,
                        execution_result.error
                    );

                    let response = BotResponse {
                        bot_id: message.bot_id.clone(),
                        user_id: message.user_id.clone(),
                        session_id: message.session_id.clone(),
                        channel: message.channel.clone(),
                        content: error_msg,
                        message_type: MessageType::BOT_RESPONSE,
                        stream_token: None,
                        is_complete: false,
                        suggestions: Vec::new(),
                        context_name: None,
                        context_length: 0,
                        context_max_length: 0,
                    };

                    if response_tx.send(response).await.is_err() {
                        warn!("Response channel closed during tool error");
                        break;
                    }
                }

                // Don't add tool_call JSON to full_response or analysis_buffer
                // Clear the tool_call_buffer since we found and executed a tool call
                tool_call_buffer.clear();
                accumulating_tool_call = false; // Reset accumulation flag
                // Continue to next chunk
                continue;
            }

            // Clear tool_call_buffer if it's getting too large and no tool call was found
            // This prevents memory issues from accumulating JSON fragments
            // Increased limit to 50000 to handle large tool calls with many parameters
            if tool_call_buffer.len() > 50000 {
                // Flush accumulated content to client since it's too large to be a tool call
                full_response.push_str(&tool_call_buffer);

                let response = BotResponse {
                    bot_id: message.bot_id.clone(),
                    user_id: message.user_id.clone(),
                    session_id: message.session_id.clone(),
                    channel: message.channel.clone(),
                    content: tool_call_buffer.clone(),
                    message_type: MessageType::BOT_RESPONSE,
                    stream_token: None,
                    is_complete: false,
                    suggestions: Vec::new(),
                    context_name: None,
                    context_length: 0,
                    context_max_length: 0,
                };

                tool_call_buffer.clear();
                accumulating_tool_call = false; // Reset accumulation flag after flush

                if response_tx.send(response).await.is_err() {
                    warn!("Response channel closed");
                    break;
                }
            }

            // If this chunk was added to tool_call_buffer and no tool call was found yet,
            // skip processing (it's part of an incomplete tool call JSON)
            if chunk_in_tool_buffer {
                continue;
            }
            // ===== END TOOL EXECUTION =====

            analysis_buffer.push_str(&chunk);

            if !in_analysis && handler.has_analysis_markers(&analysis_buffer) {
                in_analysis = true;
                log::debug!(
                    "Detected start of thinking/analysis content for model {}",
                    model
                );

                let processed = handler.process_content(&analysis_buffer);
                if !processed.is_empty() && processed != analysis_buffer {
                    full_response.push_str(&processed);

                    let response = BotResponse {
                        bot_id: message.bot_id.clone(),
                        user_id: message.user_id.clone(),
                        session_id: message.session_id.clone(),
                        channel: message.channel.clone(),
                        content: processed,
                        message_type: MessageType::BOT_RESPONSE,
                        stream_token: None,
                        is_complete: false,
                        suggestions: Vec::new(),
                        context_name: None,
                        context_length: 0,
                        context_max_length: 0,
                    };

                    if response_tx.send(response).await.is_err() {
                        warn!("Response channel closed");
                        break;
                    }
                }
                continue;
            }

            if in_analysis && handler.is_analysis_complete(&analysis_buffer) {
                in_analysis = false;
                trace!("Detected end of thinking for model {}", model);

                let processed = handler.process_content(&analysis_buffer);
                if !processed.is_empty() {
                    full_response.push_str(&processed);

                    let response = BotResponse {
                        bot_id: message.bot_id.clone(),
                        user_id: message.user_id.clone(),
                        session_id: message.session_id.clone(),
                        channel: message.channel.clone(),
                        content: processed,
                        message_type: MessageType::BOT_RESPONSE,
                        stream_token: None,
                        is_complete: false,
                        suggestions: Vec::new(),
                        context_name: None,
                        context_length: 0,
                        context_max_length: 0,
                    };

                    if response_tx.send(response).await.is_err() {
                        warn!("Response channel closed");
                        break;
                    }
                }

                analysis_buffer.clear();
                continue;
            }

            if in_analysis {
                trace!("Accumulating thinking content, not sending to user");
                continue;
            }

            if !in_analysis {
                full_response.push_str(&chunk);

                let response = BotResponse {
                    bot_id: message.bot_id.clone(),
                    user_id: message.user_id.clone(),
                    session_id: message.session_id.clone(),
                    channel: message.channel.clone(),
                    content: chunk,
                    message_type: MessageType::BOT_RESPONSE,
                    stream_token: None,
                    is_complete: false,
                    suggestions: Vec::new(),
                    context_name: None,
                    context_length: 0,
                    context_max_length: 0,
                };

                if response_tx.send(response).await.is_err() {
                    warn!("Response channel closed");
                    break;
                }
            }
        }

        trace!("LLM stream complete. Full response: {}", full_response);

        let state_for_save = self.state.clone();
        let full_response_clone = full_response.clone();
        tokio::task::spawn_blocking(
            move || -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
                let mut sm = state_for_save.session_manager.blocking_lock();
                sm.save_message(session.id, user_id, 2, &full_response_clone, 2)?;
                Ok(())
            },
        )
        .await??;

        // Extract bot_id and session_id before moving them into BotResponse
        let bot_id_str = message.bot_id.clone();
        let session_id_str = message.session_id.clone();

        #[cfg(feature = "chat")]
        let suggestions = get_suggestions(self.state.cache.as_ref(), &bot_id_str, &session_id_str);
        #[cfg(not(feature = "chat"))]
        let suggestions: Vec<crate::core::shared::models::Suggestion> = Vec::new();

        // Content was already sent as streaming chunks.
        // Sending full_response again would duplicate it (especially for WhatsApp which accumulates buffer).
        // The final response is just a signal that streaming is complete - it should not contain content.
        let final_content = String::new();

        let final_response = BotResponse {
            bot_id: message.bot_id,
            user_id: message.user_id,
            session_id: message.session_id,
            channel: message.channel,
            content: final_content,
            message_type: MessageType::BOT_RESPONSE,
            stream_token: None,
            is_complete: true,
            suggestions,
            context_name: None,
            context_length: 0,
            context_max_length: 0,
        };

        response_tx.send(final_response).await?;
        Ok(())
    }

    #[cfg(not(feature = "llm"))]
    pub async fn stream_response(
        &self,
        message: UserMessage,
        response_tx: mpsc::Sender<BotResponse>,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        warn!("LLM feature not enabled, cannot stream response");

        let error_response = BotResponse {
            bot_id: message.bot_id,
            user_id: message.user_id,
            session_id: message.session_id,
            channel: message.channel,
            content: "LLM feature is not enabled in this build".to_string(),
            message_type: MessageType::BOT_RESPONSE,
            stream_token: None,
            is_complete: true,
            suggestions: Vec::new(),
            context_name: None,
            context_length: 0,
            context_max_length: 0,
        };

        response_tx.send(error_response).await?;
        Ok(())
    }

    pub async fn get_user_sessions(
        &self,
        user_id: Uuid,
    ) -> Result<Vec<UserSession>, Box<dyn std::error::Error + Send + Sync>> {
        let mut session_manager = self.state.session_manager.lock().await;
        let sessions = session_manager.get_user_sessions(user_id)?;
        Ok(sessions)
    }

    pub async fn get_conversation_history(
        &self,
        session_id: Uuid,
        user_id: Uuid,
    ) -> Result<Vec<(String, String)>, Box<dyn std::error::Error + Send + Sync>> {
        let mut session_manager = self.state.session_manager.lock().await;
        let history = session_manager.get_conversation_history(session_id, user_id)?;
        Ok(history)
    }
}

pub async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    Query(params): Query<HashMap<String, String>>,
) -> impl IntoResponse {
    info!("websocket_handler: Received request with params: {:?}", params);
    let session_id = params
        .get("session_id")
        .and_then(|s| Uuid::parse_str(s).ok());
    let user_id = params.get("user_id").and_then(|s| Uuid::parse_str(s).ok());

    // Extract bot_name from query params
    let bot_name = params
        .get("bot_name")
        .cloned()
        .unwrap_or_else(|| "default".to_string());

    // Allow anonymous connections for desktop UI - create UUIDs if not provided
    let session_id = session_id.unwrap_or_else(Uuid::new_v4);
    let user_id = user_id.unwrap_or_else(Uuid::new_v4);

    info!("WebSocket: session_id from params = {:?}, user_id = {:?}", session_id, user_id);

    // Look up bot_id from bot_name
    let bot_id = {
        let conn = state.conn.get().ok();
        if let Some(mut db_conn) = conn {
            use crate::core::shared::models::schema::bots::dsl::*;

            // Try to parse as UUID first, if that fails treat as bot name
            let result: Result<Uuid, _> = if let Ok(uuid) = Uuid::parse_str(&bot_name) {
                // Parameter is a UUID, look up by id
                bots.filter(id.eq(uuid)).select(id).first(&mut db_conn)
            } else {
                // Parameter is a bot name, look up by name
                bots.filter(name.eq(&bot_name))
                    .select(id)
                    .first(&mut db_conn)
            };

            result.unwrap_or_else(|_| {
                log::warn!("Bot not found: {}, using nil bot_id", bot_name);
                Uuid::nil()
            })
        } else {
            log::warn!("Could not get database connection, using nil bot_id");
            Uuid::nil()
        }
    };

    ws.on_upgrade(move |socket| handle_websocket(socket, state, session_id, user_id, bot_id))
        .into_response()
}

pub async fn websocket_handler_with_bot(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    Path(bot_name): Path<String>,
    Query(params): Query<HashMap<String, String>>,
) -> impl IntoResponse {
    let mut params = params;
    if !bot_name.is_empty() {
        params.insert("bot_name".to_string(), bot_name);
    }
    websocket_handler(ws, State(state), Query(params)).await
}

async fn handle_websocket(
    socket: WebSocket,
    state: Arc<AppState>,
    session_id: Uuid,
    user_id: Uuid,
    bot_id: Uuid,
) {
    let (mut sender, mut receiver) = socket.split();
    let (tx, mut rx) = mpsc::channel::<BotResponse>(100);

    state
        .web_adapter
        .add_connection(session_id.to_string(), tx.clone())
        .await;

    {
        let mut channels = state.response_channels.lock().await;
        channels.insert(session_id.to_string(), tx.clone());
    }

    info!(
        "WebSocket connected for session: {}, user: {}, bot: {}",
        session_id, user_id, bot_id
    );

    // Get bot_name for tools loading
    let bot_name_result = {
        let conn = state.conn.get().ok();
        if let Some(mut db_conn) = conn {
            use crate::core::shared::models::schema::bots::dsl::*;
            bots.filter(id.eq(bot_id))
                .select(name)
                .first::<String>(&mut db_conn)
                .ok()
        } else {
            None
        }
    };

    // Load session tools
    let tools = if let Some(bot_name) = bot_name_result {
        match get_session_tools(&state.conn, &bot_name, &session_id) {
            Ok(tools_vec) => {
                info!(
                    "[WEBSOCKET] Loaded {} session tools for bot {}, session {}",
                    tools_vec.len(),
                    bot_name,
                    session_id
                );
                tools_vec
            }
            Err(e) => {
                error!(
                    "[WEBSOCKET] Failed to load session tools for bot {}, session {}: {}",
                    bot_name, session_id, e
                );
                vec![]
            }
        }
    } else {
        warn!(
            "[WEBSOCKET] Could not get bot name for bot_id {}, no session tools loaded",
            bot_id
        );
        vec![]
    };

    let welcome = serde_json::json!({
        "type": "connected",
        "session_id": session_id,
        "user_id": user_id,
        "bot_id": bot_id,
        "message": "Connected to bot server",
        "tools": tools
    });

    if let Ok(welcome_str) = serde_json::to_string(&welcome) {
        if sender.send(Message::Text(welcome_str)).await.is_err() {
            error!("Failed to send welcome message");
        }
    }

    let (send_ready_tx, send_ready_rx) = tokio::sync::mpsc::channel::<()>(1);

    // Execute start.bas automatically on connection (similar to auth.ast pattern)
    {
        let bot_name_result = {
            let conn = state.conn.get().ok();
            if let Some(mut db_conn) = conn {
                use crate::core::shared::models::schema::bots::dsl::*;
                bots.filter(id.eq(bot_id))
                    .select(name)
                    .first::<String>(&mut db_conn)
                    .ok()
            } else {
                None
            }
        };

        // DEBUG: Log start script execution attempt
        info!(
            "Checking for start.bas: bot_id={}, bot_name_result={:?}",
            bot_id,
            bot_name_result
        );

        if let Some(bot_name) = bot_name_result {
            // Web clients expect start.bas to execute their first screen every time they connect/reload.
            // We always run it, but we SET the start_bas_key flag right after so stream_response skips execution.
            let start_bas_key = format!("start_bas_executed:{}", session_id);
            let should_execute_start_bas = true;

            if should_execute_start_bas {
                let data_dir = "/opt/gbo/data";
                let start_script_path = format!("{}/{}.gbai/{}.gbdialog/start.bas", data_dir, bot_name, bot_name);

                info!("Looking for start.bas at: {}", start_script_path);

                if let Ok(metadata) = tokio::fs::metadata(&start_script_path).await {
                    if metadata.is_file() {
                        info!("Found start.bas file, reading contents...");
                        if let Ok(start_script) = tokio::fs::read_to_string(&start_script_path).await {
                        info!(
                            "Executing start.bas for bot {} on session {}",
                            bot_name, session_id
                        );

                        let state_for_start = state.clone();
                        let tx_for_start = tx.clone();
                        let bot_id_str = bot_id.to_string();
                        let session_id_str = session_id.to_string();
                        let mut send_ready_rx = send_ready_rx;

                        tokio::spawn(async move {
                            let _ = send_ready_rx.recv().await;
                            
                            let session_result = {
                                let mut sm = state_for_start.session_manager.lock().await;
                                let by_id = sm.get_session_by_id(session_id);
                                match by_id {
                                    Ok(Some(s)) => Ok(Some(s)),
                                    _ => sm.get_or_create_user_session(user_id, bot_id, "Chat Session"),
                                }
                            };

                            if let Ok(Some(mut session)) = session_result {
                                info!("start.bas: Found session {} for websocket session {}", session.id, session_id);

                                // Store WebSocket session_id in context so TALK can route messages correctly
                                if let serde_json::Value::Object(ref mut map) = session.context_data {
                                    map.insert("websocket_session_id".to_string(), serde_json::Value::String(session_id.to_string()));
                                } else {
                                    let mut map = serde_json::Map::new();
                                    map.insert("websocket_session_id".to_string(), serde_json::Value::String(session_id.to_string()));
                                    session.context_data = serde_json::Value::Object(map);
                                }

                                // Clone state_for_start for use in Redis SET after execution
                                let state_for_redis = state_for_start.clone();

                                let result = tokio::task::spawn_blocking(move || {
                                    info!("start.bas: Creating ScriptService with session.id={}", session.id);
                                    let mut script_service = crate::basic::ScriptService::new(
                                        state_for_start.clone(),
                                        session.clone()
                                    );
                                    script_service.load_bot_config_params(&state_for_start, bot_id);

                                    match script_service.compile(&start_script) {
                                        Ok(ast) => match script_service.run(&ast) {
                                            Ok(_) => Ok(()),
                                            Err(e) => Err(format!("Script execution error: {}", e)),
                                        },
                                        Err(e) => Err(format!("Script compilation error: {}", e)),
                                    }
                                }).await;

                                match result {
                                    Ok(Ok(())) => {
                                        info!("start.bas executed successfully for bot {}", bot_name);

                                        // Mark start.bas as executed for this session to prevent re-running
                                        if let Some(cache) = &state_for_redis.cache {
                                            if let Ok(mut conn) = cache.get_multiplexed_async_connection().await {
                                                let _: Result<(), redis::RedisError> = redis::cmd("SET")
                                                    .arg(&start_bas_key)
                                                    .arg("1")
                                                    .arg("EX")
                                                    .arg("86400") // Expire after 24 hours
                                                    .query_async(&mut conn)
                                                    .await;
                                                info!("Marked start.bas as executed for session {}", session_id);
                                            }
                                        }

                                        // Fetch suggestions from Redis and send to frontend
                                        let user_id_str = user_id.to_string();
                                        let suggestions = get_suggestions(state_for_redis.cache.as_ref(), &bot_id_str, &session_id_str);
                                        if !suggestions.is_empty() {
                                            info!("Sending {} suggestions to frontend for session {}", suggestions.len(), session_id);
                                            let response = BotResponse {
                                                bot_id: bot_id_str.clone(),
                                                user_id: user_id_str.clone(),
                                                session_id: session_id_str.clone(),
                                                channel: "Chat".to_string(),
                                                content: String::new(),
                                                message_type: MessageType::BOT_RESPONSE,
                                                stream_token: None,
                                                is_complete: true,
                                                suggestions,
                                                context_name: None,
                                                context_length: 0,
                                                context_max_length: 0,
                                            };
                                            let _ = tx_for_start.send(response).await;


                                        }
                                    }
                                    Ok(Err(e)) => {
                                        error!("start.bas error for bot {}: {}", bot_name, e);
                                    }
                                    Err(e) => {
                                        error!("start.bas task error for bot {}: {}", bot_name, e);
                                    }
                                }
                            }
                        });
                    }
                }
            }
            } // End of if should_execute_start_bas
        }
    }

    let mut send_task = tokio::spawn(async move {
        while let Some(response) = rx.recv().await {
            if let Ok(json_str) = serde_json::to_string(&response) {
                if sender.send(Message::Text(json_str)).await.is_err() {
                    break;
                }
            }
        }
    });

    let _ = send_ready_tx.send(()).await;

    let state_clone = state.clone();
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            match msg {
                Message::Text(text) => {
                    if let Ok(user_msg) = serde_json::from_str::<UserMessage>(&text) {
                        let orchestrator = BotOrchestrator::new(state_clone.clone());
                        if let Some(tx_clone) = state_clone
                            .response_channels
                            .lock()
                            .await
                            .get(&session_id.to_string())
                        {
                            // Ensure session exists - create if not
                            let session_result = {
                                let mut sm = state_clone.session_manager.lock().await;
                                sm.get_session_by_id(session_id)
                            };

                            let session = match session_result {
                                Ok(Some(sess)) => sess,
                                Ok(None) => {
                                    // Use session manager to create session (will generate new UUID)
                                    let mut sm = state_clone.session_manager.lock().await;
                                    match sm.create_session(user_id, bot_id, "WebSocket Chat") {
                                        Ok(new_session) => new_session,
                                        Err(e) => {
                                            error!("Failed to create session: {}", e);
                                            continue;
                                        }
                                    }
                                }
                                Err(e) => {
                                    error!("Error getting session: {}", e);
                                    continue;
                                }
                            };

                            // Use bot_id from WebSocket connection instead of from message
                            let corrected_msg = UserMessage {
                                bot_id: bot_id.to_string(),
                                user_id: session.user_id.to_string(),
                                session_id: session.id.to_string(),
                                ..user_msg
                            };
                            if let Err(e) = orchestrator
                                .stream_response(corrected_msg, tx_clone.clone())
                                .await
                            {
                                error!("Failed to stream response: {}", e);
                            }
                        } else {
                            warn!("Response channel NOT found for session: {}", session_id);
                        }
                    } else {
                        warn!("Failed to parse UserMessage from: {}", text);
                    }
                }
                Message::Close(_) => {
                    trace!("WebSocket close message received");
                    break;
                }
                _ => {}
            }
        }
    });

    tokio::select! {
        _ = (&mut send_task) => { recv_task.abort(); }
        _ = (&mut recv_task) => { send_task.abort(); }
    }

    state
        .web_adapter
        .remove_connection(&session_id.to_string())
        .await;

    {
        let mut channels = state.response_channels.lock().await;
        channels.remove(&session_id.to_string());
    }

    info!("WebSocket disconnected for session: {}", session_id);
}

pub fn create_bot_handler(
    Extension(state): Extension<Arc<AppState>>,
    Json(payload): Json<HashMap<String, String>>,
) -> impl IntoResponse {
    let bot_name = payload
        .get("bot_name")
        .cloned()
        .unwrap_or_else(|| "default".to_string());

    let orchestrator = BotOrchestrator::new(state);
    if let Err(e) = orchestrator.mount_all_bots() {
        error!("Failed to mount bots: {}", e);
    }

    (
        StatusCode::OK,
        Json(serde_json::json!({ "status": format!("bot '{}' created", bot_name) })),
    )
}

pub fn mount_bot_handler(
    Extension(state): Extension<Arc<AppState>>,
    Json(payload): Json<HashMap<String, String>>,
) -> impl IntoResponse {
    let bot_guid = payload.get("bot_guid").cloned().unwrap_or_default();

    let orchestrator = BotOrchestrator::new(state);
    if let Err(e) = orchestrator.mount_all_bots() {
        error!("Failed to mount bot: {}", e);
    }

    (
        StatusCode::OK,
        Json(serde_json::json!({ "status": format!("bot '{}' mounted", bot_guid) })),
    )
}

pub async fn handle_user_input_handler(
    Extension(state): Extension<Arc<AppState>>,
    Json(payload): Json<HashMap<String, String>>,
) -> impl IntoResponse {
    let session_id = payload.get("session_id").cloned().unwrap_or_default();
    let user_input = payload.get("input").cloned().unwrap_or_default();

    info!(
        "Processing user input: {} for session: {}",
        user_input, session_id
    );

    let orchestrator = BotOrchestrator::new(state);
    if let Ok(sessions) = orchestrator.get_user_sessions(Uuid::nil()).await {
        info!("Found {} sessions", sessions.len());
    }

    (
        StatusCode::OK,
        Json(serde_json::json!({ "status": format!("processed: {}", user_input) })),
    )
}

pub async fn get_user_sessions_handler(
    Extension(state): Extension<Arc<AppState>>,
    Json(payload): Json<HashMap<String, String>>,
) -> impl IntoResponse {
    let user_id = payload
        .get("user_id")
        .and_then(|id| Uuid::parse_str(id).ok())
        .unwrap_or_else(Uuid::nil);

    let orchestrator = BotOrchestrator::new(state);
    match orchestrator.get_user_sessions(user_id).await {
        Ok(sessions) => (
            StatusCode::OK,
            Json(serde_json::json!({ "sessions": sessions })),
        ),
        Err(e) => {
            error!("Failed to get sessions: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": e.to_string() })),
            )
        }
    }
}

pub async fn get_conversation_history_handler(
    Extension(state): Extension<Arc<AppState>>,
    Json(payload): Json<HashMap<String, String>>,
) -> impl IntoResponse {
    let session_id = payload
        .get("session_id")
        .and_then(|id| Uuid::parse_str(id).ok())
        .unwrap_or_else(Uuid::nil);
    let user_id = payload
        .get("user_id")
        .and_then(|id| Uuid::parse_str(id).ok())
        .unwrap_or_else(Uuid::nil);

    let orchestrator = BotOrchestrator::new(state);
    match orchestrator
        .get_conversation_history(session_id, user_id)
        .await
    {
        Ok(history) => (
            StatusCode::OK,
            Json(serde_json::json!({ "history": history })),
        ),
        Err(e) => {
            error!("Failed to get history: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": e.to_string() })),
            )
        }
    }
}

pub async fn send_warning_handler(
    Extension(state): Extension<Arc<AppState>>,
    Json(payload): Json<HashMap<String, String>>,
) -> impl IntoResponse {
    let message = payload
        .get("message")
        .cloned()
        .unwrap_or_else(|| "Warning".to_string());
    let session_id = payload.get("session_id").cloned().unwrap_or_default();

    warn!("Warning for session {}: {}", session_id, message);

    let orchestrator = BotOrchestrator::new(state);
    info!("Orchestrator created for warning");

    if let Ok(sessions) = orchestrator.get_user_sessions(Uuid::nil()).await {
        info!("Current active sessions: {}", sessions.len());
    }

    (
        StatusCode::OK,
        Json(serde_json::json!({ "status": "warning sent", "message": message })),
    )
}
