use oxc_ast::ast::*;
use oxc_span::GetSpan;

use crate::{
    JsLabels,
    ast_nodes::{AstNode, AstNodes},
    format_args,
    formatter::{Buffer, Format, JsFormatter, prelude::*, trivia::FormatLeadingComments},
    utils::{
        expression::as_call_expression_without_chain_wrappers,
        member_chain::chain_member::FormatComputedMemberExpressionWithoutObject,
    },
    write,
};

use super::FormatWrite;

impl<'a> FormatWrite<'a> for AstNode<'a, ComputedMemberExpression<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        write!(f, self.object());
        FormatComputedMemberExpressionWithoutObject(self).fmt(f);
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, StaticMemberExpression<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let is_member_chain = {
            let mut recording = f.start_recording();
            write!(recording, [self.object()]);
            recording.stop().has_label(LabelId::of(JsLabels::MemberChain))
        };

        match layout(self, is_member_chain, f) {
            StaticMemberLayout::NoBreak => {
                let format_no_break =
                    format_with(|f| write!(f, [operator_token(self.optional()), self.property()]));

                if is_member_chain {
                    write!(f, [labelled(LabelId::of(JsLabels::MemberChain), &format_no_break)]);
                } else {
                    write!(f, [format_no_break]);
                }
            }
            StaticMemberLayout::BreakAfterObject => {
                write!(
                    f,
                    [group(&indent(&format_args!(
                        soft_line_break(),
                        &format_with(|f| {
                            if f.comments().has_leading_own_line_comment(self.property.span.start) {
                                let comments = f
                                    .context()
                                    .comments()
                                    .comments_before(self.property.span.start);
                                write!(
                                    f,
                                    [FormatLeadingComments::Comments(comments), soft_line_break()]
                                );
                            }
                        }),
                        operator_token(self.optional()),
                        self.property(),
                    )))]
                );
            }
        }
    }
}

#[derive(Debug, Copy, Clone)]
enum StaticMemberLayout {
    /// Forces that there's no line break between the object, operator, and member
    NoBreak,

    /// Breaks the static member expression after the object if the whole expression doesn't fit on a single line
    BreakAfterObject,
}

fn operator_token(optional: bool) -> &'static str {
    if optional { "?." } else { "." }
}

fn layout<'a>(
    node: &AstNode<'a, StaticMemberExpression<'a>>,
    is_member_chain: bool,
    f: &JsFormatter<'_, 'a>,
) -> StaticMemberLayout {
    if f.comments().has_leading_own_line_comment(node.property.span.start) {
        return StaticMemberLayout::BreakAfterObject;
    }

    // `a.b.c!` and `a.b?.c`
    // `TSNonNullExpression` is a wrapper node for `!`, and `ChainExpression` is a wrapper node for `?.`,
    // so we need to skip them to find the real parent node.
    let mut parent = node.parent();
    while matches!(parent, AstNodes::TSNonNullExpression(_) | AstNodes::ChainExpression(_)) {
        parent = parent.parent();
    }
    let object = &node.object;

    let is_nested = match parent {
        AstNodes::AssignmentExpression(_) | AstNodes::VariableDeclarator(_) => {
            let no_break = as_call_expression_without_chain_wrappers(object)
                .is_some_and(|call| !call.arguments.is_empty());

            if no_break || is_member_chain {
                return StaticMemberLayout::NoBreak;
            }

            false
        }
        AstNodes::StaticMemberExpression(_) | AstNodes::ComputedMemberExpression(_) => true,
        _ => false,
    };

    if !is_nested && matches!(object, Expression::Identifier(_)) {
        return StaticMemberLayout::NoBreak;
    }

    let mut first_non_static_member_ancestor = parent;
    while matches!(
        first_non_static_member_ancestor,
        AstNodes::StaticMemberExpression(_)
            | AstNodes::ComputedMemberExpression(_)
            | AstNodes::ChainExpression(_)
            | AstNodes::TSNonNullExpression(_)
    ) {
        first_non_static_member_ancestor = first_non_static_member_ancestor.parent();
    }

    match first_non_static_member_ancestor {
        AstNodes::NewExpression(expr) if expr.callee.span().contains_inclusive(node.span()) => {
            StaticMemberLayout::NoBreak
        }
        AstNodes::AssignmentExpression(assignment) => {
            if matches!(assignment.left, AssignmentTarget::AssignmentTargetIdentifier(_)) {
                StaticMemberLayout::BreakAfterObject
            } else {
                StaticMemberLayout::NoBreak
            }
        }
        _ => StaticMemberLayout::BreakAfterObject,
    }
}

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