use oxc_allocator::{ArenaBox, ArenaVec, CloneIn, GetAllocator, ReplaceWith};
use oxc_ast::{ast::*, builder::NONE};
use oxc_span::{GetSpan, SPAN};
use oxc_str::Str;

use crate::{IsolatedDeclarations, diagnostics::default_export_inferred};

impl<'a> IsolatedDeclarations<'a> {
    pub(crate) fn transform_export_named_declaration(
        &mut self,
        prev_decl: &ExportNamedDeclaration<'a>,
    ) -> Option<ArenaBox<'a, ExportNamedDeclaration<'a>>> {
        let decl = self.transform_declaration(prev_decl.declaration.as_ref()?, false)?;

        Some(ExportNamedDeclaration::boxed(
            prev_decl.span,
            Some(decl),
            ArenaVec::new_in(self),
            None,
            ImportOrExportKind::Value,
            NONE,
            self,
        ))
    }

    pub(crate) fn create_unique_name(&self, name: &str) -> Str<'a> {
        let mut binding = Str::from_str_in(name, self);
        let mut i = 1;
        while self.scope.has_reference(&binding) {
            binding = Str::from_str_in(format!("{name}_{i}").as_str(), self);
            i += 1;
        }
        binding
    }

    pub(crate) fn transform_export_default_declaration(
        &self,
        decl: &ExportDefaultDeclaration<'a>,
    ) -> Option<(Option<Statement<'a>>, Statement<'a>)> {
        let declaration = match &decl.declaration {
            ExportDefaultDeclarationKind::FunctionDeclaration(decl) => Some((
                None,
                ExportDefaultDeclarationKind::FunctionDeclaration(
                    self.transform_function(decl, Some(false)),
                ),
            )),
            ExportDefaultDeclarationKind::ClassDeclaration(decl) => Some((
                None,
                ExportDefaultDeclarationKind::ClassDeclaration(
                    self.transform_class(decl, Some(false)),
                ),
            )),
            ExportDefaultDeclarationKind::TSInterfaceDeclaration(_) => {
                Some((None, decl.declaration.clone_in(self.allocator())))
            }
            declaration @ match_expression!(ExportDefaultDeclarationKind) => self
                .transform_export_expression(decl.span, declaration.to_expression())
                .map(|(var_decl, expr)| (var_decl, ExportDefaultDeclarationKind::from(expr))),
        };

        declaration.map(|(var_decl, declaration)| {
            // When `var_decl` is Some, the comments are moved to the variable declaration, otherwise
            // keep the comments on the export default declaration to avoid losing them.
            // ```ts
            // // comment
            // export default function(): void {}
            //
            // // comment
            // export default 1;
            // ```
            //
            // to
            //
            // ```ts
            // // comment
            // export default function(): void;
            //
            // // comment
            // const _default = 1;
            // export default _default;
            // ```

            let span = if var_decl.is_some() { SPAN } else { decl.span };
            let declaration = Statement::new_export_default_declaration(span, declaration, self);
            (var_decl, declaration)
        })
    }

    fn transform_export_expression(
        &self,
        decl_span: Span,
        expr: &Expression<'a>,
    ) -> Option<(Option<Statement<'a>>, Expression<'a>)> {
        if matches!(expr, Expression::Identifier(_)) {
            None
        } else {
            // declare const _default: Type
            let kind = VariableDeclarationKind::Const;
            let name = self.create_unique_name("_default");
            let id = BindingPattern::new_binding_identifier(SPAN, name, self);
            let type_annotation = self
                .infer_type_from_expression(expr)
                .map(|ts_type| TSTypeAnnotation::new(SPAN, ts_type, self));

            if type_annotation.is_none() {
                self.error(default_export_inferred(expr.span()));
            }

            let declarations = ArenaVec::from_value_in(
                VariableDeclarator::new(SPAN, kind, id, type_annotation, None, false, self),
                self,
            );

            let variable_statement = Statement::new_variable_declaration(
                decl_span,
                kind,
                declarations,
                self.is_declare(),
                self,
            );
            Some((Some(variable_statement), Expression::new_identifier(SPAN, name, self)))
        }
    }

    pub(crate) fn transform_ts_export_assignment(
        &self,
        decl: &TSExportAssignment<'a>,
    ) -> Option<(Option<Statement<'a>>, Statement<'a>)> {
        self.transform_export_expression(decl.span, &decl.expression).map(|(var_decl, expr)| {
            (var_decl, Statement::new_ts_export_assignment(SPAN, expr, self))
        })
    }

    pub(crate) fn transform_import_declaration(
        &self,
        decl: &ImportDeclaration<'a>,
    ) -> Option<ArenaBox<'a, ImportDeclaration<'a>>> {
        let specifiers = decl.specifiers.as_ref()?;

        let mut new_specifiers = ArenaVec::with_capacity_in(specifiers.len(), self);
        specifiers.iter().for_each(|specifier| {
            let is_referenced = match specifier {
                ImportDeclarationSpecifier::ImportSpecifier(specifier) => {
                    self.scope.has_reference(&specifier.local.name)
                }
                ImportDeclarationSpecifier::ImportDefaultSpecifier(specifier) => {
                    self.scope.has_reference(&specifier.local.name)
                }
                ImportDeclarationSpecifier::ImportNamespaceSpecifier(_) => {
                    self.scope.has_reference(&specifier.name())
                }
            };
            if is_referenced {
                new_specifiers.push(specifier.clone_in(self.allocator()));
            }
        });

        if new_specifiers.is_empty() {
            // We don't need to print this import statement
            None
        } else {
            Some(ImportDeclaration::boxed(
                decl.span,
                Some(new_specifiers),
                decl.source.clone(),
                None,
                decl.with_clause.clone_in(self.allocator()),
                decl.import_kind,
                self,
            ))
        }
    }

    /// Strip export keyword from ExportNamedDeclaration
    ///
    /// ```ts
    /// export const a = 1;
    /// export function b() {}
    /// ```
    /// to
    /// ```ts
    /// const a = 1;
    /// function b() {}
    /// ```
    pub(crate) fn strip_export_keyword(stmts: &mut ArenaVec<'a, Statement<'a>>) {
        stmts.iter_mut().for_each(|stmt| {
            if let Statement::ExportNamedDeclaration(decl) = stmt
                && decl.declaration.is_some()
            {
                stmt.replace_with(|stmt| {
                    let Statement::ExportNamedDeclaration(decl) = stmt else { unreachable!() };
                    Statement::from(decl.unbox().declaration.unwrap())
                });
            }
        });
    }
}
