opencrabs 0.3.24

//! Telegram Agent
//!
//! Agent struct and startup logic.

use super::TelegramState;
use super::handler::handle_message;
use crate::brain::agent::AgentService;
use crate::config::Config;
use crate::db::ChannelMessageRepository;
use crate::services::{ServiceContext, SessionService};
use std::collections::HashMap;
use std::sync::Arc;
use teloxide::prelude::*;
use tokio::sync::Mutex;
use uuid::Uuid;

/// Telegram bot that forwards messages to the agent
pub struct TelegramAgent {
    agent_service: Arc<AgentService>,
    session_service: SessionService,
    /// Shared session ID from the TUI — owner user shares the terminal session
    shared_session_id: Arc<Mutex<Option<Uuid>>>,
    telegram_state: Arc<TelegramState>,
    config_rx: tokio::sync::watch::Receiver<Config>,
    channel_msg_repo: ChannelMessageRepository,
}

impl TelegramAgent {
    pub fn new(
        agent_service: Arc<AgentService>,
        service_context: ServiceContext,
        shared_session_id: Arc<Mutex<Option<Uuid>>>,
        telegram_state: Arc<TelegramState>,
        config_rx: tokio::sync::watch::Receiver<Config>,
        channel_msg_repo: ChannelMessageRepository,
    ) -> Self {
        Self {
            agent_service,
            session_service: SessionService::new(service_context),
            shared_session_id,
            telegram_state,
            config_rx,
            channel_msg_repo,
        }
    }

    /// Start the bot as a background task. Returns a JoinHandle.
    pub fn start(self, token: String) -> tokio::task::JoinHandle<()> {
        tokio::spawn(async move {
            // Validate token format BEFORE creating Bot: "numbers:alphanumeric"
            // e.g., "123456789:ABCdefGHIjklMNOpqrsTUVwxyz"
            if token.is_empty() {
                tracing::debug!("Telegram bot token is empty, skipping bot start");
                return;
            }

            if !token.contains(':') {
                tracing::debug!("Telegram bot token missing ':' separator, skipping bot start");
                return;
            }

            let parts: Vec<&str> = token.splitn(2, ':').collect();
            if parts.len() != 2 {
                tracing::debug!("Telegram bot token has invalid format, skipping bot start");
                return;
            }

            // First part must be numeric (bot ID)
            if parts[0].parse::<u64>().is_err() {
                tracing::debug!("Telegram bot token has invalid bot ID, skipping bot start");
                return;
            }

            // Second part must be at least 30 chars (API key)
            if parts[1].len() < 30 {
                tracing::debug!("Telegram bot token has too short API key, skipping bot start");
                return;
            }

            // Read initial config for logging
            let cfg = self.config_rx.borrow().clone();
            tracing::info!(
                "Starting Telegram bot with {} allowed user(s), STT={}, TTS={}",
                cfg.channels.telegram.allowed_users.len(),
                cfg.voice_config().stt_enabled,
                cfg.voice_config().tts_enabled,
            );

            let bot = Bot::new(token.clone());

            // Verify token works with Telegram API before setting up dispatcher
            match bot.get_me().await {
                Ok(me) => {
                    if let Some(ref username) = me.username {
                        tracing::info!("Telegram: bot username is @{}", username);
                        self.telegram_state.set_bot_username(username.clone()).await;
                    }
                    // Store bot in state for proactive messaging only after successful auth
                    self.telegram_state.set_bot(bot.clone()).await;

                    // Register slash commands so they appear in Telegram's / menu
                    register_bot_commands(&bot).await;
                }
                Err(e) => {
                    tracing::warn!("Telegram: token validation failed: {}. Bot not started.", e);
                    return;
                }
            }

            // Per-user session tracking for non-owner users (owner shares TUI session)
            let extra_sessions: Arc<Mutex<HashMap<i64, (Uuid, std::time::Instant)>>> =
                Arc::new(Mutex::new(HashMap::new()));
            let agent = self.agent_service.clone();
            let session_svc = self.session_service.clone();
            let bot_token = Arc::new(token);
            let shared_session = self.shared_session_id.clone();
            let telegram_state = self.telegram_state.clone();
            let config_rx = self.config_rx.clone();
            let channel_msg_repo = self.channel_msg_repo.clone();

            // ── Message handler ───────────────────────────────────────────────
            let msg_handler = Update::filter_message().endpoint({
                let agent = agent.clone();
                let session_svc = session_svc.clone();
                let bot_token = bot_token.clone();
                let shared_session = shared_session.clone();
                let telegram_state = telegram_state.clone();
                let config_rx = config_rx.clone();
                let channel_msg_repo = channel_msg_repo.clone();
                move |bot: Bot, msg: Message| {
                    let agent = agent.clone();
                    let session_svc = session_svc.clone();
                    let bot_token = bot_token.clone();
                    let shared_session = shared_session.clone();
                    let telegram_state = telegram_state.clone();
                    let config_rx = config_rx.clone();
                    let channel_msg_repo = channel_msg_repo.clone();
                    async move {
                        // Spawn in background so the dispatcher is free to
                        // process callback queries (approval button clicks)
                        // while the agent is running.
                        tokio::spawn(async move {
                            let result = tokio::task::spawn(async move {
                                handle_message(
                                    bot,
                                    msg,
                                    agent,
                                    session_svc,
                                    bot_token,
                                    shared_session,
                                    telegram_state,
                                    config_rx,
                                    channel_msg_repo,
                                )
                                .await
                            })
                            .await;
                            match result {
                                Ok(Ok(())) => {}
                                Ok(Err(e)) => {
                                    tracing::error!("Telegram handle_message error: {e}");
                                }
                                Err(panic_err) => {
                                    tracing::error!(
                                        "Telegram handle_message panicked: {:?}",
                                        panic_err
                                    );
                                }
                            }
                        });
                        ResponseResult::Ok(())
                    }
                }
            });

            // ── Callback query handler (for Approve / Deny inline buttons) ────
            let cb_handler = Update::filter_callback_query().endpoint({
                let telegram_state = telegram_state.clone();
                let agent = agent.clone();
                let session_svc = session_svc.clone();
                let shared_session = shared_session.clone();
                let extra_sessions = extra_sessions.clone();
                let config_rx = config_rx.clone();
                move |bot: Bot, query: CallbackQuery| {
                    let state = telegram_state.clone();
                    let agent = agent.clone();
                    let session_svc = session_svc.clone();
                    let shared_session = shared_session.clone();
                    let extra_sessions = extra_sessions.clone();
                    let config_rx = config_rx.clone();
                    async move {
                        if let Some(data) = query.data.as_deref() {
                            tracing::info!("Telegram callback query received: data={}", data);

                            // Provider picker callback → show models for that provider
                            if let Some(provider_name) = data.strip_prefix("provider:") {
                                let resp = crate::channels::commands::models_for_provider(provider_name).await;

                                // Agent-handled providers (OpenRouter 300+ models, custom)
                                // Switch to default if set, then let the agent follow up.
                                if resp.agent_handled {
                                    // Resolve session from the chat where the button was pressed,
                                    // not from shared_session (which is the TUI session).
                                    let session_id = resolve_callback_session(&query, &state, &shared_session).await;
                                    let display = crate::channels::commands::provider_display_name(provider_name);
                                    // Switch to this provider with its default model. Pin the
                                    // provider to THIS session so another channel/session
                                    // doesn't get yanked onto it — the model callback's
                                    // switch_model then reads from the same per-session slot.
                                    if let Ok(config) = crate::config::Config::load()
                                        && let Ok(new_provider) = crate::brain::provider::factory::create_provider_by_name(&config, provider_name).await
                                    {
                                        match session_id {
                                            Some(sid) => agent.swap_provider_for_session(sid, new_provider),
                                            None => agent.swap_provider(new_provider),
                                        }
                                    }
                                    if !resp.current_model.is_empty() {
                                        let _ = crate::channels::commands::switch_model(&agent, &resp.current_model, session_id, Some(provider_name)).await;
                                    }
                                    let _ = bot.answer_callback_query(&query.id).await;
                                    // Send synthetic message to agent so it handles follow-up
                                    let prompt = if resp.current_model.is_empty() {
                                        format!(
                                            "[System: User selected {} provider but no default model is set. \
                                             Ask them which model they want. Use config_manager tool to read \
                                             providers section, then set the default_model. Keep current provider \
                                             until a model is chosen.]",
                                            display
                                        )
                                    } else {
                                        format!(
                                            "[System: User switched to {} provider with model {}. \
                                             Confirm the switch. Ask if they want a different model — \
                                             if so, use config_manager to update providers.{}.default_model \
                                             and confirm.]",
                                            display, resp.current_model,
                                            if provider_name == "openrouter" { "openrouter" } else { provider_name }
                                        )
                                    };
                                    if let Some(sid) = session_id {
                                        let agent_clone = agent.clone();
                                        let bot_clone = bot.clone();
                                        let chat_id = query.message.as_ref().map(|m| m.chat().id).unwrap_or(teloxide::types::ChatId(0));
                                        tokio::spawn(async move {
                                            match agent_clone.send_message(sid, prompt, None).await {
                                                Ok(resp) => {
                                                    let clean = crate::utils::sanitize::strip_llm_artifacts(&resp.content);
                                                    let html = crate::channels::telegram::handler::md_to_html(&clean);
                                                    let _ = bot_clone.send_message(chat_id, html)
                                                        .parse_mode(teloxide::types::ParseMode::Html)
                                                        .await;
                                                }
                                                Err(e) => {
                                                    tracing::error!("Agent follow-up failed: {}", e);
                                                }
                                            }
                                        });
                                    }
                                    return ResponseResult::Ok(());
                                }

                                if resp.models.is_empty() {
                                    let _ = bot
                                        .answer_callback_query(&query.id)
                                        .text("No models available for this provider")
                                        .await;
                                    return ResponseResult::Ok(());
                                }
                                let _ = bot.answer_callback_query(&query.id).await;
                                if let Some(msg) = &query.message {
                                    use teloxide::payloads::EditMessageTextSetters;
                                    use teloxide::prelude::Requester;
                                    use teloxide::types::{InlineKeyboardButton, InlineKeyboardMarkup};
                                    let rows: Vec<Vec<InlineKeyboardButton>> = resp
                                        .models
                                        .iter()
                                        .map(|m| {
                                            let display = if *m == resp.current_model {
                                                format!("✓ {}", m)
                                            } else {
                                                m.clone()
                                            };
                                            vec![InlineKeyboardButton::callback(
                                                display,
                                                // Pipe separator because BOTH provider_name and
                                                // model can contain `:` — custom providers
                                                // are `custom:<name>` (e.g. `custom:dialagram`)
                                                // and OpenRouter models carry `:free`/`:thinking`
                                                // suffixes. Splitting on `:` here put the
                                                // provider's tail into the model name and the
                                                // session got persisted with broken metadata
                                                // (`provider=custom`, `model=dialagram:qwen-3.7-…`
                                                // — seen 2026-05-18T23:39 sync_provider trace).
                                                // Generator + parser MUST stay in lock-step.
                                                format!("model:{}|{}", resp.provider_name, m),
                                            )]
                                        })
                                        .collect();
                                    let keyboard = InlineKeyboardMarkup::new(rows);
                                    let text = crate::channels::telegram::handler::md_to_html(&resp.text);
                                    let _ = bot
                                        .edit_message_text(msg.chat().id, msg.id(), &text)
                                        .parse_mode(teloxide::types::ParseMode::Html)
                                        .reply_markup(keyboard)
                                        .await;
                                }
                                return ResponseResult::Ok(());
                            }

                            // Model switch callback (format: model:<provider>|<model>).
                            // Pipe — not colon — between provider and model so a
                            // custom-provider name (`custom:dialagram`) and an
                            // OpenRouter-style model suffix (`:free`, `:thinking`)
                            // don't fold into each other on parse.
                            if let Some(rest) = data.strip_prefix("model:") {
                                let (provider_name, model_name) = if let Some((p, m)) = rest.split_once('|') {
                                    (Some(p), m)
                                } else {
                                    (None, rest)
                                };
                                // Resolve session BEFORE the provider swap so the swap
                                // lands on the right per-session slot. Leaving the
                                // resolve below the swap (as it was) made the provider
                                // change visible to other sessions via the global slot
                                // until the per-session pin from switch_model landed.
                                let session_id = resolve_callback_session(&query, &state, &shared_session).await;
                                // Switch provider if specified and different
                                let mut provider_err: Option<String> = None;
                                if let Some(pname) = provider_name {
                                    match crate::config::Config::load() {
                                        Ok(config) => match crate::brain::provider::factory::create_provider_by_name(&config, pname).await {
                                            Ok(new_provider) => match session_id {
                                                Some(sid) => agent.swap_provider_for_session(sid, new_provider),
                                                None => agent.swap_provider(new_provider),
                                            },
                                            Err(e) => provider_err = Some(format!("Failed to create provider '{}': {}", pname, e)),
                                        },
                                        Err(e) => provider_err = Some(format!("Failed to load config: {}", e)),
                                    }
                                }
                                let (switch_ok, display_text) = if let Some(err) = provider_err {
                                    (false, format!("⚠️ {}", err))
                                } else {
                                    match crate::channels::commands::switch_model(&agent, model_name, session_id, provider_name).await {
                                        Ok(_) => (true, format!("✅ Model switched to <code>{}</code>", model_name)),
                                        Err(e) => (false, format!("⚠️ {}", e)),
                                    }
                                };
                                let _ = bot.answer_callback_query(&query.id).await;
                                if let Some(msg) = &query.message {
                                    use teloxide::payloads::EditMessageTextSetters;
                                    use teloxide::prelude::Requester;
                                    let _ = bot
                                        .edit_message_text(msg.chat().id, msg.id(), &display_text)
                                        .parse_mode(teloxide::types::ParseMode::Html)
                                        .reply_markup(
                                            teloxide::types::InlineKeyboardMarkup::default(),
                                        )
                                        .await;
                                }
                                if !switch_ok {
                                    tracing::warn!("Telegram model switch failed: {}", display_text);
                                }
                                return ResponseResult::Ok(());
                            }

                            // Session switch callback
                            if let Some(session_id_str) = data.strip_prefix("session:") {
                                if let Ok(new_id) = session_id_str.parse::<Uuid>() {
                                    // Determine if caller is owner
                                    let cfg = config_rx.borrow().clone();
                                    let caller_id = query.from.id.0 as i64;
                                    let owner_id = cfg
                                        .channels
                                        .telegram
                                        .allowed_users
                                        .first()
                                        .and_then(|s| s.parse::<i64>().ok());
                                    let is_owner = cfg.channels.telegram.allowed_users.is_empty()
                                        || owner_id == Some(caller_id);

                                    if is_owner {
                                        *shared_session.lock().await = Some(new_id);
                                    } else {
                                        extra_sessions.lock().await.insert(
                                            caller_id,
                                            (new_id, std::time::Instant::now()),
                                        );
                                    }
                                    state
                                        .register_session_chat(new_id, query.message.as_ref().map(|m| m.chat().id.0).unwrap_or(caller_id))
                                        .await;

                                    // Touch updated_at so find_session_by_title_suffix returns this session on next message
                                    if let Ok(Some(s)) = session_svc.get_session(new_id).await {
                                        let _ = session_svc.update_session(&s).await;
                                    }

                                    let _ = bot
                                        .answer_callback_query(&query.id)
                                        .text("Session switched")
                                        .await;
                                    if let Some(msg) = &query.message {
                                        use teloxide::payloads::EditMessageTextSetters;
                                        use teloxide::prelude::Requester;
                                        let _ = bot
                                            .edit_message_text(
                                                msg.chat().id,
                                                msg.id(),
                                                {
                                                    let display = match session_svc.get_session(new_id).await {
                                                        Ok(Some(s)) => s.title.unwrap_or_else(|| session_id_str[..8.min(session_id_str.len())].to_string()),
                                                        _ => session_id_str[..8.min(session_id_str.len())].to_string(),
                                                    };
                                                    format!("✅ Switched to session <code>{}</code>", display)
                                                },
                                            )
                                            .parse_mode(teloxide::types::ParseMode::Html)
                                            .reply_markup(
                                                teloxide::types::InlineKeyboardMarkup::default(),
                                            )
                                            .await;
                                    }
                                } else {
                                    let _ = bot
                                        .answer_callback_query(&query.id)
                                        .text("Invalid session ID")
                                        .await;
                                }
                                return ResponseResult::Ok(());
                            }

                            // Follow-up question callback: `q:<id>:<idx>`.
                            // Handled separately from the approve/deny chain
                            // because it returns an option string, not a
                            // boolean.
                            if let Some(rest) = data.strip_prefix("q:") {
                                let mut parts = rest.splitn(2, ':');
                                let q_id = parts.next().unwrap_or("");
                                let idx_str = parts.next().unwrap_or("");
                                let idx: usize = idx_str.parse().unwrap_or(usize::MAX);
                                let resolved = state
                                    .resolve_pending_question(q_id, idx)
                                    .await;
                                tracing::info!(
                                    "Telegram follow_up_question resolved: id={} idx={} answer={:?}",
                                    q_id,
                                    idx,
                                    resolved
                                );
                                let _ = bot.answer_callback_query(&query.id).await;
                                if let Some(answer) = resolved
                                    && let Some(msg) = &query.message
                                {
                                    let original_text = match msg {
                                        teloxide::types::MaybeInaccessibleMessage::Regular(m) => {
                                            m.text().unwrap_or("").to_string()
                                        }
                                        _ => String::new(),
                                    };
                                    let updated =
                                        format!("{}\n\n✅ {}", original_text, answer);
                                    use teloxide::payloads::EditMessageTextSetters;
                                    use teloxide::prelude::Requester;
                                    if let Err(e) = bot
                                        .edit_message_text(msg.chat().id, msg.id(), &updated)
                                        .reply_markup(
                                            teloxide::types::InlineKeyboardMarkup::default(),
                                        )
                                        .await
                                    {
                                        tracing::error!(
                                            "Telegram: failed to edit question message: {}",
                                            e
                                        );
                                    }
                                }
                                return ResponseResult::Ok(());
                            }

                            let (approved, always, yolo, id) =
                                if let Some(id) = data.strip_prefix("approve:") {
                                    (true, false, false, id.to_string())
                                } else if let Some(id) = data.strip_prefix("always:") {
                                    (true, true, false, id.to_string())
                                } else if let Some(id) = data.strip_prefix("yolo:") {
                                    (true, true, true, id.to_string())
                                } else if let Some(id) = data.strip_prefix("deny:") {
                                    (false, false, false, id.to_string())
                                } else {
                                    tracing::warn!("Telegram: unknown callback data: {}", data);
                                    let _ = bot.answer_callback_query(&query.id).await;
                                    return ResponseResult::Ok(());
                                };

                            // Persist YOLO (permanent) directly from callback
                            if yolo {
                                crate::utils::persist_auto_always_policy();
                            }

                            let resolved = state.resolve_pending_approval(&id, approved, always).await;
                            tracing::info!(
                                "Telegram approval resolved: id={}, approved={}, always={}, found_pending={}",
                                id, approved, always, resolved
                            );
                            if !resolved {
                                tracing::warn!(
                                    "Telegram: no pending approval found for id={} — may have timed out or already resolved",
                                    id
                                );
                            }
                            let _ = bot.answer_callback_query(&query.id).await;

                            // Edit the approval message: keep original context, append outcome, remove buttons
                            if let Some(msg) = &query.message {
                                let label = if yolo {
                                    "\n\n🔥 YOLO — always approved"
                                } else if always {
                                    "\n\n🔁 Always approved (session)"
                                } else if approved {
                                    "\n\n✅ Approved"
                                } else {
                                    "\n\n❌ Denied"
                                };
                                let original_text = match msg {
                                    teloxide::types::MaybeInaccessibleMessage::Regular(m) => {
                                        m.text().unwrap_or("").to_string()
                                    }
                                    _ => String::new(),
                                };
                                let updated = format!("{}{}", original_text, label);
                                use teloxide::payloads::EditMessageTextSetters;
                                use teloxide::prelude::Requester;
                                if let Err(e) = bot
                                    .edit_message_text(msg.chat().id, msg.id(), &updated)
                                    .reply_markup(teloxide::types::InlineKeyboardMarkup::default())
                                    .await
                                {
                                    tracing::error!("Telegram: failed to edit approval message: {}", e);
                                }
                            } else {
                                tracing::warn!("Telegram: callback query has no message — cannot edit");
                            }
                        } else {
                            tracing::warn!("Telegram: callback query with no data");
                            let _ = bot.answer_callback_query(&query.id).await;
                        }
                        ResponseResult::Ok(())
                    }
                }
            });

            let tree = dptree::entry().branch(msg_handler).branch(cb_handler);

            // Retry loop: if the dispatcher exits (network hiccup, Telegram conflict
            // from another process using the same token, etc.), wait and reconnect.
            // Without this, daemon mode silently loses the Telegram connection forever.
            loop {
                tracing::info!("Telegram: starting dispatcher polling loop");
                Dispatcher::builder(bot.clone(), tree.clone())
                    .build()
                    .dispatch()
                    .await;
                tracing::warn!("Telegram: dispatcher exited unexpectedly — reconnecting in 5s");
                tokio::time::sleep(std::time::Duration::from_secs(5)).await;
            }
        })
    }
}

/// Resolve the correct session ID for a callback query.
///
/// Callbacks from inline buttons (e.g. `/models` picker) fire in the chat
/// where the button was pressed. We look up the session that's registered for
/// that chat — this is the session the message handler resolved. Only falls
/// back to the shared TUI session when no chat→session mapping exists (e.g.
/// first-ever interaction before any message was processed).
async fn resolve_callback_session(
    query: &CallbackQuery,
    state: &super::TelegramState,
    shared_session: &tokio::sync::Mutex<Option<Uuid>>,
) -> Option<Uuid> {
    // Try to get the session registered for this chat
    if let Some(msg) = &query.message {
        let chat_id = msg.chat().id.0;
        if let Some(session_id) = state.chat_session(chat_id).await {
            return Some(session_id);
        }
    }
    // Fallback: shared TUI session (owner DMs before any message handler ran)
    *shared_session.lock().await
}

/// Register bot commands with Telegram so they appear in the `/` menu.
async fn register_bot_commands(bot: &Bot) {
    use teloxide::types::BotCommand;

    let commands = vec![
        BotCommand::new("help", "Show available commands"),
        BotCommand::new("models", "Switch AI model or provider"),
        BotCommand::new("usage", "Session token and cost stats"),
        BotCommand::new("new", "Start a new session"),
        BotCommand::new("sessions", "List and switch sessions"),
        BotCommand::new("stop", "Cancel the current operation"),
        BotCommand::new("compact", "Compact conversation context"),
        BotCommand::new("doctor", "Run connection health check"),
        BotCommand::new("evolve", "Check for updates"),
    ];

    match bot.set_my_commands(commands).await {
        Ok(_) => tracing::info!("Telegram: registered {} bot commands", 9),
        Err(e) => tracing::warn!("Telegram: failed to register bot commands: {}", e),
    }
}