Untitled

// Task 1
#include "primer/trie.h"
#include <iostream>
#include <memory>
#include <string_view>
#include "common/exception.h"

namespace bustub {

template <class T>
auto Trie::Get(std::string_view key) const -> const T * {
  std::cout << "get key = " << key << std::endl;
  auto cur = this->root_;
  if(cur == nullptr) {
    return nullptr;
  }
  
  for(char i : key){
    if(cur->children_.find(i) != cur->children_.end()){ 
      cur = cur->children_.find(i)->second;
    } else {
      return nullptr;
    }
    // if(cur == nullptr){ // ??
    //   return nullptr;
    // }
  }

  if(cur->is_value_node_){
    auto temp = std::dynamic_pointer_cast<const TrieNodeWithValue<T>>(cur); // ???
    if(temp == nullptr){
      return nullptr;
    }
    return temp->value_.get();
  }
  return nullptr;

  // You should walk through the trie to find the node corresponding to the key. If the node doesn't exist, return
  // nullptr. After you find the node, you should use `dynamic_cast` to cast it to `const TrieNodeWithValue<T> *`. If
  // dynamic_cast returns `nullptr`, it means the type of the value is mismatched, and you should return nullptr.
  // Otherwise, return the value.
}

template <class T>
auto Trie::Put(std::string_view key, T value) const -> Trie {
  // Note that `T` might be a non-copyable type. Always use `std::move` when creating `shared_ptr` on that value.

  std::cout << "put key = " << key << std::endl;

  

  
  if(!this->root_){ // empty trie
    auto new_root = std::make_shared<TrieNode>();

    if(key.empty()){
      auto value_node = std::make_shared<TrieNodeWithValue<T>>(std::make_shared<T>(std::move(value)));
      return Trie(value_node);
    }

    int count = 0;
    auto cur = new_root;
    
    for(char i : key) {
      count++;
      auto it = cur->children_.find(i);
      if(it != cur->children_.end()){ // clone since it is exist
        // auto new_node = std::shared_ptr<TrieNode>(it->second->Clone()); 
        auto new_node = std::make_shared<TrieNode>(it->second->children_); 
        auto old_children = it->second->children_; 
        std::cout << "old_children size = " << old_children.size() << std::endl;
        if(count == static_cast<int>(key.size())){
          auto new_node2 = std::make_shared<TrieNodeWithValue<T>>(old_children, std::make_shared<T>(std::move(value))); // children!!!
          cur->children_[i] = new_node2; 
          cur = new_node2; 
          return Trie(new_root); 
          continue; 
        }
        cur->children_[i] = new_node;
        cur = new_node;             
      }else{
        auto new_node = std::make_shared<TrieNode>();
        if(count == static_cast<int>(key.size())){
          new_node = std::make_shared<TrieNodeWithValue<T>>(std::make_shared<T>(std::move(value)));
        }
        cur->children_[i] = new_node;
        cur = new_node;
      }
    }
    
    return Trie(new_root);
  } 
  
  
  
  auto new_root = std::make_shared<TrieNode>(this->root_->children_);
  auto cur = new_root;

  if(key.empty()){
    auto old_children = new_root->children_;
    auto value_node = std::make_shared<TrieNodeWithValue<T>>(old_children, std::make_shared<T>(std::move(value)));
    return Trie(value_node);
  }

  
  int count2 = 0;
  for(char i : key) {
    count2++;
    auto it = cur->children_.find(i);
    if(it != cur->children_.end()){ // clone since it is exist
      // auto new_node = std::shared_ptr<TrieNode>(it->second->Clone()); 
      
      auto new_node = std::make_shared<TrieNode>(it->second->children_); // problem!!!
      if(it->second->is_value_node_ && count2 != static_cast<int>(key.size())){ // add!!!!!!
        
        auto old_children = it->second->children_;
        // auto new_node3 = std::make_shared<TrieNodeWithValue<T>>(old_children, std::make_shared<T>(std::move(*temp_value2)));
        auto new_node3 = it->second->Clone();
        cur->children_[i] = new_node3;
        cur = new_node3;  
        continue;
      }
      auto old_children = it->second->children_; 
      
      if(count2 == static_cast<int>(key.size())){
        // auto new_node2 = std::make_shared<TrieNodeWithValue<T>>(old_children, std::make_shared<T>(std::move(value)));
        auto new_node2 = std::make_shared<TrieNodeWithValue<T>>(old_children, std::make_shared<T>(std::move(value)));
        cur->children_[i] = new_node2;
        cur = new_node2;
        return Trie(new_root); 
        continue;
      }
      cur->children_[i] = new_node;
      cur = new_node;             
    }else{
      auto new_node = std::make_shared<TrieNode>();
      if(count2 == static_cast<int>(key.size())){
        new_node = std::make_shared<TrieNodeWithValue<T>>(std::make_shared<T>(std::move(value)));
      }
      cur->children_[i] = new_node;
      cur = new_node;
    }
  }

  auto old = cur->children_;
  auto value_node = std::make_shared<TrieNodeWithValue<T>>(old, std::make_shared<T>(std::move(value)));
  cur->children_[key.back()] = value_node;

  return Trie(new_root);
  

  // You should walk through the trie and create new nodes if necessary. If the node corresponding to the key already
  // exists, you should create a new `TrieNodeWithValue`.
}

auto Trie::Remove(std::string_view key) const -> Trie {
  std::cout << "remove key = " << key << std::endl;

  if(!this->root_){
    return Trie(nullptr);
  }
  
  auto new_root = std::make_shared<TrieNode>(this->root_->children_);
  auto cur = new_root;
  int count = 0;

  for(char i : key){
    count++;
    auto it = cur->children_.find(i);
    if(it != cur->children_.end()){ 
      auto new_node = std::make_shared<TrieNode>(it->second->children_);
      if(count == static_cast<int>(key.size())){
        if(new_node->children_.empty()){
          std::map<char, std::shared_ptr<const TrieNode>> m;
          cur->children_ = m;
        }
      }else{
        cur->children_[i] = new_node;
        cur = new_node;
      }
    } 
  }

  return Trie(new_root);

  // You should walk through the trie and remove nodes if necessary. If the node doesn't contain a value any more,
  // you should convert it to `TrieNode`. If a node doesn't have children any more, you should remove it.
  
}

// Below are explicit instantiation of template functions.
//
// Generally people would write the implementation of template classes and functions in the header file. However, we
// separate the implementation into a .cpp file to make things clearer. In order to make the compiler know the
// implementation of the template functions, we need to explicitly instantiate them here, so that they can be picked up
// by the linker.

template auto Trie::Put(std::string_view key, uint32_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint32_t *;

template auto Trie::Put(std::string_view key, uint64_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint64_t *;

template auto Trie::Put(std::string_view key, std::string value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const std::string *;

// If your solution cannot compile for non-copy tests, you can remove the below lines to get partial score.

using Integer = std::unique_ptr<uint32_t>;

template auto Trie::Put(std::string_view key, Integer value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const Integer *;

template auto Trie::Put(std::string_view key, MoveBlocked value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const MoveBlocked *;

}  // namespace bustub