use oxc_allocator::ArenaVec;
use oxc_ast::ast::*;

pub mod child_list;
pub mod element;
pub mod opening_element;

use child_list::{FormatChildrenResult, FormatJsxChildList, JsxChildListLayout};
use element::AnyJsxTagWithChildren;
use opening_element::FormatOpeningElement;
use oxc_span::GetSpan;

use crate::{
    AttributePosition, Format,
    ast_nodes::{AstNode, AstNodes},
    format_args,
    formatter::{
        TailwindContextEntry,
        prelude::*,
        trivia::{DanglingIndentMode, FormatDanglingComments, FormatTrailingComments},
    },
    utils::{
        expression::as_call_expression_without_chain_wrappers,
        tailwindcss::is_tailwind_jsx_attribute,
    },
    write,
};

use super::FormatWrite;

impl<'a> FormatWrite<'a> for AstNode<'a, JSXElement<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        AnyJsxTagWithChildren::Element(self).fmt(f);
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXOpeningElement<'a>> {
    fn write(&self, _f: &mut JsFormatter<'_, 'a>) {
        unreachable!("`AnyJsxTagWithChildren` will print it.")
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXClosingElement<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let name = self.name();
        let mut name_has_leading_comment = false;
        let mut name_has_own_line_leading_comment = false;
        for leading_comment in f.comments().comments_before(name.span().start) {
            name_has_leading_comment = true;
            name_has_own_line_leading_comment =
                name_has_own_line_leading_comment || leading_comment.is_line();
        }

        let format_name = format_with(|f| {
            if name_has_own_line_leading_comment {
                write!(f, [hard_line_break()]);
            } else if name_has_leading_comment {
                write!(f, [space()]);
            }
            if name_has_own_line_leading_comment {
                write!(f, [block_indent(&name), hard_line_break()]);
            } else {
                write!(f, [name]);
            }
        });

        write!(f, ["</", &format_name, ">",]);
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXFragment<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        AnyJsxTagWithChildren::Fragment(self).fmt(f);
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXOpeningFragment> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let comments = f.context().comments().comments_before(self.span.end);

        if comments.is_empty() {
            write!(f, "<>");
            return;
        }

        let has_own_line_comment = comments.iter().any(|c| c.is_line());

        let format_comments = format_with(|f| {
            if has_own_line_comment {
                write!(f, [hard_line_break()]);
            }

            write!(
                f,
                [FormatDanglingComments::Comments { comments, indent: DanglingIndentMode::None }]
            );
        });

        write!(
            f,
            ["<", indent(&format_comments), has_own_line_comment.then_some(hard_line_break()), ">"]
        );
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXClosingFragment> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let comments = f.context().comments().comments_before(self.span.end);

        if comments.is_empty() {
            write!(f, "</>");
            return;
        }

        let has_own_line_comment = comments.iter().any(|c| c.is_line());

        let format_comments = format_with(|f| {
            if has_own_line_comment {
                write!(f, [hard_line_break()]);
            } else if !comments.is_empty() {
                write!(f, [space()]);
            }

            write!(
                f,
                [FormatDanglingComments::Comments { comments, indent: DanglingIndentMode::None }]
            );
        });

        write!(
            f,
            [
                "</",
                indent(&format_comments),
                has_own_line_comment.then_some(hard_line_break()),
                ">"
            ]
        );
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXNamespacedName<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        write!(f, [self.namespace(), ":", self.name()]);
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXMemberExpression<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        write!(f, [self.object(), ".", self.property()]);
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXExpressionContainer<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let has_comment = |f: &mut JsFormatter<'_, '_>| {
            let expression_span = self.expression.span();
            f.comments().has_comment_before(expression_span.start)
                || f.comments().has_comment_in_range(expression_span.end, self.span.end)
        };

        // Expression child
        if matches!(self.parent(), AstNodes::JSXElement(_) | AstNodes::JSXFragment(_)) {
            if let JSXExpression::EmptyExpression(_) = self.expression {
                let comments = f.context().comments().comments_before(self.span.end);
                let has_line_comment = comments.iter().any(|c| c.is_line());

                write!(f, ["{"]);

                if has_line_comment {
                    write!(
                        f,
                        [
                            FormatDanglingComments::Comments {
                                comments,
                                indent: DanglingIndentMode::Block
                            },
                            format_dangling_comments(self.span).with_block_indent(),
                            hard_line_break()
                        ]
                    );
                } else {
                    write!(
                        f,
                        [FormatDanglingComments::Comments {
                            comments,
                            indent: DanglingIndentMode::None
                        },]
                    );
                }

                write!(f, ["}"]);
            } else {
                let is_conditional_or_binary = matches!(
                    self.expression,
                    JSXExpression::ConditionalExpression(_)
                        | JSXExpression::LogicalExpression(_)
                        | JSXExpression::BinaryExpression(_)
                );

                let should_inline = !has_comment(f)
                    && (is_conditional_or_binary || should_inline_jsx_expression(self));

                let format_expression = format_with(|f| {
                    if should_inline {
                        write!(f, self.expression());
                    } else {
                        write!(
                            f,
                            soft_block_indent(&format_with(|f| {
                                write!(f, [self.expression()]);
                                let comments =
                                    f.context().comments().comments_before(self.span.end);
                                write!(f, [FormatTrailingComments::Comments(comments)]);
                            }))
                        );
                    }
                });

                write!(
                    f,
                    [group(&format_args!("{", format_expression, line_suffix_boundary(), "}"))]
                );
            }
        } else {
            // JSXAttributeValue
            let should_inline = !has_comment(f) && should_inline_jsx_expression(self);

            if should_inline {
                write!(f, ["{", self.expression(), line_suffix_boundary(), "}"]);
            } else {
                write!(
                    f,
                    [group(&format_args!(
                        "{",
                        soft_block_indent(&format_with(|f| {
                            write!(f, [self.expression()]);
                            let comments = f.context().comments().comments_before(self.span.end);
                            write!(f, [FormatTrailingComments::Comments(comments)]);
                        })),
                        line_suffix_boundary(),
                        "}"
                    ))]
                );
            }
        }
    }
}

/// Tests if an expression inside of a [`JSXExpressionContainer`] should be inlined.
/// Good:
/// ```jsx
///  <ColorPickerPage
///     colors={[
///        "blue",
///        "brown",
///        "green",
///        "orange",
///        "purple",
///     ]} />
/// ```
///
/// Bad:
/// ```jsx
///  <ColorPickerPage
///     colors={
///       [
///         "blue",
///          "brown",
///         "green",
///         "orange",
///         "purple",
///       ]
///     } />
/// ```
pub fn should_inline_jsx_expression(container: &JSXExpressionContainer<'_>) -> bool {
    match &container.expression {
        JSXExpression::ArrayExpression(_)
        | JSXExpression::ObjectExpression(_)
        | JSXExpression::ArrowFunctionExpression(_)
        | JSXExpression::CallExpression(_)
        | JSXExpression::ImportExpression(_)
        | JSXExpression::MetaProperty(_)
        | JSXExpression::FunctionExpression(_)
        | JSXExpression::TemplateLiteral(_)
        | JSXExpression::TaggedTemplateExpression(_) => true,
        // A call wrapped in `?.` chains and/or `!` assertions still inlines.
        JSXExpression::ChainExpression(_) | JSXExpression::TSNonNullExpression(_) => {
            as_call_expression_without_chain_wrappers(container.expression.to_expression())
                .is_some()
        }
        JSXExpression::AwaitExpression(await_expression) => {
            matches!(
                await_expression.argument,
                Expression::ArrayExpression(_)
                    | Expression::ObjectExpression(_)
                    | Expression::ArrowFunctionExpression(_)
                    | Expression::CallExpression(_)
                    | Expression::ImportExpression(_)
                    | Expression::MetaProperty(_)
                    | Expression::FunctionExpression(_)
                    | Expression::TemplateLiteral(_)
                    | Expression::TaggedTemplateExpression(_)
                    | Expression::JSXElement(_)
                    | Expression::JSXFragment(_)
            )
        }
        _ => false,
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXEmptyExpression> {
    fn write(&self, _f: &mut JsFormatter<'_, 'a>) {}
}

impl<'a> Format<'a, JsFormatContext<'a>> for AstNode<'a, ArenaVec<'a, JSXAttributeItem<'a>>> {
    fn fmt(&self, f: &mut JsFormatter<'_, 'a>) {
        if f.options().attribute_position == AttributePosition::Multiline {
            f.join_nodes_with_hardline().entries(self.iter());
        } else {
            f.join_nodes_with_soft_line().entries(self.iter());
        }
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXAttribute<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        write!(f, self.name());

        if let Some(value) = &self.value() {
            // Check if this is a Tailwind attribute and push context
            // Extract context entry before mutating f
            let tailwind_ctx_to_push = f
                .options()
                .sort_tailwindcss
                .as_ref()
                .filter(|opts| is_tailwind_jsx_attribute(&self.name, opts))
                .map(|opts| TailwindContextEntry::new(opts.preserve_whitespace));

            if let Some(ctx) = tailwind_ctx_to_push {
                f.context_mut().push_tailwind_context(ctx);
            }

            write!(f, ["=", value]);

            if tailwind_ctx_to_push.is_some() {
                f.context_mut().pop_tailwind_context();
            }
        }
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXSpreadAttribute<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let comments = f.context().comments();
        let has_comment = comments.has_comment_before(self.argument.span().start)
            || comments.has_comment_in_range(self.argument.span().end, self.span.end);
        let format_inner = format_with(|f| {
            write!(f, [format_leading_comments(self.argument.span()), "..."]);
            self.argument().fmt(f);
        });

        write!(f, ["{"]);

        if has_comment {
            write!(f, [soft_block_indent(&format_inner)]);
        } else {
            write!(f, [format_inner]);
        }

        write!(f, "}");
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXIdentifier<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        write!(f, text_without_whitespace(self.name().as_str()));
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXSpreadChild<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let comments = f.context().comments();
        let has_comment = comments.has_comment_before(self.expression.span().start)
            || comments.has_comment_in_range(self.expression.span().end, self.span.end);
        let format_inner = format_with(|f| {
            write!(f, [format_leading_comments(self.expression.span()), "..."]);
            self.expression().fmt(f);
        });

        write!(f, "{");

        if has_comment {
            write!(f, [soft_block_indent(&format_inner)]);
        } else {
            write!(f, [format_inner]);
        }

        write!(f, "}");
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, JSXText<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        write!(f, text(self.value().as_str()));
    }
}
