/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
/*!
 * \file source_map.cc
 * \brief The implementation of the source map data structure.
 */
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/source_map.h>
#include <tvm/ir/transform.h>
#include <tvm/runtime/logging.h>

#include <algorithm>

namespace tvm {

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  SourceNameNode::RegisterReflection();
  SpanNode::RegisterReflection();
  SequentialSpanNode::RegisterReflection();
  SourceNode::RegisterReflection();
  SourceMapObj::RegisterReflection();
  // overrride SourceNameNode to serialization mechanism
  refl::TypeAttrDef<SourceNameNode>()
      .def("__data_to_json__",
           [](const SourceNameNode* node) {
             // simply save as the string
             return node->name;
           })
      .def("__data_from_json__", SourceName::Get);
}

ffi::ObjectPtr<SourceNameNode> GetSourceNameNode(const ffi::String& name) {
  // always return pointer as the reference can change as map re-allocate.
  // or use another level of indirection by creating a unique_ptr
  static std::unordered_map<ffi::String, ffi::ObjectPtr<SourceNameNode>> source_map;

  auto sn = source_map.find(name);
  if (sn == source_map.end()) {
    ffi::ObjectPtr<SourceNameNode> n = ffi::make_object<SourceNameNode>();
    source_map[name] = n;
    n->name = std::move(name);
    return n;
  } else {
    return sn->second;
  }
}

ffi::ObjectPtr<SourceNameNode> GetSourceNameNodeByStr(const std::string& name) {
  return GetSourceNameNode(name);
}

SourceName SourceName::Get(const ffi::String& name) { return SourceName(GetSourceNameNode(name)); }

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("ir.SourceName", SourceName::Get);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::TypeAttrDef<SourceNameNode>().def(
      refl::type_attr::kRepr, [](SourceName sn, ffi::Function) -> ffi::String {
        std::ostringstream os;
        os << "SourceName(" << sn->name << ", " << static_cast<const void*>(sn.get()) << ")";
        return os.str();
      });
}

Span::Span(SourceName source_name, int line, int end_line, int column, int end_column) {
  auto n = ffi::make_object<SpanNode>();
  n->source_name = std::move(source_name);
  n->line = line;
  n->end_line = end_line;
  n->column = column;
  n->end_column = end_column;
  data_ = std::move(n);
}

Span Span::Merge(const Span& other) const {
  TVM_FFI_ICHECK(this->defined() && other.defined()) << "Span::Merge: both spans must be defined";

  TVM_FFI_ICHECK((*this)->source_name == other->source_name);
  return Span((*this)->source_name, std::min((*this)->line, other->line),
              std::max((*this)->end_line, other->end_line),
              std::min((*this)->column, other->column),
              std::max((*this)->end_column, other->end_column));
}

SequentialSpan::SequentialSpan(tvm::ffi::Array<Span> spans) {
  auto n = ffi::make_object<SequentialSpanNode>();
  tvm::ffi::Array<Span> tmp_spans;
  for (const Span& s : spans) {
    if (const SequentialSpanNode* seq_s = s.as<SequentialSpanNode>()) {
      tmp_spans.insert(tmp_spans.end(), seq_s->spans.begin(), seq_s->spans.end());
    } else {
      tmp_spans.push_back(s);
    }
  }
  n->spans = std::move(tmp_spans);

  n->line = 0;
  n->end_line = 0;
  n->column = 0;
  n->end_column = 0;

  data_ = std::move(n);
}

SequentialSpan::SequentialSpan(std::initializer_list<Span> init) {
  auto n = ffi::make_object<SequentialSpanNode>();
  tvm::ffi::Array<Span> spans = tvm::ffi::Array<Span>(init);
  tvm::ffi::Array<Span> tmp_spans;
  for (const Span& s : spans) {
    if (const SequentialSpanNode* seq_s = s.as<SequentialSpanNode>()) {
      tmp_spans.insert(tmp_spans.end(), seq_s->spans.begin(), seq_s->spans.end());
    } else {
      tmp_spans.push_back(s);
    }
  }
  n->spans = std::move(tmp_spans);

  n->line = 0;
  n->end_line = 0;
  n->column = 0;
  n->end_column = 0;

  data_ = std::move(n);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("ir.Span",
           [](SourceName source_name, int line, int end_line, int column, int end_column) {
             return Span(source_name, line, end_line, column, end_column);
           })
      .def("ir.SequentialSpan", [](tvm::ffi::Array<Span> spans) { return SequentialSpan(spans); });
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::TypeAttrDef<SpanNode>().def(
      refl::type_attr::kRepr, [](Span span, ffi::Function fn_repr) -> ffi::String {
        std::ostringstream os;
        os << "Span(" << fn_repr(ffi::AnyView(span->source_name)).cast<ffi::String>() << ", "
           << span->line << ", " << span->end_line << ", " << span->column << ", "
           << span->end_column << ")";
        return os.str();
      });
  refl::TypeAttrDef<SequentialSpanNode>().def(
      refl::type_attr::kRepr, [](SequentialSpan seq, ffi::Function fn_repr) -> ffi::String {
        // Fix typo: was "SequentailSpan", now "SequentialSpan"
        std::ostringstream os;
        os << "SequentialSpan([ ";
        const int last = static_cast<int>(seq->spans.size()) - 1;
        for (int i = 0; i < last; ++i) {
          os << fn_repr(ffi::AnyView(seq->spans[i])).cast<ffi::String>() << ", ";
        }
        if (last >= 0) {
          os << fn_repr(ffi::AnyView(seq->spans[last])).cast<ffi::String>();
        }
        os << " ])";
        return os.str();
      });
}

/*! \brief Construct a source from a string. */
Source::Source(SourceName src_name, std::string source) {
  auto n = ffi::make_object<SourceNode>();
  n->source_name = std::move(src_name);
  n->source = std::move(source);

  int index = 0;
  int length = 0;
  n->line_map.push_back({index, length});
  // NB(@jroesch):
  std::string source_str = n->source;
  for (auto c : source_str) {
    if (c == '\n') {
      // Record the length of the line.
      n->line_map.back().second = length;
      // Bump past the newline.
      index += 1;
      // Record the start of the next line, and put placeholder for length.
      n->line_map.push_back({index, 0});
      // Reset length to zero.
      length = 0;
    } else {
      length += 1;
      index += 1;
    }
  }
  n->line_map.back().second = length;

  data_ = n;
}

tvm::ffi::String Source::GetLine(int line) {
  VLOG(1) << "Source::GetLine: line=" << line;
  TVM_FFI_ICHECK(line - 1 < static_cast<int64_t>((*this)->line_map.size()))
      << "requested line: " << line << "at index: " << (line - 1)
      << "line_map size: " << (*this)->line_map.size() << "source: " << (*this)->source;

  // Adjust for zero indexing, now have (line_start, line_length);
  auto range = (*this)->line_map.at(line - 1);
  int line_start = range.first;
  int line_length = range.second;
  VLOG(1) << "Source::GetLine: line_start=" << line_start << " line_length=" << line_length;
  // TODO(@jroesch): expose substring on tvm::ffi::String.
  auto line_text = std::string((*this)->source).substr(line_start, line_length);
  VLOG(1) << "Source::GetLine: line_text=" << line_text;
  return line_text;
}

SourceMap::SourceMap(ffi::Map<SourceName, Source> source_map) {
  auto n = ffi::make_object<SourceMapObj>();
  n->source_map = std::move(source_map);
  data_ = std::move(n);
}

void SourceMap::Add(const Source& source) { (*this)->source_map.Set(source->source_name, source); }

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("SourceMapAdd", [](SourceMap map, ffi::String name, ffi::String content) {
    auto src_name = SourceName::Get(name);
    Source source(src_name, content);
    map.Add(source);
    return src_name;
  });
}

}  // namespace tvm
