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

use crate::{
    ast_nodes::{AstNode, AstNodes},
    format_args,
    formatter::{
        Buffer, Format, GroupId, JsFormatContext, JsFormatter,
        prelude::*,
        trivia::{DanglingIndentMode, FormatDanglingComments},
    },
    options::{FormatTrailingCommas, TrailingSeparator},
    utils::{
        call_expression::is_test_call_expression,
        typescript::{is_object_like_type, is_simple_type, should_hug_type},
    },
    write,
};

use super::FormatWrite;

impl<'a> FormatWrite<'a> for AstNode<'a, TSTypeParameter<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        if self.r#const() {
            write!(f, ["const", space()]);
        }
        if self.r#in() {
            write!(f, ["in", space()]);
        }
        if self.out() {
            write!(f, ["out", space()]);
        }
        write!(f, self.name());

        if let Some(constraint) = &self.constraint() {
            let group_id = f.group_id("constraint");

            write!(
                f,
                [
                    space(),
                    "extends",
                    group(&indent(&format_args!(
                        line_suffix_boundary(),
                        soft_line_break_or_space()
                    )))
                    .with_group_id(Some(group_id)),
                    indent_if_group_breaks(&constraint, group_id)
                ]
            );
        }
        if let Some(default) = &self.default() {
            let group_id = f.group_id("default");
            write!(
                f,
                [
                    space(),
                    "=",
                    group(&indent(&soft_line_break_or_space())).with_group_id(Some(group_id)),
                    line_suffix_boundary(),
                    indent_if_group_breaks(&default, group_id)
                ]
            );
        }
    }
}

impl<'a> Format<'a, JsFormatContext<'a>> for AstNode<'a, ArenaVec<'a, TSTypeParameter<'a>>> {
    fn fmt(&self, f: &mut JsFormatter<'_, 'a>) {
        // Type parameter lists of arrow function expressions have to include at least one comma
        // to avoid any ambiguity with JSX elements, and in `.mts`/`.cts` sources.
        // Thus, we have to add a trailing comma when there is a single type parameter.
        // The comma can be omitted in the case where the single parameter has a constraint,
        // i.i. an `extends` clause.
        let trailing_separator = if should_force_trailing_comma_for_arrow_function(self, f) {
            TrailingSeparator::Mandatory
        } else {
            FormatTrailingCommas::ES5.trailing_separator(f.options())
        };

        f.join_with(soft_line_break_or_space()).entries_with_trailing_separator(
            self.iter(),
            ",",
            trailing_separator,
        );
    }
}

/// Matches Prettier's `shouldForceTrailingComma` behavior for arrow functions.
///
/// <https://github.com/prettier/prettier/blob/070c89bba46235f4948560ed612a11e89ccd2da9/src/language-js/print/type-parameters.js#L33-L42>
fn should_force_trailing_comma_for_arrow_function(
    params: &AstNode<'_, ArenaVec<'_, TSTypeParameter<'_>>>,
    f: &JsFormatter<'_, '_>,
) -> bool {
    if params.len() != 1 {
        return false;
    }

    if !matches!(params.grand_parent(), AstNodes::ArrowFunctionExpression(_)) {
        return false;
    }

    // Ignore type parameters with a constraint (extends clause).
    // A default type alone does not disambiguate from JSX, so the comma is still required.
    if params.first().is_some_and(|t| t.constraint().is_some()) {
        return false;
    }

    let source_type = f.context().source_type();
    let is_ts_extension = matches!(source_type.extension(), Some(FileExtension::Ts));

    // Force trailing comma for non-.ts sources (e.g. .tsx, .mts, .cts) or when extension is unknown.
    !is_ts_extension
}

#[derive(Default)]
pub struct FormatTSTypeParametersOptions {
    pub group_id: Option<GroupId>,
    pub is_type_or_interface_decl: bool,
}

pub struct FormatTSTypeParameters<'a, 'b> {
    decl: &'b AstNode<'a, TSTypeParameterDeclaration<'a>>,
    options: FormatTSTypeParametersOptions,
}

impl<'a, 'b> FormatTSTypeParameters<'a, 'b> {
    pub fn new(
        decl: &'b AstNode<'a, TSTypeParameterDeclaration<'a>>,
        options: FormatTSTypeParametersOptions,
    ) -> Self {
        Self { decl, options }
    }
}

impl<'a> Format<'a, JsFormatContext<'a>> for FormatTSTypeParameters<'a, '_> {
    fn fmt(&self, f: &mut JsFormatter<'_, 'a>) {
        let params = self.decl.params();
        if params.is_empty() && self.options.is_type_or_interface_decl {
            write!(f, "<>");
        } else {
            write!(
                f,
                [group(&format_args!("<", format_with(|f| {
                    if matches!(self.decl.grand_parent(), AstNodes::CallExpression(call) if is_test_call_expression(call))
                    {
                        f.join_nodes_with_space().entries_with_trailing_separator(params, ",", TrailingSeparator::Omit);
                    } else {
                        soft_block_indent(&params).fmt(f);
                    }

                    format_dangling_comments(self.decl.span).with_soft_block_indent().fmt(f);
                }), ">"))
                    .with_group_id(self.options.group_id)]
            );
        }
    }
}

impl<'a> FormatWrite<'a> for AstNode<'a, TSTypeParameterInstantiation<'a>> {
    fn write(&self, f: &mut JsFormatter<'_, 'a>) {
        let params = self.params();

        if params.is_empty() {
            // This shouldn't happen in valid TypeScript code, but handle it gracefully
            let comments = f.context().comments().comments_before(self.span.end);
            let indent = if comments.iter().any(|c| c.is_line()) {
                DanglingIndentMode::Soft
            } else {
                DanglingIndentMode::None
            };
            return write!(f, ["<", FormatDanglingComments::Comments { comments, indent }, ">"]);
        }

        // Check if this is in the context of an arrow function variable
        let is_arrow_function_vars = is_arrow_function_variable_type_argument(self);

        // Check if the first (and only) argument can be hugged
        let first_arg_can_be_hugged = if params.len() == 1 {
            if let Some(first_type) = params.first() {
                matches!(first_type.as_ref(), TSType::TSNullKeyword(_))
                    || should_hug_single_type(first_type.as_ref(), f)
            } else {
                false
            }
        } else {
            false
        };

        let format_params = format_with(|f| {
            f.join_with(&soft_line_break_or_space()).entries_with_trailing_separator(
                params,
                ",",
                TrailingSeparator::Disallowed,
            );
        });

        let should_inline = !is_arrow_function_vars && first_arg_can_be_hugged;

        if should_inline {
            write!(f, ["<", format_params, ">"]);
        } else {
            write!(f, [group(&format_args!("<", soft_block_indent(&format_params), ">"))]);
        }
    }
}

/// Check if a single type should be "hugged" (kept inline)
fn should_hug_single_type(ty: &TSType, f: &JsFormatter<'_, '_>) -> bool {
    // Simple types and object-like types can be hugged
    if is_simple_type(ty) || is_object_like_type(ty) {
        return true;
    }

    // Check for union types with mostly void types and one object type
    // (e.g., `SomeType<ObjectType | null | undefined>`)
    matches!(ty, TSType::TSUnionType(union) if should_hug_type(union, f))
}

/// Check if this type parameter instantiation is in an arrow function variable context
///
/// This detects patterns like:
/// ```typescript
/// const foo: SomeThing<{ [P in "x" | "y"]: number }> = () => {};
/// ```
fn is_arrow_function_variable_type_argument<'a>(
    node: &AstNode<'a, TSTypeParameterInstantiation<'a>>,
) -> bool {
    let Some(first) = node.params().first() else { unreachable!() };

    // Skip check for single object-like types
    if node.params().len() == 1 && is_object_like_type(first.as_ref()) {
        return false;
    }

    // `node.parent` is `TSTypeReference`
    matches!(
        &node.grand_parent(),
        AstNodes::TSTypeAnnotation(type_annotation)
            if matches!(
                &type_annotation.parent(),
                AstNodes::VariableDeclarator(var_decl)
                    if matches!(&var_decl.init, Some(Expression::ArrowFunctionExpression(_)))
            )
    )
}
