use indexmap::map::Entry;
use oxc::{
  allocator::{Allocator, TakeIn},
  ast::ast::{self, Expression},
  semantic::{SemanticBuilder, Stats},
  span::SPAN,
};
use oxc_str::CompactStr;
use rolldown_common::{
  EcmaModuleAstUsage, ExportsKind, GetLocalDbMut, LocalExport, Module, ModuleIdx, ModuleType,
  NormalModule, StmtInfo, StmtInfoIdx, StmtInfos, SymbolOrMemberExprRef, SymbolRef,
  SymbolRefDbForModule, TaggedSymbolRef, WrapKind,
};
use rolldown_ecmascript::EcmaAst;
use rolldown_ecmascript_utils::AstFactory;
use rolldown_utils::ecmascript::legitimize_json_local_binding_name;
#[cfg(not(target_family = "wasm"))]
use rolldown_utils::rayon::IndexedParallelIterator;
use rolldown_utils::{
  indexmap::FxIndexMap,
  rayon::{IntoParallelRefMutIterator, ParallelIterator},
};

use super::LinkStage;

/// Index of the first statement after the namespace statement (index 0).
const FIRST_TOP_LEVEL_STMT_IDX: StmtInfoIdx = StmtInfoIdx::from_raw_unchecked(1);

struct LazyModuleInfo {
  idx: ModuleIdx,
  exports_kind: ExportsKind,
  is_json: bool,
}

impl LinkStage<'_> {
  #[tracing::instrument(level = "debug", skip_all)]
  pub(super) fn generate_lazy_export(&mut self) {
    let lazy_modules = append_only_vec::AppendOnlyVec::new();
    self
      .module_table
      .modules
      .par_iter_mut()
      .zip(self.stmt_infos.par_iter_mut())
      .filter_map(|(m, stmt_infos)| m.as_normal_mut().map(|m| (m, stmt_infos)))
      .filter(|(module, _)| module.meta.has_lazy_export())
      .for_each(|(module, stmt_infos)| {
        let default_symbol_ref = module.default_export_ref;
        let is_json = matches!(module.module_type, ModuleType::Json);
        if !is_json || module.exports_kind == ExportsKind::CommonJs {
          update_module_default_export_info(module, stmt_infos, default_symbol_ref);
        }
        lazy_modules.push(LazyModuleInfo {
          idx: module.idx,
          exports_kind: module.exports_kind,
          is_json,
        });

        // generate `module.exports = expr`
        if module.exports_kind == ExportsKind::CommonJs {
          // since the wrap arguments are generate on demand, we need to insert the module ref usage here.
          stmt_infos.infos[FIRST_TOP_LEVEL_STMT_IDX].side_effect = true.into();
          module.ecma_view.ast_usage.insert(EcmaModuleAstUsage::ModuleRef);
        }
      });

    for LazyModuleInfo { idx: module_idx, exports_kind, is_json } in lazy_modules {
      let Some(ecma_ast) = &mut self.ast_table[module_idx] else { unreachable!() };
      if matches!(exports_kind, ExportsKind::CommonJs) {
        replace_first_expr_stmt(ecma_ast, LazyExportWrap::CjsExport);
        continue;
      }
      // ExportsKind == Esm && ModuleType == Json
      if is_json {
        if json_object_expr_to_esm(self, module_idx) {
          continue;
        }
        // if json is not a ObjectExpression, we will fallback to normal esm lazy export transform
        let stmt_infos = &mut self.stmt_infos[module_idx];
        let module = &mut self.module_table[module_idx];
        let module = module.as_normal_mut().unwrap();
        let default_export_ref = module.default_export_ref;
        update_module_default_export_info(module, stmt_infos, default_export_ref);
      }

      // shadowing the previous mutable ref, to avoid reference mutable ref twice at the same time.
      let Some(ecma_ast) = &mut self.ast_table[module_idx] else { unreachable!() };
      replace_first_expr_stmt(ecma_ast, LazyExportWrap::EsmDefault);

      // Ensure exports_kind is set to Esm for all modules that generate ESM export syntax.
      // This is needed for proper CJS export rendering in preserveModules mode.
      let module = &mut self.module_table[module_idx];
      let module = module.as_normal_mut().unwrap();
      module.exports_kind = ExportsKind::Esm;
    }
  }
}

#[derive(Clone, Copy)]
enum LazyExportWrap {
  CjsExport,
  EsmDefault,
}

/// Takes the expression of the first statement (which must be an `ExpressionStatement`)
/// and replaces the statement with either `module.exports = expr` or `export default expr`.
fn replace_first_expr_stmt(ecma_ast: &mut EcmaAst, kind: LazyExportWrap) {
  ecma_ast.program.with_mut(|fields| {
    let ast_factory = AstFactory::new(fields.allocator);
    let Some(stmt) = fields.program.body.first_mut() else { unreachable!() };
    let expr = match stmt {
      ast::Statement::ExpressionStatement(stmt) => stmt.expression.take_in(ast_factory.allocator),
      _ => unreachable!(),
    };
    *stmt = match kind {
      LazyExportWrap::CjsExport => ast_factory.make_module_exports_stmt(expr),
      LazyExportWrap::EsmDefault => ast_factory.make_export_default_stmt(expr),
    };
  });
}

/// Takes `expr` (leaving a dummy in its place) and returns the owned inner
/// expression with any wrapping `(...)` parentheses removed.
fn take_without_parentheses<'ast>(
  expr: &mut Expression<'ast>,
  allocator: &'ast Allocator,
) -> Expression<'ast> {
  let mut inner_expr = expr.take_in(allocator);
  while let Expression::ParenthesizedExpression(mut paren_expr) = inner_expr {
    inner_expr = paren_expr.expression.take_in(allocator);
  }
  inner_expr
}

fn update_module_default_export_info(
  module: &mut NormalModule,
  stmt_infos: &mut StmtInfos,
  default_symbol_ref: SymbolRef,
) {
  module.named_exports.insert(
    "default".into(),
    LocalExport { span: SPAN, referenced: default_symbol_ref, came_from_commonjs: false },
  );
  stmt_infos
    .declare_symbol_for_stmt(FIRST_TOP_LEVEL_STMT_IDX, TaggedSymbolRef::Normal(default_symbol_ref));
}

/// return true if the json is a ObjectExpression
fn json_object_expr_to_esm(link_staged: &mut LinkStage, module_idx: ModuleIdx) -> bool {
  let module = &mut link_staged.module_table[module_idx];
  let Module::Normal(module) = module else {
    return false;
  };

  let ecma_ast = link_staged.ast_table[module_idx].as_mut().unwrap();
  // (local, (exported, legal_ident))
  let mut declaration_binding_names: FxIndexMap<CompactStr, (CompactStr, bool)> =
    FxIndexMap::default();
  let transformed = ecma_ast.program.with_mut(|fields| {
    let mut index_map = FxIndexMap::default();
    let ast_factory = AstFactory::new(fields.allocator);
    let program = fields.program;
    let Some(stmts) = program.body.first_mut() else { unreachable!() };
    let expr = match stmts {
      ast::Statement::ExpressionStatement(stmt) => &mut stmt.expression,
      _ => {
        unreachable!()
      }
    };
    if !matches!(expr.without_parentheses(), Expression::ObjectExpression(_)) {
      return false;
    }
    let Expression::ObjectExpression(mut obj_expr) =
      take_without_parentheses(expr, ast_factory.allocator)
    else {
      unreachable!();
    };
    // clean program body, since we already take it and left a dummy expr
    program.body.clear();

    // convert {"a": "b", "c": "d"} to
    // {"a": b, "c": d}
    // and collect related info
    for property in &mut obj_expr.properties {
      match property {
        ast::ObjectPropertyKind::ObjectProperty(property) => {
          let key = property.key.static_name().expect("should be static name");
          if key.is_empty() {
            continue;
          }
          let legitimized_ident =
            legitimize_json_local_binding_name(&key, &declaration_binding_names);

          let is_legal_ident = legitimized_ident.as_str() == key;

          declaration_binding_names
            .insert(legitimized_ident.clone(), (CompactStr::new(&key), is_legal_ident));

          let value = std::mem::replace(
            &mut property.value,
            ast_factory.make_id_ref_expr(SPAN, legitimized_ident.as_str()),
          );
          // TODO(shulaoda): Waiting for oxc transform to support the ES feature `ShorthandProperties`.
          if key == "__proto__" {
            property.computed = true;
          } else if is_legal_ident {
            property.shorthand = is_legal_ident;
            property.key = ast::PropertyKey::StaticIdentifier(
              ast_factory.alloc_identifier_name(SPAN, ast_factory.str(legitimized_ident.as_ref())),
            );
          }
          match index_map.entry(legitimized_ident) {
            Entry::Occupied(mut occ) => {
              *occ.get_mut() = value;
            }
            Entry::Vacant(vac) => {
              vac.insert(value);
            }
          }
        }
        ast::ObjectPropertyKind::SpreadProperty(_) => unreachable!(),
      }
    }
    // recreate Json Module
    let stmts = index_map
      .into_iter()
      // declaration
      .map(|(local, v)| ast_factory.make_var_decl(local.as_str(), v))
      // export default json module
      .chain(std::iter::once(
        ast_factory.make_export_default_stmt(Expression::ObjectExpression(obj_expr)),
      ))
      // export all declaration
      .chain(std::iter::once(
        ast_factory.make_export_named_stmt(None, declaration_binding_names.iter()),
      ));
    program.body.extend(stmts);
    true
  });

  if !transformed {
    return false;
  }
  let original_symbol_ref_db = std::mem::take(link_staged.symbols.local_db_mut(module_idx));
  // recreate semantic data
  #[expect(clippy::cast_possible_truncation)]
  let scoping = ecma_ast.make_symbol_table_and_scope_tree_with_semantic_builder(
    SemanticBuilder::new().with_stats(Stats {
      nodes: declaration_binding_names.len().next_power_of_two() as u32,
      scopes: 1,
      symbols: declaration_binding_names.len() as u32,
      references: declaration_binding_names.len() as u32 * 2u32,
    }),
  );

  // update semantic data of module
  let root_scope_id = scoping.root_scope_id();
  let mut symbol_ref_db = SymbolRefDbForModule::new(scoping, module_idx, root_scope_id);
  // Re-create facade symbols in the new scoping. The JSON module was re-parsed above,
  // producing a new Scoping with fresh symbol IDs, so old facade IDs are invalid.
  // We allocate new IDs by name and update the module's references.
  let mut recreate_facade = |old_ref: SymbolRef| -> SymbolRef {
    symbol_ref_db.create_facade_root_symbol_ref(original_symbol_ref_db.symbol_name(old_ref.symbol))
  };
  module.namespace_object_ref = recreate_facade(module.namespace_object_ref);
  module.default_export_ref = recreate_facade(module.default_export_ref);
  if let Some(hot_ref) = module.hmr_hot_ref {
    module.hmr_hot_ref = Some(recreate_facade(hot_ref));
  }
  let namespace_object_ref = module.namespace_object_ref;
  let default_export_ref = module.default_export_ref;

  // update module stmts info
  // clear stmt info, since we need to split `ObjectExpression` into multiple decl, the original stmt info is invalid.
  // preserve the first one, which is `NamespaceRef`
  let stmt_infos = &mut link_staged.stmt_infos[module_idx];
  let stmt_info = stmt_infos.drain(FIRST_TOP_LEVEL_STMT_IDX..);
  let mut all_declared_symbols =
    stmt_info.flat_map(|info| info.referenced_symbols).collect::<Vec<_>>();
  for (local, (exported, _)) in &declaration_binding_names {
    let symbol_id =
      symbol_ref_db.scoping().get_root_binding(local.as_str().into()).expect("should have binding");
    let symbol_ref: SymbolRef = (module_idx, symbol_id).into();
    all_declared_symbols.push(SymbolOrMemberExprRef::from(symbol_ref));
    let stmt_info =
      StmtInfo::default().with_declared_symbols(vec![TaggedSymbolRef::Normal(symbol_ref)]);
    stmt_infos.add_stmt_info(stmt_info);
    module.named_exports.insert(
      exported.clone(),
      LocalExport { span: SPAN, referenced: symbol_ref, came_from_commonjs: false },
    );
  }
  // declare default export statement
  let stmt_info = StmtInfo::default()
    .with_declared_symbols(vec![TaggedSymbolRef::Normal(default_export_ref)])
    .with_referenced_symbols(all_declared_symbols.clone());

  stmt_infos.add_stmt_info(stmt_info);
  module.named_exports.insert(
    "default".into(),
    LocalExport { span: SPAN, referenced: default_export_ref, came_from_commonjs: false },
  );

  // declare namespace object statement
  module.exports_kind = ExportsKind::Esm;
  stmt_infos.replace_namespace_stmt_info(
    StmtInfo::default()
      .with_declared_symbols(vec![TaggedSymbolRef::Normal(namespace_object_ref)])
      .with_referenced_symbols(all_declared_symbols),
  );
  // for a es json module it did not needs to be wrapped anyway.
  link_staged.metas[module_idx].wrapper_stmt_info = None;
  link_staged.metas[module_idx].wrapper_ref = None;
  link_staged.metas[module_idx].sync_wrap_kind(WrapKind::None);

  link_staged.symbols.store_local_db(module_idx, symbol_ref_db);
  true
}
