magi-code 0.96.1

use super::{DisplayLine, DisplayRole, DisplaySpan, RenderInputKind, render};
use pulldown_cmark::{Alignment, CodeBlockKind, Event, Options, Parser, Tag, TagEnd};

/// Semantic references only: no filesystem access and no terminal control codes.
/// Keep the work and number of previews bounded independently of transcript size.
pub(crate) fn image_references(markdown: &str) -> Vec<String> {
    if markdown.len() > 64 * 1024 {
        return Vec::new();
    }
    Parser::new(markdown)
        .filter_map(|event| match event {
            Event::Start(Tag::Image { dest_url, .. }) if dest_url.len() <= 4096 => {
                Some(dest_url.into_string())
            }
            _ => None,
        })
        .take(4)
        .collect()
}

struct MarkdownRenderer {
    lines: Vec<DisplayLine>,
    current: Vec<DisplaySpan>,
    list_stack: Vec<Option<u64>>,
    blockquote_depth: usize,
    heading_depth: u8,
    in_code_block: Option<Option<String>>,
    code: String,
    table_alignments: Vec<Alignment>,
    table_row: Vec<Vec<DisplaySpan>>,
    table_rows: Vec<Vec<Vec<DisplaySpan>>>,
    link_url_stack: Vec<String>,
}

impl MarkdownRenderer {
    fn new() -> Self {
        Self {
            lines: Vec::new(),
            current: Vec::new(),
            list_stack: Vec::new(),
            blockquote_depth: 0,
            heading_depth: 0,
            in_code_block: None,
            code: String::new(),
            table_alignments: Vec::new(),
            table_row: Vec::new(),
            table_rows: Vec::new(),
            link_url_stack: Vec::new(),
        }
    }

    fn handle_event(&mut self, event: Event) {
        match event {
            Event::Start(Tag::Paragraph) => {
                flush_current(&mut self.lines, &mut self.current);
                if self.blockquote_depth > 0 {
                    self.current.push(DisplaySpan::new(
                        "> ".repeat(self.blockquote_depth),
                        DisplayRole::BlockQuote,
                    ));
                }
            }
            Event::End(TagEnd::Paragraph) => {
                flush_current(&mut self.lines, &mut self.current);
                if self.list_stack.is_empty() {
                    push_blank_if_needed(&mut self.lines);
                }
            }
            Event::Start(Tag::Heading { level, .. }) => {
                flush_current(&mut self.lines, &mut self.current);
                self.heading_depth = level as u8;
                self.current.push(DisplaySpan::new(
                    "#".repeat(self.heading_depth as usize),
                    DisplayRole::Heading,
                ));
                self.current
                    .push(DisplaySpan::new(" ", DisplayRole::Heading));
            }
            Event::End(TagEnd::Heading(_)) => {
                set_role(&mut self.current, DisplayRole::Heading);
                flush_current(&mut self.lines, &mut self.current);
                self.heading_depth = 0;
                push_blank_if_needed(&mut self.lines);
            }
            Event::Start(Tag::List(start)) => self.list_stack.push(start),
            Event::End(TagEnd::List(_)) => {
                self.list_stack.pop();
                push_blank_if_needed(&mut self.lines);
            }
            Event::Start(Tag::Item) => {
                flush_current(&mut self.lines, &mut self.current);
                let marker = match self.list_stack.last_mut() {
                    Some(Some(next)) => {
                        let marker = format!("{}. ", *next);
                        *next = next.saturating_add(1);
                        marker
                    }
                    _ => "• ".to_string(),
                };
                self.current
                    .push(DisplaySpan::new(marker, DisplayRole::ListMarker));
            }
            Event::End(TagEnd::Item) => flush_current(&mut self.lines, &mut self.current),
            Event::Start(Tag::BlockQuote(_)) => {
                flush_current(&mut self.lines, &mut self.current);
                self.blockquote_depth = self.blockquote_depth.saturating_add(1);
                self.current.push(DisplaySpan::new(
                    "> ".repeat(self.blockquote_depth),
                    DisplayRole::BlockQuote,
                ));
            }
            Event::End(TagEnd::BlockQuote(_)) => {
                flush_current(&mut self.lines, &mut self.current);
                self.blockquote_depth = self.blockquote_depth.saturating_sub(1);
                push_blank_if_needed(&mut self.lines);
            }
            Event::Start(Tag::Table(alignments)) => {
                flush_current(&mut self.lines, &mut self.current);
                self.table_alignments = alignments;
                self.table_row.clear();
            }
            Event::End(TagEnd::Table) => {
                let table = super::DisplayTable {
                    alignments: std::mem::take(&mut self.table_alignments)
                        .into_iter()
                        .map(|alignment| match alignment {
                            Alignment::None => super::TableAlignment::None,
                            Alignment::Left => super::TableAlignment::Left,
                            Alignment::Center => super::TableAlignment::Center,
                            Alignment::Right => super::TableAlignment::Right,
                        })
                        .collect(),
                    rows: std::mem::take(&mut self.table_rows),
                };
                self.lines.push(table_display_line(table));
                push_blank_if_needed(&mut self.lines);
            }
            Event::Start(Tag::TableHead | Tag::TableRow) => {
                flush_current(&mut self.lines, &mut self.current);
                self.table_row.clear();
            }
            Event::End(TagEnd::TableHead | TagEnd::TableRow) => {
                self.table_rows.push(std::mem::take(&mut self.table_row));
            }
            Event::Start(Tag::TableCell) => {
                self.current.clear();
            }
            Event::End(TagEnd::TableCell) => {
                self.table_row.push(std::mem::take(&mut self.current));
            }
            Event::Start(Tag::CodeBlock(kind)) => {
                flush_current(&mut self.lines, &mut self.current);
                self.in_code_block = Some(match kind {
                    CodeBlockKind::Fenced(language) => {
                        let language = language.trim().to_string();
                        if language.is_empty() {
                            None
                        } else {
                            Some(language)
                        }
                    }
                    CodeBlockKind::Indented => None,
                });
                self.code.clear();
            }
            Event::End(TagEnd::CodeBlock) => {
                let language = self.in_code_block.take().flatten();
                self.lines.push(DisplayLine::from_span(
                    format!("```{}", language.as_deref().unwrap_or("")),
                    DisplayRole::CodeFence,
                ));
                if let Some(language) = &language {
                    self.lines.push(DisplayLine::from_span(
                        format!("language: {language}"),
                        DisplayRole::CodeLanguageLabel,
                    ));
                }
                let rendered_code = if language.as_deref().is_some_and(is_diff_fence_language) {
                    render(self.code.trim_end_matches('\n'), RenderInputKind::Diff)
                } else {
                    render(
                        self.code.trim_end_matches('\n'),
                        RenderInputKind::Code { language },
                    )
                };
                self.lines.extend(add_code_gutter(rendered_code));
                self.lines
                    .push(DisplayLine::from_span("```", DisplayRole::CodeFence));
                push_blank_if_needed(&mut self.lines);
            }
            Event::Text(text) => {
                if self.in_code_block.is_some() {
                    self.code.push_str(&text);
                } else {
                    for (index, part) in text.split('\n').enumerate() {
                        if index > 0 {
                            flush_current(&mut self.lines, &mut self.current);
                            if self.blockquote_depth > 0 {
                                self.current.push(DisplaySpan::new(
                                    "> ".repeat(self.blockquote_depth),
                                    DisplayRole::BlockQuote,
                                ));
                            }
                        }
                        if !part.is_empty() {
                            let role = if self.heading_depth > 0 {
                                DisplayRole::Heading
                            } else {
                                DisplayRole::Plain
                            };
                            self.current.push(DisplaySpan::new(part, role));
                        }
                    }
                }
            }
            Event::Start(Tag::Link { dest_url, .. }) => {
                self.link_url_stack.push(dest_url.to_string());
            }
            Event::End(TagEnd::Link) => {
                if let Some(url) = self.link_url_stack.pop()
                    && !url.is_empty()
                {
                    self.current
                        .push(DisplaySpan::new(format!(" ({url})"), DisplayRole::Plain));
                }
            }
            Event::Code(text) => self
                .current
                .push(DisplaySpan::new(text, DisplayRole::InlineCode)),
            Event::SoftBreak => self.current.push(DisplaySpan::new(" ", DisplayRole::Plain)),
            Event::HardBreak => flush_current(&mut self.lines, &mut self.current),
            Event::Rule => {
                flush_current(&mut self.lines, &mut self.current);
                self.lines
                    .push(DisplayLine::from_span("───", DisplayRole::Plain));
            }
            _ => {}
        }
    }

    fn finish(mut self) -> Vec<DisplayLine> {
        flush_current(&mut self.lines, &mut self.current);
        trim_trailing_blank(&mut self.lines);
        self.lines
    }
}

pub(crate) fn render_markdown(input: &str) -> Vec<DisplayLine> {
    if input.is_empty() {
        return vec![DisplayLine::plain("")];
    }
    let lines = if input_has_unclosed_fence(input) {
        render_unclosed_fence_fallback(input)
    } else {
        render_with_pulldown(input)
    };
    if lines.is_empty() {
        super::render_plain(input)
    } else {
        lines
    }
}

fn render_with_pulldown(input: &str) -> Vec<DisplayLine> {
    let mut options = Options::empty();
    options.insert(Options::ENABLE_TABLES);
    let parser = Parser::new_ext(input, options);
    let mut renderer = MarkdownRenderer::new();
    for event in parser {
        renderer.handle_event(event);
    }
    renderer.finish()
}

fn render_unclosed_fence_fallback(input: &str) -> Vec<DisplayLine> {
    let lines: Vec<&str> = input.split('\n').collect();
    let Some(unclosed_start) = find_unclosed_fence_start(&lines) else {
        return render_fallback_fence_lines(&lines);
    };

    let mut rendered = Vec::new();
    if unclosed_start > 0 {
        rendered.extend(super::render_plain(&lines[..unclosed_start].join("\n")));
    }
    rendered.extend(render_fallback_fence_lines(&lines[unclosed_start..]));
    rendered
}

fn find_unclosed_fence_start(lines: &[&str]) -> Option<usize> {
    let mut opener: Option<(char, usize, usize)> = None;

    for (index, line) in lines.iter().enumerate() {
        match opener {
            Some((opener_char, opener_len, _)) => {
                if is_closing_fence(line, opener_char, opener_len) {
                    opener = None;
                }
            }
            None => {
                if let Some((marker_char, marker_len)) = parse_fence_marker(line) {
                    opener = Some((marker_char, marker_len, index));
                }
            }
        }
    }

    opener.map(|(_, _, start)| start)
}

fn render_fallback_fence_lines(input_lines: &[&str]) -> Vec<DisplayLine> {
    let mut lines = Vec::new();
    let mut opener: Option<(char, usize, Option<String>)> = None;
    let mut code = Vec::new();

    for line in input_lines {
        match opener.take() {
            Some((opener_char, opener_len, language)) => {
                if is_closing_fence(line, opener_char, opener_len) {
                    render_fallback_code_lines(&mut lines, &mut code, language);
                    lines.push(DisplayLine::from_span(*line, DisplayRole::CodeFence));
                } else {
                    code.push(*line);
                    opener = Some((opener_char, opener_len, language));
                }
            }
            None => {
                if let Some((marker_char, marker_len)) = parse_fence_marker(line) {
                    opener = Some((marker_char, marker_len, fence_language(line, marker_len)));
                    lines.push(DisplayLine::from_span(*line, DisplayRole::CodeFence));
                } else {
                    lines.push(DisplayLine::plain(*line));
                }
            }
        }
    }

    if let Some((_, _, language)) = opener
        && !code.is_empty()
    {
        render_fallback_code_lines(&mut lines, &mut code, language);
    }
    lines
}

fn fence_language(line: &str, marker_len: usize) -> Option<String> {
    let indent = line.chars().take_while(|c| *c == ' ').take(4).count();
    line[indent + marker_len..]
        .split_whitespace()
        .next()
        .filter(|s| !s.is_empty())
        .map(str::to_string)
}

fn render_fallback_code_lines(
    lines: &mut Vec<DisplayLine>,
    code: &mut Vec<&str>,
    language: Option<String>,
) {
    if code.is_empty() {
        return;
    }
    let code = std::mem::take(code).join("\n");
    let rendered_code = if language.as_deref().is_some_and(is_diff_fence_language) {
        render(&code, RenderInputKind::Diff)
    } else {
        render(&code, RenderInputKind::Code { language })
    };
    lines.extend(add_code_gutter(rendered_code));
}

fn add_code_gutter(lines: Vec<DisplayLine>) -> Vec<DisplayLine> {
    lines
        .into_iter()
        .map(|line| {
            let mut spans = Vec::with_capacity(line.spans.len().saturating_add(1));
            spans.push(DisplaySpan::new("│ ", DisplayRole::CodeBlockGutter));
            spans.extend(line.spans);
            DisplayLine { spans, table: None }
        })
        .collect()
}

fn is_diff_fence_language(language: &str) -> bool {
    matches!(language.trim(), "diff" | "patch")
}

fn input_has_unclosed_fence(input: &str) -> bool {
    let mut opener: Option<(char, usize)> = None;
    for line in input.split('\n') {
        match opener {
            Some((opener_char, opener_len)) => {
                if is_closing_fence(line, opener_char, opener_len) {
                    opener = None;
                }
            }
            None => {
                if let Some(parsed) = parse_fence_marker(line) {
                    opener = Some(parsed);
                }
            }
        }
    }
    opener.is_some()
}

/// Returns (marker_char, marker_len) if the line opens a CommonMark fence.
/// Recognizes 0..=3 leading spaces, backtick or tilde marker runs of length >= 3.
fn parse_fence_marker(line: &str) -> Option<(char, usize)> {
    let indent = line.chars().take_while(|c| *c == ' ').take(4).count();
    if indent >= 4 {
        return None;
    }
    let rest = &line[indent..];
    let marker_char = rest.chars().next()?;
    if marker_char != '`' && marker_char != '~' {
        return None;
    }
    let run_len = rest.chars().take_while(|c| *c == marker_char).count();
    if run_len < 3 {
        return None;
    }
    if marker_char == '`' {
        let info = &rest[run_len..];
        if info.contains('`') {
            return None;
        }
    }
    Some((marker_char, run_len))
}

/// Checks whether the line closes a fence opened with the given marker char and length.
fn is_closing_fence(line: &str, opener_char: char, opener_len: usize) -> bool {
    let indent = line.chars().take_while(|c| *c == ' ').take(4).count();
    if indent >= 4 {
        return false;
    }
    let rest = &line[indent..];
    if !rest.starts_with(opener_char) {
        return false;
    }
    let run_len = rest.chars().take_while(|c| *c == opener_char).count();
    if run_len < opener_len {
        return false;
    }
    let after = rest[run_len..].trim();
    after.is_empty()
}

fn table_display_line(table: super::DisplayTable) -> DisplayLine {
    let mut spans = Vec::new();
    for (index, cells) in table.rows.iter().enumerate() {
        if index > 0 {
            spans.push(DisplaySpan::new("\n", DisplayRole::Plain));
        }
        spans.push(DisplaySpan::new("|", DisplayRole::Plain));
        for cell in cells {
            spans.push(DisplaySpan::new(" ", DisplayRole::Plain));
            spans.extend(cell.iter().cloned());
            spans.push(DisplaySpan::new(" |", DisplayRole::Plain));
        }
        if index == 0 {
            spans.push(DisplaySpan::new("\n|", DisplayRole::Plain));
            for alignment in &table.alignments {
                let marker = match alignment {
                    super::TableAlignment::None => " --- |",
                    super::TableAlignment::Left => " :--- |",
                    super::TableAlignment::Center => " :---: |",
                    super::TableAlignment::Right => " ---: |",
                };
                spans.push(DisplaySpan::new(marker, DisplayRole::Plain));
            }
        }
    }
    DisplayLine {
        spans,
        table: Some(table),
    }
}

fn flush_current(lines: &mut Vec<DisplayLine>, current: &mut Vec<DisplaySpan>) {
    if !current.is_empty() {
        lines.push(DisplayLine {
            spans: std::mem::take(current),
            table: None,
        });
    }
}

fn set_role(spans: &mut [DisplaySpan], role: DisplayRole) {
    for span in spans {
        if span.role == DisplayRole::Plain {
            span.role = role;
        }
    }
}

fn push_blank_if_needed(lines: &mut Vec<DisplayLine>) {
    if !lines.last().is_some_and(|line| line.is_blank()) {
        lines.push(DisplayLine::plain(""));
    }
}

fn trim_trailing_blank(lines: &mut Vec<DisplayLine>) {
    while lines.last().is_some_and(|line| line.is_blank()) {
        lines.pop();
    }
}