owlen/xtask/src/screenshots.rs

use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

use anyhow::{Context, Result, bail};
use async_trait::async_trait;
use crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
use futures_util::FutureExt;
use futures_util::future::BoxFuture;
use owlen_core::{
    Config, McpMode, Mode, Provider,
    session::SessionController,
    storage::StorageManager,
    types::{Message, ModelInfo, ToolCall},
    ui::{NoOpUiController, UiController},
};
use owlen_tui::ChatApp;
use owlen_tui::events::Event;
use owlen_tui::ui::render_chat;
use ratatui::{Terminal, backend::TestBackend};
use serde_json::json;
use tokio::runtime::Runtime;
use tokio::sync::mpsc;
use which::which;

use crate::workspace_root;

#[derive(Clone, Copy)]
struct SceneSpec {
    name: &'static str,
    width: u16,
    height: u16,
    builder: fn() -> BoxFuture<'static, Result<ChatApp>>,
}

pub(crate) fn run(output: Option<PathBuf>, chafa: Option<PathBuf>, no_png: bool) -> Result<()> {
    let output_dir = output.unwrap_or_else(|| workspace_root().join("images/generated"));
    fs::create_dir_all(&output_dir)
        .with_context(|| format!("failed to create output directory {}", output_dir.display()))?;

    let config_home = workspace_root().join("dev/xtask-config");
    fs::create_dir_all(&config_home).with_context(|| {
        format!(
            "failed to create config home directory {}",
            config_home.display()
        )
    })?;
    unsafe {
        // Safe: we control the value and keep it within the workspace for deterministic snapshots.
        env::set_var("XDG_CONFIG_HOME", &config_home);
        env::set_var("XDG_DATA_HOME", &config_home);
    }

    let png_requested = !no_png;
    let mut png_skip_reason: Option<String> = None;
    let chafa_binary: Option<PathBuf> = if png_requested {
        let candidate = if let Some(ref path) = chafa {
            if path.exists() {
                Some(path.clone())
            } else {
                bail!("specified chafa binary '{}' does not exist", path.display());
            }
        } else {
            which("chafa").ok()
        };

        if let Some(path) = candidate {
            match chafa_supports_save(&path) {
                Ok(true) => Some(path),
                Ok(false) => {
                    png_skip_reason = Some(format!(
                        "chafa at '{}' lacks --save support; install chafa 1.14+ or rerun with --no-png",
                        path.display()
                    ));
                    None
                }
                Err(err) => {
                    eprintln!("warning: failed to probe chafa capabilities: {err}");
                    png_skip_reason = Some("failed to verify chafa capabilities".to_string());
                    None
                }
            }
        } else {
            png_skip_reason =
                Some("chafa not found in PATH; install it or rerun with --no-png".to_string());
            None
        }
    } else {
        None
    };

    let runtime = Runtime::new().context("failed to create tokio runtime")?;

    for scene in scenes() {
        let mut app = runtime.block_on((scene.builder)())?;
        let ansi = render_snapshot(&mut app, scene.width, scene.height)?;
        let ans_path = output_dir.join(format!("{}.ans", scene.name));
        fs::write(&ans_path, ansi.as_bytes())
            .with_context(|| format!("failed to write ANSI dump {}", ans_path.display()))?;

        if let Some(ref binary) = chafa_binary
            && let Err(err) = convert_to_png(
                &ans_path,
                scene.width,
                scene.height,
                &output_dir,
                scene.name,
                binary.as_path(),
            )
        {
            eprintln!("warning: {}", err);
        }
    }

    println!("Screenshots written to {}", output_dir.display());
    if png_requested && chafa_binary.is_none() {
        if let Some(reason) = png_skip_reason {
            println!("PNG conversion skipped ({reason})");
        } else {
            println!("PNG conversion skipped");
        }
    } else if !png_requested {
        println!("PNG conversion skipped (use --no-png=false or omit flag to enable)");
    }

    Ok(())
}

fn convert_to_png(
    ans_path: &Path,
    width: u16,
    height: u16,
    output_dir: &Path,
    name: &str,
    chafa_path: &Path,
) -> Result<()> {
    let png_path = output_dir.join(format!("{}.png", name));
    let mut command = Command::new(chafa_path);
    let status = command
        .arg("--size")
        .arg(format!("{}x{}", width, height))
        .arg("--save")
        .arg(&png_path)
        .arg(ans_path)
        .status()
        .with_context(|| "failed to spawn chafa".to_string())?;
    if !status.success() {
        bail!(
            "chafa exited with status {} when rendering {}",
            status.code().unwrap_or_default(),
            name
        );
    }
    Ok(())
}

fn chafa_supports_save(chafa_path: &Path) -> Result<bool> {
    let output = Command::new(chafa_path)
        .arg("--help")
        .output()
        .with_context(|| {
            format!(
                "failed to run '{}' to probe capabilities",
                chafa_path.display()
            )
        })?;
    let help = String::from_utf8_lossy(&output.stdout);
    Ok(help.contains("--save"))
}

fn scenes() -> &'static [SceneSpec] {
    &[
        SceneSpec {
            name: "chat-idle-dark",
            width: 120,
            height: 32,
            builder: || async move { Ok(build_chat_app(|_| {}, |_| {}).await) }.boxed(),
        },
        SceneSpec {
            name: "chat-tool-call",
            width: 120,
            height: 32,
            builder: || {
                async move {
                let app = build_chat_app(
                    |_| {},
                    |session| {
                        let conversation = session.conversation_mut();
                        conversation.push_user_message("What happened in the Rust ecosystem today?");
                        let stream_id = conversation.start_streaming_response();
                        conversation
                            .set_stream_placeholder(stream_id, "Consulting the knowledge base…")
                            .expect("placeholder");
                        let tool_call = ToolCall {
                            id: "call-search-1".into(),
                            name: "web_search".into(),
                            arguments: json!({ "query": "Rust language news" }),
                        };
                        conversation
                            .set_tool_calls_on_message(stream_id, vec![tool_call.clone()])
                            .expect("tool call metadata");
                        conversation
                            .append_stream_chunk(stream_id, "Found multiple articles…", false)
                            .expect("stream chunk");
                        conversation.push_message(Message::tool(
                            tool_call.id.clone(),
                            "Rust 1.85 released with generics cleanups and faster async compilation.".into(),
                        ));
                        conversation.push_message(Message::assistant(
                            "Summarising the latest Rust release and the async runtime updates.".into(),
                        ));
                    },
                )
                .await;
                Ok(app)
            }
            .boxed()
            },
        },
        SceneSpec {
            name: "command-palette",
            width: 110,
            height: 28,
            builder: || {
                async move {
                    let mut app = build_chat_app(|_| {}, |_| {}).await;
                    send_key(&mut app, KeyCode::Char(':'), KeyModifiers::NONE).await;
                    type_text(&mut app, "focus").await;
                    send_key(&mut app, KeyCode::Down, KeyModifiers::NONE).await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "guidance-onboarding",
            width: 100,
            height: 28,
            builder: || {
                async move {
                    let app = build_chat_app(
                        |cfg| {
                            cfg.ui.show_onboarding = true;
                            cfg.ui.guidance.coach_marks_complete = false;
                        },
                        |_| {},
                    )
                    .await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "guidance-cheatsheet",
            width: 120,
            height: 30,
            builder: || {
                async move {
                    let mut app =
                        build_chat_app(|cfg| cfg.ui.guidance.coach_marks_complete = true, |_| {})
                            .await;
                    send_key(&mut app, KeyCode::Char('?'), KeyModifiers::NONE).await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "editing-mode",
            width: 120,
            height: 32,
            builder: || {
                async move {
                    let mut app = build_chat_app(|_| {}, |_| {}).await;
                    send_key(&mut app, KeyCode::Char('i'), KeyModifiers::NONE).await;
                    type_text(&mut app, "Editing mode demonstration").await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "visual-mode",
            width: 120,
            height: 32,
            builder: || {
                async move {
                    let mut app = build_chat_app(
                        |_| {},
                        |session| {
                            let conversation = session.conversation_mut();
                            conversation.push_user_message(
                                "Render visual selection across multiple lines.",
                            );
                            conversation.push_message(Message::assistant(
                                "Assistant reply for visual mode highlighting.".into(),
                            ));
                        },
                    )
                    .await;
                    send_key(&mut app, KeyCode::Char('v'), KeyModifiers::NONE).await;
                    send_key(&mut app, KeyCode::Char('j'), KeyModifiers::NONE).await;
                    send_key(&mut app, KeyCode::Char('j'), KeyModifiers::NONE).await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "accessibility-high-contrast",
            width: 120,
            height: 32,
            builder: || {
                async move {
                    let app = build_chat_app(
                        |cfg| {
                            cfg.ui.accessibility.high_contrast = true;
                            cfg.ui.guidance.coach_marks_complete = true;
                        },
                        |_| {},
                    )
                    .await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "accessibility-reduced-chrome",
            width: 120,
            height: 32,
            builder: || {
                async move {
                    let app = build_chat_app(
                        |cfg| {
                            cfg.ui.accessibility.reduced_chrome = true;
                            cfg.ui.guidance.coach_marks_complete = true;
                        },
                        |_| {},
                    )
                    .await;
                    Ok(app)
                }
                .boxed()
            },
        },
        SceneSpec {
            name: "emacs-profile",
            width: 120,
            height: 32,
            builder: || {
                async move {
                    let mut app = build_chat_app(
                        |cfg| {
                            cfg.ui.keymap_profile = Some("emacs".into());
                            cfg.ui.guidance.coach_marks_complete = true;
                        },
                        |_| {},
                    )
                    .await;
                    send_key(&mut app, KeyCode::Char(':'), KeyModifiers::NONE).await;
                    type_text(&mut app, "help").await;
                    Ok(app)
                }
                .boxed()
            },
        },
    ]
}

fn render_snapshot(app: &mut ChatApp, width: u16, height: u16) -> Result<String> {
    let backend = TestBackend::new(width, height);
    let mut terminal = Terminal::new(backend).context("failed to create test terminal")?;
    terminal
        .draw(|frame| render_chat(frame, app))
        .context("failed to render chat frame")?;
    let buffer = terminal.backend().buffer();
    Ok(buffer_to_string(buffer))
}

fn buffer_to_string(buffer: &ratatui::buffer::Buffer) -> String {
    let mut output = String::new();
    for y in 0..buffer.area.height {
        for x in 0..buffer.area.width {
            output.push_str(buffer[(x, y)].symbol());
        }
        output.push('\n');
    }
    output
}

async fn send_key(app: &mut ChatApp, code: KeyCode, modifiers: KeyModifiers) {
    app.handle_event(Event::Key(KeyEvent::new(code, modifiers)))
        .await
        .expect("send key event");
}

async fn type_text(app: &mut ChatApp, text: &str) {
    for ch in text.chars() {
        send_key(app, KeyCode::Char(ch), KeyModifiers::NONE).await;
    }
}

async fn build_chat_app<C, F>(configure_config: C, configure_session: F) -> ChatApp
where
    C: FnOnce(&mut Config) + Send + 'static,
    F: FnOnce(&mut SessionController) + Send + 'static,
{
    let temp_dir = tempfile::tempdir().expect("temp dir");
    let storage =
        StorageManager::with_database_path(temp_dir.path().join("owlen-tui-screenshots.db"))
            .await
            .expect("storage");
    let storage = std::sync::Arc::new(storage);

    let mut config = Config::default();
    configure_config(&mut config);
    config.general.default_model = Some("stub-model".into());
    config.general.enable_streaming = true;
    config.privacy.encrypt_local_data = false;
    config.privacy.require_consent_per_session = false;
    config.ui.show_timestamps = false;
    config.mcp.mode = McpMode::LocalOnly;
    config.mcp.allow_fallback = false;
    config.mcp.warn_on_legacy = false;

    let provider: std::sync::Arc<dyn Provider> = std::sync::Arc::new(StubProvider);
    let ui: std::sync::Arc<dyn UiController> = std::sync::Arc::new(NoOpUiController);
    let (event_tx, controller_event_rx) = mpsc::unbounded_channel();

    let mut session = SessionController::new(provider, config, storage, ui, true, Some(event_tx))
        .await
        .expect("session controller");

    session
        .set_operating_mode(Mode::Chat)
        .await
        .expect("chat mode");

    configure_session(&mut session);

    let (app, mut session_rx) = ChatApp::new(session, controller_event_rx)
        .await
        .expect("chat app");
    session_rx.close();

    app
}

#[derive(Default)]
struct StubProvider;

#[async_trait]
impl Provider for StubProvider {
    fn name(&self) -> &str {
        "stub-provider"
    }

    async fn list_models(&self) -> owlen_core::Result<Vec<owlen_core::types::ModelInfo>> {
        Ok(vec![ModelInfo {
            id: "stub-model".into(),
            name: "Stub Model".into(),
            description: Some("Stub model for screenshot generation".into()),
            provider: self.name().into(),
            context_window: Some(8_192),
            capabilities: vec!["chat".into(), "tool-use".into()],
            supports_tools: true,
        }])
    }

    async fn send_prompt(
        &self,
        _request: owlen_core::types::ChatRequest,
    ) -> owlen_core::Result<owlen_core::types::ChatResponse> {
        Ok(owlen_core::types::ChatResponse {
            message: Message::assistant("stub completion".into()),
            usage: None,
            is_streaming: false,
            is_final: true,
        })
    }

    async fn stream_prompt(
        &self,
        _request: owlen_core::types::ChatRequest,
    ) -> owlen_core::Result<owlen_core::ChatStream> {
        Ok(Box::pin(futures_util::stream::empty()))
    }

    async fn health_check(&self) -> owlen_core::Result<()> {
        Ok(())
    }

    fn as_any(&self) -> &(dyn std::any::Any + Send + Sync) {
        self
    }
}