#include "cached_shortest_sequences.h"

#include "find_shortest_sequences.h"
#include "try_split_sequence.h"
#include "segmented_sequences.h"

#include <iostream>

std::shared_ptr<SegmentedSequences> CachedShortestSequences::Find(
    std::string_view code,
    std::function<std::optional<char>(char, char)> apply_cmd) {
  std::string cache_key(code.data(), code.size());
  auto cache_iter = cache_.find(cache_key);
  if (cache_iter != cache_.end()) {
    return cache_iter->second;
  }

  auto maybe_split = TrySplitSequence(code);
  if (!maybe_split) {
    auto result_strs = FindShortestSequences(code, apply_cmd);
    std::vector<std::shared_ptr<SegmentedSequences>> segmented_sequences_leaves;
    for (auto &result_str : result_strs) {
      segmented_sequences_leaves.emplace_back(
          std::make_shared<SegmentedSequencesLeaf>(
              std::make_shared<std::string>(std::move(result_str))));
    }

    auto segmented_sequences_alts =
        std::make_shared<SegmentedSequencesAlternatives>(
            segmented_sequences_leaves);
    cache_.emplace(cache_key, segmented_sequences_alts);
    return segmented_sequences_alts;
  }

  auto segmented_sequences_product =
      std::make_shared<SegmentedSequencesProduct>(
          Find(maybe_split->first, apply_cmd),
          Find(maybe_split->second, apply_cmd));
  cache_.emplace(cache_key, segmented_sequences_product);

  return segmented_sequences_product;
}

std::shared_ptr<SegmentedSequences> CachedShortestSequences::Find(
    std::shared_ptr<SegmentedSequences> codes,
    std::function<std::optional<char>(char, char)> apply_cmd) {
  auto segmented_cache_iter = segmented_cache_.find(codes.get());
  if (segmented_cache_iter != segmented_cache_.end()) {
    return segmented_cache_iter->second;
  }

  auto leaf = std::dynamic_pointer_cast<SegmentedSequencesLeaf>(codes);
  if (leaf != nullptr) {
    auto leaf_result = Find(leaf->sequence(), apply_cmd);
    segmented_cache_.emplace(codes.get(), leaf_result);
    return leaf_result;
  }

  auto alts = std::dynamic_pointer_cast<SegmentedSequencesAlternatives>(codes);
  if (alts != nullptr) {
    std::vector<std::shared_ptr<SegmentedSequences>> alts_results;
    for (auto alt : alts->alternatives()) {
      alts_results.emplace_back(Find(alt, apply_cmd));
    }
    auto alt_result = std::make_shared<SegmentedSequencesAlternatives>(
        std::move(alts_results));
    segmented_cache_.emplace(codes.get(), alt_result);
    return alt_result;
  }

  auto product = std::dynamic_pointer_cast<SegmentedSequencesProduct>(codes);
  if (product != nullptr) {
    auto product_result = std::make_shared<SegmentedSequencesProduct>(
        Find(product->a(), apply_cmd), Find(product->b(), apply_cmd));
    segmented_cache_.emplace(codes.get(), product_result);
    return product_result;
  }

  return nullptr;
}