tail-fin-daemon 0.7.8

use std::sync::Arc;

use async_trait::async_trait;
use night_fury_core::BrowserSession;
use night_fury_daemon_core::protocol::Response;
use serde_json::{json, Value};
use tail_fin_grok::GrokClient;
use tokio::sync::Mutex;

use crate::handlers::params::{
    optional_bool, optional_positive_usize, optional_string_array, required_nonempty_str,
    required_str,
};
use crate::handlers::response::{err_str, ok_json, ok_value};
use crate::handlers::SiteHandler;

/// Per-session Grok handler. Lazily initialises a [`GrokClient`] on the first
/// `grok.ask` — which triggers one page reload to harvest `x-statsig-id`.
/// Subsequent asks reuse the same client and cached statsig, so no reload is
/// needed after the first request.
///
/// The `GrokClient` is wrapped in `Arc` so the mutex guard can be dropped
/// before any async network calls, avoiding lock contention on concurrent
/// or cancelled requests.
pub struct GrokHandler {
    session: BrowserSession,
    client: Mutex<Option<Arc<GrokClient>>>,
}

impl GrokHandler {
    pub fn new(session: BrowserSession) -> Self {
        Self {
            session,
            client: Mutex::new(None),
        }
    }

    /// Best-effort: ensure the attached tab points at the target conversation.
    ///
    /// Daemon sessions are launched via a generic browser connector (no Grok-
    /// specific tab filter), so the first attached page may not be `/c/{cid}`.
    /// Statsig harvesting relies on reloading a real Grok conversation tab;
    /// otherwise no `/rest/` traffic appears and harvest times out.
    async fn ensure_on_conversation(
        &self,
        cid: &str,
        rid_hint: Option<&str>,
    ) -> Result<(), String> {
        let current_url = self.session.get_url().await.map_err(|e| e.to_string())?;
        let expected = format!("/c/{cid}");
        if current_url.contains("grok.com") && current_url.contains(&expected) {
            return Ok(());
        }

        let target = match rid_hint.filter(|r| !r.is_empty()) {
            Some(rid) => format!("https://grok.com/c/{cid}?rid={rid}"),
            None => format!("https://grok.com/c/{cid}"),
        };
        self.session
            .navigate(&target)
            .await
            .map_err(|e| e.to_string())?;
        let _ = self.session.wait_for_network_idle(15_000, 800).await;
        Ok(())
    }

    /// Ensures the `GrokClient` is initialised. On the first call this
    /// triggers `prepare()` (one page reload). Returns an `Arc` clone so
    /// the caller can drop the lock before making async network calls.
    async fn ensure_client(&self) -> Result<Arc<GrokClient>, String> {
        let mut guard = self.client.lock().await;
        if guard.is_none() {
            let statsig = std::env::var("TAIL_FIN_GROK_STATSIG_ID").unwrap_or_default();
            let c = GrokClient::new(self.session.clone()).with_statsig_id(statsig);
            c.prepare().await.map_err(|e| e.to_string())?;
            *guard = Some(Arc::new(c));
        }
        Ok(guard.as_ref().expect("just initialised").clone())
    }
}

fn build_ask_output(
    resp: tail_fin_grok::GrokResponse,
    debug: Vec<tail_fin_grok::ImageAttachmentDebug>,
    include_debug: bool,
) -> serde_json::Value {
    let mut out = json!({
        "response": resp.response,
        "conversation_id": resp.conversation_id,
        "response_id": resp.response_id,
    });
    if include_debug {
        out["debug"] = json!({
            "image_count": debug.len(),
            "images": debug,
        });
    }
    out
}

#[async_trait]
impl SiteHandler for GrokHandler {
    async fn handle(&self, id: &str, cmd: &str, params: &Value) -> Response {
        match cmd {
            "grok.ask" => {
                let prompt = match required_str(params, "prompt") {
                    Ok(p) => p,
                    Err(e) => return err_str(id, e),
                };
                let cid = match required_nonempty_str(params, "cid") {
                    Ok(c) => c,
                    Err(e) => {
                        return err_str(
                            id,
                            format!("{e} — start the conversation manually in Chrome first"),
                        )
                    }
                };

                // rid: prefer explicit param; fall back to tab URL.
                let rid_param = params.get("rid").and_then(|v| v.as_str()).unwrap_or("");

                if let Err(e) = self.ensure_on_conversation(&cid, Some(rid_param)).await {
                    return err_str(id, e);
                }

                // Lock released before any async network calls.
                let client = match self.ensure_client().await {
                    Ok(c) => c,
                    Err(e) => return err_str(id, e),
                };

                let rid = if rid_param.is_empty() {
                    match client.read_rid_from_tab().await {
                        Ok(r) => r,
                        Err(e) => return err_str(id, e.to_string()),
                    }
                } else {
                    rid_param.to_string()
                };

                let image_paths = match optional_string_array(params, "images") {
                    Ok(v) => v,
                    Err(e) => return err_str(id, e),
                };
                let image_debug = optional_bool(params, "image_debug", true);
                let max_image_mb = match optional_positive_usize(params, "max_image_mb") {
                    Ok(v) => v,
                    Err(e) => return err_str(id, e),
                };
                let max_total_image_mb = match optional_positive_usize(params, "max_total_image_mb")
                {
                    Ok(v) => v,
                    Err(e) => return err_str(id, e),
                };

                match client
                    .ask_continue_with_ids_and_images_debug_with_limits(
                        &prompt,
                        &cid,
                        &rid,
                        &image_paths,
                        max_image_mb,
                        max_total_image_mb,
                    )
                    .await
                {
                    Ok((resp, debug)) => ok_value(id, build_ask_output(resp, debug, image_debug)),
                    Err(e) => err_str(id, e.to_string()),
                }
            }
            "grok.conversations" => {
                // Reuse cached client when available; list_conversations needs no statsig
                // but avoids creating a redundant session wrapper.
                let client = {
                    let guard = self.client.lock().await;
                    guard
                        .as_ref()
                        .map(Arc::clone)
                        .unwrap_or_else(|| Arc::new(GrokClient::new(self.session.clone())))
                };
                match client.list_conversations().await {
                    Ok(convs) => ok_json(id, &json!({"conversations": convs})),
                    Err(e) => err_str(id, e.to_string()),
                }
            }
            other => err_str(id, format!("unknown grok cmd: {other}")),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn build_ask_output_hides_debug_when_disabled() {
        let resp = tail_fin_grok::GrokResponse {
            response: "ok".into(),
            conversation_id: "c1".into(),
            response_id: "r1".into(),
        };
        let out = build_ask_output(resp, vec![], false);
        assert!(out.get("debug").is_none());
    }

    #[test]
    fn build_ask_output_includes_debug_when_enabled() {
        let resp = tail_fin_grok::GrokResponse {
            response: "ok".into(),
            conversation_id: "c1".into(),
            response_id: "r1".into(),
        };
        let dbg = tail_fin_grok::ImageAttachmentDebug {
            original_path: "/tmp/a.png".into(),
            sent_name: "a.jpg".into(),
            original_bytes: 100,
            sent_bytes: 80,
            mime: "image/jpeg".into(),
            compressed: true,
        };
        let out = build_ask_output(resp, vec![dbg], true);
        assert_eq!(out["debug"]["image_count"].as_u64(), Some(1));
    }
}