//! Rich blocks in agent answers. Agents write diagrams and small UIs as plain
//! fenced code blocks (` ```mermaid `, ` ```openui `, ` ```svg `), so they
//! read fine anywhere else; here the Markdown view swaps them for native
//! renders. Each keeps a header row with Copy and Show source, so a render
//! that goes wrong never hides what the agent wrote.

mod mermaid;
mod openui;

use std::collections::HashMap;
use std::sync::Arc;
use std::hash::Hasher as _;

use gpui_kit::component::{
    Icon, IconName, h_flex,
    message_scroller::MessageScrollerState,
    text::{MarkdownNode, MarkdownParseContext, TextView, markdown_ast},
    v_flex,
};
use gpui_kit::prelude::FluentBuilder as _;
use gpui_kit::*;

use crate::inline_ui::lang::{self, Program};
use crate::session::Session;
use crate::ui::icons::Lucide;
use crate::ui::workspace::Workspace;
use crate::ui::metal::{self, MetalButton, MetalStyled as _, Palette};

/// The one custom node name every rich block parses into. A constant set of
/// parsers and names keeps the Markdown view from reparsing every frame.
const NODE: &str = "inline-ui";

/// Where a Markdown view sits, so a block's renderer knows which thread it
/// belongs to and which transcript row to remeasure when its height changes.
#[derive(Clone)]
pub struct InlineCx {
    pub session: WeakEntity<Session>,
    pub ix: usize,
    pub scroller: WeakEntity<MessageScrollerState>,
}

impl InlineCx {
    /// Tells the transcript a block changed height, so the virtual list
    /// doesn't jump on the next frame.
    fn remeasure(&self, cx: &mut App) {
        let ix = self.ix;
        if let Some(scroller) = self.scroller.upgrade() {
            scroller.update(cx, |scroller, cx| {
                scroller.remeasure_items(ix..ix + 1, cx);
            });
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
enum Lang {
    Mermaid,
    OpenUi,
    Svg,
}

impl Lang {
    fn parse(lang: &str) -> Option<Self> {
        match lang.trim().to_ascii_lowercase().as_str() {
            "mermaid" => Some(Self::Mermaid),
            "openui" | "openui-lang" => Some(Self::OpenUi),
            "svg" => Some(Self::Svg),
            _ => None,
        }
    }

    fn label(self) -> &'static str {
        match self {
            Self::Mermaid => "Diagram",
            Self::OpenUi => "Interactive",
            Self::Svg => "Image",
        }
    }

    fn fence(self) -> &'static str {
        match self {
            Self::Mermaid => "mermaid",
            Self::OpenUi => "openui",
            Self::Svg => "svg",
        }
    }
}

/// A fenced block as the parser hands it to the renderer. Parsing runs on a
/// background thread, so this holds only plain data.
struct Fence {
    lang: Lang,
    code: String,
    /// Whether the closing fence has arrived. Streaming answers show
    /// unfinished blocks as skeletons, and never as errors.
    closed: bool,
    /// OpenUI's program, parsed here off the main thread.
    program: Option<Arc<Program>>,
}

/// Whether answers render rich blocks at all.
pub fn enabled() -> bool {
    true
}

/// The workspace, so a block's buttons can open previews in its panels.
struct Hub(WeakEntity<Workspace>);

impl Global for Hub {}

pub fn set_workspace(workspace: WeakEntity<Workspace>, cx: &mut App) {
    cx.set_global(Hub(workspace));
}

/// Opens an HTML block as a page: written to the data folder's
/// `artifacts`, then shown in the thread's browser panel, or the system
/// browser without one.
pub fn preview_html(code: &str, window: &mut Window, cx: &mut App) {
    let Some(dir) = crate::settings::data_dir().map(|dir| dir.join("artifacts")) else { return };
    let path = dir.join(format!("{:016x}.html", hash(code)));
    if let Err(err) = std::fs::create_dir_all(&dir).and_then(|_| std::fs::write(&path, code)) {
        eprintln!("vectrum-code: couldn't write the preview {}: {err}", path.display());
        return;
    }
    let url = format!("file://{}", path.display());
    match cx.try_global::<Hub>().and_then(|hub| hub.0.upgrade()) {
        Some(workspace) => workspace.update(cx, |workspace, cx| workspace.preview(url, window, cx)),
        None => crate::ui::workspace::open_url(&url),
    }
}

/// Adds the rich block hooks to an answer's Markdown view.
pub fn extend(view: TextView, icx: InlineCx) -> TextView {
    view.markdown_block_parser(parse).markdown_block_renderer(NODE, move |node, window, cx| render(node, &icx, window, cx))
}

fn parse(node: &markdown_ast::Node, cx: &MarkdownParseContext) -> Option<MarkdownNode> {
    let markdown_ast::Node::Code(code) = node else {
        return None;
    };
    let lang = Lang::parse(code.lang.as_deref()?)?;
    let source = cx.node_source(node)?;
    let fence = Fence {
        lang,
        code: code.value.clone(),
        closed: is_closed(source),
        program: (lang == Lang::OpenUi).then(|| Arc::new(lang::parse(&code.value))),
    };
    Some(MarkdownNode::new(NODE, fence).text(source.to_string()).markdown(source.to_string()))
}

/// Whether a fence's source ends with a closing fence. An unclosed fence runs
/// to the end of the document, which is what a block mid-stream looks like.
fn is_closed(source: &str) -> bool {
    let mut lines = source.trim_end().lines();
    let Some(open) = lines.next() else { return false };
    let open = open.trim_start();
    let Some(marker) = open.chars().next().filter(|c| *c == '`' || *c == '~') else {
        return false;
    };
    let width = open.chars().take_while(|c| *c == marker).count();
    lines.next_back().is_some_and(|close| {
        let close = close.trim();
        close.len() >= width && close.chars().all(|c| c == marker)
    })
}

fn hash(value: impl std::hash::Hash) -> u64 {
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    value.hash(&mut hasher);
    hasher.finish()
}

/// What blocks remember across frames. Rows of the virtual list drop their
/// element state when they scroll away, so this lives in a global.
#[derive(Default)]
struct Store {
    blocks: HashMap<Key, BlockState>,
}

/// A block: its thread, and a hash of where it starts and what it holds.
type Key = (EntityId, u64);

impl Global for Store {}

#[derive(Default)]
struct BlockState {
    show_source: bool,
    ui: openui::UiState,
}

/// A rendered block body, or why it couldn't be drawn.
enum Body {
    Ready(AnyElement),
    Failed(String),
}

fn render(node: &MarkdownNode, icx: &InlineCx, window: &mut Window, cx: &mut App) -> AnyElement {
    let Some(fence) = node.data::<Fence>() else {
        return div().into_any_element();
    };
    let p = metal::palette(cx);
    let key = (icx.session.entity_id(), hash((node.source_range().map(|range| range.start), &fence.code)));
    let id = ElementId::NamedInteger("inline-ui".into(), key.1);
    let show_source = cx.try_global::<Store>().and_then(|store| store.blocks.get(&key)).is_some_and(|block| block.show_source);
    let body = if show_source {
        Body::Ready(source_view(id.clone(), fence, p))
    } else {
        match fence.lang {
            Lang::Mermaid => mermaid::render(id.clone(), fence, icx, window, cx),
            Lang::OpenUi => openui::render(&id, key, fence, icx, window, cx),
            Lang::Svg => mermaid::render_svg(id.clone(), fence, icx, window, cx),
        }
    };
    let error = match &body {
        Body::Failed(error) => Some(error.clone()),
        Body::Ready(_) => None,
    };
    let body = match body {
        Body::Ready(element) => element,
        Body::Failed(_) => source_view(id.clone(), fence, p),
    };
    v_flex()
        .w_full()
        .gap(px(6.))
        .child(header(id, fence, key, show_source, error.is_some(), icx, p))
        .child(body)
        .when_some(error, |this, error| {
            this.child(div().px(px(4.)).text_size(px(12.5)).text_color(p.red).truncate().child(error))
        })
        .into_any_element()
}

/// The row over a block: its kind, Show source and Copy.
fn header(
    id: ElementId,
    fence: &Fence,
    key: Key,
    show_source: bool,
    failed: bool,
    icx: &InlineCx,
    p: &Palette,
) -> impl IntoElement {
    let code = fence.code.clone();
    let icx = icx.clone();
    h_flex()
        .w_full()
        .h(px(22.))
        .px(px(4.))
        .gap(px(4.))
        .child(metal::engraved(fence.lang.label(), p))
        .child(div().flex_1())
        .when(!failed, |this| {
            this.child(
                MetalButton::new(ElementId::Name(format!("{id}-source").into()))
                    .ghost()
                    .mini()
                    .icon(if show_source { Icon::new(IconName::Eye) } else { Icon::new(Lucide::Code) })
                    .tooltip(if show_source { "Show rendered" } else { "Show source" })
                    .on_click(move |_, window, cx| {
                        let block = cx.default_global::<Store>().blocks.entry(key).or_default();
                        block.show_source = !block.show_source;
                        icx.remeasure(cx);
                        window.refresh();
                    }),
            )
        })
        .child(
            MetalButton::new(ElementId::Name(format!("{id}-copy").into()))
                .ghost()
                .mini()
                .icon(IconName::Copy)
                .tooltip("Copy source")
                .on_click(move |_, _, cx| cx.write_to_clipboard(ClipboardItem::new_string(code.clone()))),
        )
}

/// The block as the plain code it is, in the same well code blocks use.
fn source_view(id: ElementId, fence: &Fence, p: &Palette) -> AnyElement {
    let markdown = format!("```{}\n{}\n```", fence.lang.fence(), fence.code.trim_end_matches('\n'));
    TextView::markdown(ElementId::Name(format!("{id}-code").into()), markdown)
        .style(crate::ui::chat::markdown_style(p))
        .selectable(true)
        .into_any_element()
}

/// Rasterizes SVG at exactly the display's density. GPUI's own
/// `render_single_frame` draws at twice that for smoothing, which quadruples
/// the memory a large diagram takes for no visible gain at this size.
fn rasterize(renderer: &SvgRenderer, svg: &str, scale: f32) -> anyhow::Result<(std::sync::Arc<RenderImage>, Size<Pixels>)> {
    let (width, height) = svg_size(svg).ok_or_else(|| anyhow::anyhow!("the SVG has no size"))?;
    anyhow::ensure!(width * height * scale * scale <= MAX_PIXELS, "the image is too large");
    let parsed = renderer.parse_svg(svg.as_bytes())?;
    let pixels = size(DevicePixels((width * scale).ceil() as i32), DevicePixels((height * scale).ceil() as i32));
    let image = renderer.render_parsed(&parsed, SvgSize::ExactSize(pixels))?;
    Ok((image, size(px(width), px(height))))
}

/// The largest bitmap a block may rasterize to, in device pixels.
const MAX_PIXELS: f32 = 6000. * 6000.;

/// The width and height an SVG asks for: its `width` and `height`
/// attributes, else its `viewBox`.
fn svg_size(svg: &str) -> Option<(f32, f32)> {
    let start = svg.find("<svg")?;
    let tag = &svg[start..start + svg[start..].find('>')?];
    let attr = |name: &str| -> Option<&str> {
        let mut rest = tag;
        loop {
            let at = rest.find(name)?;
            let before = rest[..at].chars().next_back();
            rest = &rest[at + name.len()..];
            if before.is_some_and(char::is_whitespace) && rest.trim_start().starts_with('=') {
                let value = rest.trim_start()[1..].trim_start();
                let quote = value.chars().next().filter(|c| *c == '"' || *c == '\'')?;
                return value[1..].split(quote).next();
            }
        }
    };
    let length = |value: &str| -> Option<f32> {
        let value = value.trim().trim_end_matches("px");
        value.parse::<f32>().ok().filter(|value| value.is_finite() && *value > 0.)
    };
    let view_box = attr("viewBox").and_then(|value| {
        let numbers: Vec<f32> = value.split([' ', ',']).filter(|part| !part.is_empty()).filter_map(|part| part.parse().ok()).collect();
        (numbers.len() == 4 && numbers[2] > 0. && numbers[3] > 0.).then(|| (numbers[2], numbers[3]))
    });
    let width = attr("width").and_then(length);
    let height = attr("height").and_then(length);
    match (width, height, view_box) {
        (Some(width), Some(height), _) => Some((width, height)),
        (Some(width), None, Some((w, h))) => Some((width, width * h / w)),
        (None, Some(height), Some((w, h))) => Some((height * w / h, height)),
        (_, _, view_box) => view_box,
    }
}

/// A placeholder while a block streams in or renders.
fn skeleton(label: &'static str, height: Pixels, p: &Palette) -> AnyElement {
    div()
        .w_full()
        .h(height)
        .flex()
        .items_center()
        .justify_center()
        .rounded(px(14.))
        .region(p)
        .text_size(px(12.5))
        .text_color(p.ink3)
        .child(label)
        .into_any_element()
}

#[cfg(test)]
mod tests {
    use super::{is_closed, svg_size};

    #[::core::prelude::v1::test]
    fn svg_sizes() {
        assert_eq!(svg_size(r#"<svg xmlns="x" width="120" height="80">"#), Some((120., 80.)));
        assert_eq!(svg_size(r#"<svg viewBox="0 0 300 150" width="600px">"#), Some((600., 300.)));
        assert_eq!(svg_size(r#"<?xml?><svg stroke-width="2" viewBox='0,0,40,20'>"#), Some((40., 20.)));
        assert_eq!(svg_size(r#"<svg width="100%" viewBox="0 0 40 20">"#), Some((40., 20.)));
        assert_eq!(svg_size("<svg>"), None);
    }

    #[::core::prelude::v1::test]
    fn closed_fences() {
        assert!(is_closed("```mermaid\ngraph TD\n```"));
        assert!(is_closed("````mermaid\ngraph TD\n````\n"));
        assert!(is_closed("~~~svg\n<svg/>\n~~~"));
        assert!(!is_closed("```mermaid\ngraph TD"));
        assert!(!is_closed("```mermaid"));
        assert!(!is_closed("````mermaid\ngraph TD\n```"));
    }
}