#include "drake/multibody/parsing/detail_dmd_parser.h"

#include <filesystem>
#include <memory>
#include <set>

#include "drake/common/yaml/yaml_io.h"
#include "drake/multibody/parsing/detail_make_model_name.h"
#include "drake/multibody/parsing/detail_path_utils.h"
#include "drake/multibody/parsing/scoped_names.h"
#include "drake/multibody/tree/quaternion_floating_joint.h"
#include "drake/multibody/tree/scoped_name.h"

namespace drake {
namespace multibody {
namespace internal {

using parsing::ModelDirectives;
using parsing::ModelInstanceInfo;
using parsing::GetScopedFrameByName;

namespace {

// Adds a new weld joint to @p plant from @p parent_frame as indicated by @p
// weld (as resolved relative to @p model_namespace) and update the @p info for
// the child model accordingly.
void AddWeld(
    const Frame<double>& parent_frame,
    const Frame<double>& child_frame,
    const math::RigidTransform<double>& X_PC,
    MultibodyPlant<double>* plant,
    std::vector<ModelInstanceInfo>* added_models) {
  plant->WeldFrames(parent_frame, child_frame, X_PC);
  if (added_models) {
    // Record weld info into crappy ModelInstanceInfo struct.
    bool found = false;
    for (auto& info : *added_models) {
      if (info.model_instance == child_frame.model_instance()) {
        found = true;
        // See warning in ModelInstanceInfo about these members.
        info.parent_frame_name = parent_frame.name();
        info.child_frame_name = child_frame.name();
        info.X_PC = X_PC;
      }
    }
    DRAKE_THROW_UNLESS(found);
  }
}

void ParseModelDirectivesImpl(
    const ModelDirectives& directives,
    const std::string& model_namespace,
    const ParsingWorkspace& workspace,
    std::vector<ModelInstanceInfo>* added_models) {
  drake::log()->debug("ParseModelDirectivesImpl(MultibodyPlant)");
  DRAKE_DEMAND(added_models != nullptr);
  auto& [options, package_map, diagnostic, plant,
         collision_resolver, parser_selector] = workspace;
  DRAKE_DEMAND(plant != nullptr);
  auto get_scoped_frame = [plant = plant, &model_namespace](
                              const std::string& name) -> const Frame<double>& {
    return GetScopedFrameByName(
        *plant, DmdScopedNameJoin(model_namespace, name).to_string());
  };

  for (auto& directive : directives.directives) {
    if (directive.add_model) {
      ModelInstanceInfo info;
      auto& model = *directive.add_model;
      const std::string name =
          ScopedName::Join(model_namespace, model.name).to_string();
      drake::log()->debug("  add_model: {}\n    {}", name, model.file);
      const std::string file =
          ResolveUri(diagnostic, model.file, package_map, {});
      std::optional<ModelInstanceIndex> child_model_instance_id =
          parser_selector(diagnostic, file).AddModel(
              {DataSource::kFilename, &file},
              model.name, model_namespace, workspace);
      if (!child_model_instance_id.has_value()) {
        // Error should have already been emitted.
        continue;
      }
      for (const auto& [joint_name, positions] :
               directive.add_model->default_joint_positions) {
        plant->GetMutableJointByName(joint_name, child_model_instance_id)
            .set_default_positions(positions);
      }
      for (const auto& [element_name, pose] :
           directive.add_model->default_free_body_pose) {
        const Frame<double>& child_frame =
            plant->GetFrameByName(element_name, *child_model_instance_id);
        if (pose.base_frame.has_value() && *pose.base_frame != "world") {
          // TODO(SeanCurtis-TRI): When the new multibody graph code lands,
          // update this code to test to see if there is already a joint between
          // the two bodies.

          // Note: the logic for generating the joint name is borrowed from
          // MultibodyTree::CreateJointImplementations().
          std::string joint_name = child_frame.body().name();
          while (plant->HasJointNamed(joint_name, *child_model_instance_id)) {
            joint_name = "_" + joint_name;
          }

          const Frame<double>& parent_frame =
              get_scoped_frame(*pose.base_frame);
          plant->AddJoint(std::make_unique<QuaternionFloatingJoint<double>>(
              joint_name, parent_frame, child_frame));
          auto& joint =
              plant->GetMutableJointByName(joint_name, child_model_instance_id);
          joint.SetDefaultPose(pose.GetDeterministicValue());
        } else {
          // We can only call this if we haven't injected the floating joint.
          plant->SetDefaultFreeBodyPose(child_frame.body(),
                                        pose.GetDeterministicValue());
        }
      }
      info.model_instance = *child_model_instance_id;
      info.model_name = name;
      info.model_path = file;
      if (added_models) added_models->push_back(info);

    } else if (directive.add_model_instance) {
      auto& instance = *directive.add_model_instance;
      const std::string name =
          MakeModelName(instance.name, model_namespace, workspace);
      drake::log()->debug("  add_model_instance: {}", name);
      plant->AddModelInstance(name);

    } else if (directive.add_frame) {
      auto& frame = *directive.add_frame;
      drake::log()->debug("  add_frame: {}", frame.name);
      if (!frame.X_PF.base_frame) {
        // This would be caught elsewhere, but it is clearer to throw here.
        throw std::logic_error(
            "add_frame directive with empty base frame is ambiguous");
      }
      // Only override instance if scope is explicitly specified.
      std::optional<ModelInstanceIndex> instance;
      ScopedName parsed = ScopedName::Parse(frame.name);
      if (!parsed.get_namespace().empty()) {
        // N.B. Here we're using Join on two namespaces to make a new namespace,
        // with no "element" involved. This is an atypical use of ScopedName,
        // but effective for our needs.
        parsed.set_namespace(
            ScopedName::Join(model_namespace, parsed.get_namespace())
                .get_full());
        instance = plant->GetModelInstanceByName(parsed.get_namespace());
      }
      auto& added = plant->AddFrame(std::make_unique<FixedOffsetFrame<double>>(
          std::string{parsed.get_element()},
          get_scoped_frame(*frame.X_PF.base_frame),
          frame.X_PF.GetDeterministicValue(), instance));
      const std::string resolved_name =
          ScopedName::Join(plant->GetModelInstanceName(added.model_instance()),
                           added.name())
              .to_string();
      drake::log()->debug("    resolved_name: {}", resolved_name);

    } else if (directive.add_weld) {
      math::RigidTransform<double> X_PC{};
      if (directive.add_weld->X_PC) {
        X_PC = directive.add_weld->X_PC->GetDeterministicValue();
      }
      AddWeld(
          get_scoped_frame(directive.add_weld->parent),
          get_scoped_frame(directive.add_weld->child),
          X_PC, plant, added_models);

    } else if (directive.add_collision_filter_group) {
      // If there's no geometry registered, there's nothing to be done with
      // collision filtering.  Trying to proceed will just trigger an error
      // eventually.
      if (!plant->geometry_source_is_registered()) {
        continue;
      }
      auto& group = *directive.add_collision_filter_group;

      // Find the model instance index that corresponds to model_namespace, if
      // the name is non-empty.
      std::optional<ModelInstanceIndex> model_instance;
      if (!model_namespace.empty()) {
        model_instance = plant->GetModelInstanceByName(model_namespace);
      } else if (group.model_namespace.has_value()) {
        model_instance = plant->GetModelInstanceByName(*group.model_namespace);
      }

      drake::log()->debug("  add_collision_filter_group: {}", group.name);
      std::set<std::string> member_set(group.members.begin(),
                                       group.members.end());
      std::set<std::string> member_groups_set(group.member_groups.begin(),
                                              group.member_groups.end());
      collision_resolver->AddGroup(diagnostic, group.name, member_set,
                                   member_groups_set, model_instance);
      for (const auto& ignored_group : group.ignored_collision_filter_groups) {
        collision_resolver->AddPair(
            diagnostic, group.name, ignored_group, model_instance);
      }

    } else {
      // Recurse.
      auto& sub = *directive.add_directives;
      std::string new_model_namespace = ScopedName::Join(
          model_namespace, sub.model_namespace.value_or("")).to_string();
      // Ensure we have a model instance for this namespace.
      drake::log()->debug("  add_directives: {}", sub.file);
      drake::log()->debug("    new_model_namespace: {}", new_model_namespace);
      if (!new_model_namespace.empty() &&
          !plant->HasModelInstanceNamed(new_model_namespace)) {
        throw std::runtime_error(fmt::format(
            "Namespace '{}' does not exist as model instance",
            new_model_namespace));
      }
      std::string filename = ResolveUri(diagnostic, sub.file, package_map, {});
      auto sub_directives =
          LoadModelDirectives({DataSource::kFilename, &filename});
      ParseModelDirectivesImpl(sub_directives, new_model_namespace, workspace,
                               added_models);
    }
  }
}

}  // namespace

ScopedName DmdScopedNameJoin(const std::string& namespace_name,
                             const std::string& element_name) {
  if (element_name == "world") return ScopedName("", element_name);
  return ScopedName::Join(namespace_name, element_name);
}

ModelDirectives LoadModelDirectives(const DataSource& data_source) {
  // Even though the 'defaults' we use to start parsing here are empty, by
  // providing any defaults at all, the effect during parsing will be that any
  // of the users' ModelDirective structs and sub-structs will _also_ start
  // from their default values and allow for overwriting only a subset of the
  // data fields instead of requiring that the user provide them all.
  const ModelDirectives defaults;

  ModelDirectives directives;
  if (data_source.IsFilename()) {
    const std::string& filename = data_source.filename();
    drake::log()->debug("LoadModelDirectives: {}", filename);

    if (!std::filesystem::exists({filename})) {
      throw std::runtime_error(
          fmt::format("No such file {} during LoadModelDirectives", filename));
    }

    directives = yaml::LoadYamlFile<ModelDirectives>(
        filename, std::nullopt /* child_name */, defaults);
  } else {
    DRAKE_DEMAND(data_source.IsContents());
    directives = yaml::LoadYamlString<ModelDirectives>(
        data_source.contents(), std::nullopt /* child_name */, defaults);
  }
  DRAKE_THROW_UNLESS(directives.IsValid());
  return directives;
}

std::vector<ModelInstanceInfo> ParseModelDirectives(
    const ModelDirectives& directives,
    const std::optional<std::string>& parent_model_name,
    const ParsingWorkspace& workspace) {
  std::vector<ModelInstanceInfo> added_models;
  ParseModelDirectivesImpl(directives, parent_model_name.value_or(""),
                           workspace, &added_models);
  return added_models;
}

DmdParserWrapper::DmdParserWrapper() {}

DmdParserWrapper::~DmdParserWrapper() {}

std::optional<ModelInstanceIndex> DmdParserWrapper::AddModel(
    const DataSource& data_source, const std::string&,
    const std::optional<std::string>&,
    const ParsingWorkspace& workspace) {
  std::string display_source =
      data_source.IsFilename()
      ? data_source.GetAbsolutePath()
      : data_source.GetStem() + ".dmd.yaml";
  workspace.diagnostic.Error(fmt::format(
      "'{}' is a model directives data source; it is always an error to pass"
      " a model directives source to a single-model parser method. Use"
      " AddModels() instead.", display_source));
  return {};
}

std::vector<ModelInstanceIndex> DmdParserWrapper::AddAllModels(
    const DataSource& data_source,
    const std::optional<std::string>& parent_model_name,
    const ParsingWorkspace& workspace) {
  // TODO(#18052): diagnostic policy?
  ModelDirectives directives = LoadModelDirectives(data_source);
  // TODO(#18052): diagnostic policy?
  const std::vector<ModelInstanceInfo> infos =
      ParseModelDirectives(directives, parent_model_name, workspace);
  std::vector<ModelInstanceIndex> results;
  results.reserve(infos.size());
  for (const auto& info : infos) {
    results.push_back(info.model_instance);
  }
  return results;
}

}  // namespace internal
}  // namespace multibody
}  // namespace drake
