use heck::ToTitleCase as _;
use pretty::{Arena, DocAllocator, DocBuilder};

use crate::{
    ast::{
        Alignment, Block, Callout, CodeBlock, FootnoteDefinition, HeadingKind, List,
        ListBulletKind, ListItem, ListKind, ListOrderedKindOptions, SetextHeading, Table, TableRow,
        TaskState,
    },
    render::{Config, Context, ToDoc, highlighter, util::ArenaExt as _},
};

impl<'a> ToDoc<'a> for Vec<Block> {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        arena.concat(
            self.iter()
                .map(|block| block.to_doc(config, context, arena)),
        )
    }
}

impl<'a> ToDoc<'a> for Vec<&Block> {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        arena.concat(
            self.iter()
                .map(|block| block.to_doc(config, context, arena)),
        )
    }
}

impl<'a> ToDoc<'a> for Block {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        match self {
            Self::Paragraph(inlines) => {
                let inner = arena.concat(
                    inlines
                        .iter()
                        .map(|inline| inline.to_doc(config, context, arena)),
                );
                arena.tag("p", vec![], inner)
            }
            // TODO: add heading id
            Self::Heading(v) => {
                let htag = match v.kind {
                    HeadingKind::Atx(1) | HeadingKind::Setext(SetextHeading::Level1) => "h1",
                    HeadingKind::Atx(2) | HeadingKind::Setext(SetextHeading::Level2) => "h2",
                    HeadingKind::Atx(3) => "h3",
                    HeadingKind::Atx(4) => "h4",
                    HeadingKind::Atx(5) => "h5",
                    HeadingKind::Atx(_) => "h6",
                };
                let inner = arena.concat(
                    v.content
                        .iter()
                        .map(|inline| inline.to_doc(config, context, arena)),
                );
                arena.tag(htag, vec![], inner)
            }
            Self::ThematicBreak => arena.tag("hr", vec![], arena.nil()),
            Self::BlockQuote(inner) => {
                let inner = arena.concat(
                    inner
                        .iter()
                        .map(|inline| inline.to_doc(config, context, arena)),
                );
                arena.tag("blockquote", vec![], inner)
            }
            Self::List(v) => v.to_doc(config, context, arena),
            Self::CodeBlock(v) => v.to_doc(config, context, arena),
            Self::HtmlBlock(html) => arena.text(html.clone()),
            Self::Table(v) => v.to_doc(config, context, arena),
            Self::FootnoteDefinition(def) => def.to_doc(config, context, arena),
            Self::Callout(callout) => callout.to_doc(config, context, arena),
            Self::Empty | Self::Definition(_) => arena.nil(),
        }
    }
}

impl<'a> ToDoc<'a> for List {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        let items = arena.concat(
            self.items
                .iter()
                .map(|item| item.to_doc(config, context, arena)),
        );
        match self.kind {
            ListKind::Ordered(ListOrderedKindOptions { start }) => {
                arena.tag("ol", vec![("start".to_owned(), format!("{start}"))], items)
            }
            ListKind::Bullet(kind) => {
                let style = match kind {
                    ListBulletKind::Dash => "list-kind-dash",
                    ListBulletKind::Star => "list-kind-star",
                    ListBulletKind::Plus => "list-kind-plus",
                };
                arena.tag("ul", vec![("class".to_owned(), style.to_owned())], items)
            }
        }
    }
}

impl<'a> ToDoc<'a> for ListItem {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        let task = match self.task {
            Some(TaskState::Complete) => arena.tag(
                "span",
                vec![("class".to_owned(), "list-task-complete".to_owned())],
                arena.text("[X] "),
            ),
            Some(TaskState::Incomplete) => arena.tag(
                "span",
                vec![("class".to_owned(), "list-task-incomplete".to_owned())],
                arena.text("[ ] "),
            ),
            None => arena.nil(),
        };
        let content = task.append(
            arena.concat(
                self.blocks
                    .iter()
                    .map(|block| block.to_doc(config, context, arena)),
            ),
        );

        arena.tag("li", vec![], content)
    }
}

impl<'a> ToDoc<'a> for CodeBlock {
    fn to_doc(
        &self,
        _config: &'a Config<'a>,
        _context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        let code = highlighter::highlight(self);

        arena.tag("pre", vec![], arena.tag("code", vec![], arena.text(code)))
    }
}

impl<'a> ToDoc<'a> for Table {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        let first_row = table_row_to_doc(
            config,
            context,
            arena,
            self.rows.first().unwrap(),
            "th",
            &self.alignments,
        );
        let mut acc = arena.nil();
        for row in self.rows.iter().skip(1) {
            acc = acc.append(table_row_to_doc(
                config,
                context,
                arena,
                row,
                "td",
                &self.alignments,
            ));
        }

        let content = arena
            .tag("thead", vec![], first_row)
            .append(arena.tag("tbody", vec![], acc));

        arena.tag("table", vec![], content)
    }
}

fn table_row_to_doc<'a>(
    config: &'a Config<'a>,
    context: &'a Context,
    arena: &'a Arena<'a>,
    row: &TableRow,
    row_tag: &'static str,
    alignments: &[Alignment],
) -> DocBuilder<'a, Arena<'a>, ()> {
    let mut acc = arena.nil();
    for (i, cell) in row.iter().enumerate() {
        let alignment = match alignments.get(i) {
            None | Some(Alignment::Left | Alignment::None) => "left",
            Some(Alignment::Right) => "right",
            Some(Alignment::Center) => "center",
        };
        let alignment_class = format!("table-align-{alignment}");
        let attributes = vec![("class".to_owned(), alignment_class)];
        acc = acc.append(arena.tag(row_tag, attributes, cell.to_doc(config, context, arena)));
    }

    arena.tag("tr", vec![], acc)
}

impl<'a> ToDoc<'a> for FootnoteDefinition {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        let Some(index) = context.get_footnote_index(&self.label) else {
            return arena.nil();
        };

        arena.tag(
            "div",
            vec![
                ("class".to_owned(), "footnote-definition".to_owned()),
                ("id".to_owned(), format!("footnote-{index}")),
            ],
            arena
                .tag(
                    "span",
                    vec![("class".to_owned(), "footnote-definition-index".to_owned())],
                    arena.text(format!("{index}. ")),
                )
                .append(arena.tag(
                    "span",
                    vec![("class".to_owned(), "footnote-definition-content".to_owned())],
                    self.blocks.to_doc(config, context, arena),
                )),
        )
    }
}

impl<'a> ToDoc<'a> for Callout {
    fn to_doc(
        &self,
        config: &'a Config<'a>,
        context: &'a Context,
        arena: &'a Arena<'a>,
    ) -> DocBuilder<'a, Arena<'a>, ()> {
        arena.tag(
            "div",
            vec![
                ("class".to_owned(), "callout".to_owned()),
                ("data-callout".to_owned(), "info".to_owned()),
            ],
            arena
                .tag(
                    "div",
                    vec![("class".to_owned(), "callout-title".to_owned())],
                    arena
                        .tag(
                            "div",
                            vec![("class".to_owned(), "callout-title-icon".to_owned())],
                            arena.nil(),
                        )
                        .append(
                            arena.tag(
                                "div",
                                vec![("class".to_owned(), "callout-title-inner".to_owned())],
                                arena.text(
                                    self.title
                                        .clone()
                                        .unwrap_or_else(|| self.level.to_title_case()),
                                ),
                            ),
                        ),
                )
                .append(arena.tag(
                    "div",
                    vec![("class".to_owned(), "callout-content".to_owned())],
                    self.blocks.to_doc(config, context, arena),
                )),
        )
    }
}