// Copyright 2026 Google LLC
//
// Licensed 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
//
//    https://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.
//

#include "tracing/test_tracing.h"

#include <atomic>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "tracing/span.h"
#include "tracing/tracer.h"
#include "tracing/tracer_factory.h"

namespace credentio_tracing {

namespace {

std::string GenerateTraceId() {
  static std::atomic<uint64_t> counter{1};
  uint64_t id = counter.fetch_add(1, std::memory_order_relaxed);
  return absl::StrFormat("%032x", id);
}

std::string GenerateSpanId() {
  static std::atomic<uint64_t> counter{1};
  uint64_t id = counter.fetch_add(1, std::memory_order_relaxed);
  return absl::StrFormat("%016x", id);
}

}  // namespace

void TestRecordedSpan::SetAttribute(absl::string_view key,
                                    AttributeValue value) {
  absl::MutexLock lock(mutex_);
  for (auto& [k, v] : attributes_) {
    if (k == key) {
      v = std::move(value);
      return;
    }
  }
  attributes_.push_back({std::string(key), std::move(value)});
}

void TestRecordedSpan::AddEvent(absl::string_view name) {
  absl::MutexLock lock(mutex_);
  events_.push_back({std::string(name), absl::Now()});
}

void TestRecordedSpan::SetStatus(const absl::Status& status) {
  absl::MutexLock lock(mutex_);
  status_ = status;
}

void TestRecordedSpan::End() {
  absl::MutexLock lock(mutex_);
  ended_ = true;
  end_time_ = absl::Now();
}

bool TestRecordedSpan::ended() const {
  absl::MutexLock lock(mutex_);
  return ended_;
}

absl::Status TestRecordedSpan::status() const {
  absl::MutexLock lock(mutex_);
  return status_;
}

absl::Time TestRecordedSpan::end_time() const {
  absl::MutexLock lock(mutex_);
  return end_time_;
}

std::vector<std::pair<std::string, AttributeValue>>
TestRecordedSpan::attributes() const {
  absl::MutexLock lock(mutex_);
  return attributes_;
}

std::vector<TestRecordedEvent> TestRecordedSpan::events() const {
  absl::MutexLock lock(mutex_);
  return events_;
}

std::optional<AttributeValue> TestRecordedSpan::GetAttribute(
    absl::string_view key) const {
  absl::MutexLock lock(mutex_);
  for (const auto& [k, v] : attributes_) {
    if (k == key) {
      return v;
    }
  }
  return std::nullopt;
}

TestSpan::TestSpan(std::shared_ptr<TestRecordedSpan> recorded_span)
    : recorded_span_(std::move(recorded_span)) {}

void TestSpan::SetAttribute(absl::string_view key, AttributeValue value) {
  if (recorded_span_) {
    recorded_span_->SetAttribute(key, std::move(value));
  }
}

void TestSpan::AddEvent(absl::string_view name) {
  if (recorded_span_) {
    recorded_span_->AddEvent(name);
  }
}

void TestSpan::SetStatus(const absl::Status& status) {
  if (recorded_span_) {
    recorded_span_->SetStatus(status);
  }
}

void TestSpan::End() {
  if (recorded_span_) {
    recorded_span_->End();
  }
}

std::shared_ptr<TestRecordedSpan> TestSpan::recorded_span() const {
  return recorded_span_;
}

std::unique_ptr<Span> TestTracer::StartSpanImpl(
    absl::string_view span_name, const StartSpanOptions& options) const {
  auto recorded = std::make_shared<TestRecordedSpan>();
  recorded->name_ = std::string(span_name);
  recorded->start_time_ = absl::Now();
  recorded->span_id_ = GenerateSpanId();

  for (const auto& [k, v] : options.attributes) {
    recorded->SetAttribute(k, v);
  }

  // `Tracer::StartSpan` already resolves the implicit parent from the current
  // `ThreadContext`, so `options.parent` is authoritative here.
  recorded->parent_ = options.parent;

  // `dynamic_cast` yields nullptr for both a null parent and a parent that is
  // not a `TestSpan`, in which case the span starts a new trace.
  const auto* test_parent = dynamic_cast<const TestSpan*>(options.parent);
  if (test_parent != nullptr && test_parent->recorded_span() != nullptr) {
    recorded->parent_span_id_ = test_parent->recorded_span()->span_id_;
    recorded->trace_id_ = test_parent->recorded_span()->trace_id_;
  }
  if (recorded->trace_id_.empty()) {
    recorded->trace_id_ = GenerateTraceId();
  }

  absl::MutexLock lock(mutex_);
  spans_.push_back(recorded);
  return std::make_unique<TestSpan>(recorded);
}

std::vector<std::shared_ptr<TestRecordedSpan>> TestTracer::GetSpans() const {
  absl::MutexLock lock(mutex_);
  return spans_;
}

std::shared_ptr<TestRecordedSpan> TestTracer::GetSpan(
    absl::string_view name) const {
  absl::MutexLock lock(mutex_);
  for (const auto& span : spans_) {
    if (span->name_ == name) {
      return span;
    }
  }
  return nullptr;
}

std::vector<std::shared_ptr<TestRecordedSpan>> TestTracer::GetSpansByName(
    absl::string_view name) const {
  absl::MutexLock lock(mutex_);
  std::vector<std::shared_ptr<TestRecordedSpan>> result;
  for (const auto& span : spans_) {
    if (span->name_ == name) {
      result.push_back(span);
    }
  }
  return result;
}

void TestTracer::Clear() {
  absl::MutexLock lock(mutex_);
  spans_.clear();
}

std::shared_ptr<TestTracerFactory> TestTracerFactory::Register() {
  auto factory = std::make_shared<TestTracerFactory>();
  TracerFactory::RegisterGlobalFactory(factory);
  return factory;
}

void TestTracerFactory::Register(std::shared_ptr<TestTracerFactory> factory) {
  TracerFactory::RegisterGlobalFactory(std::move(factory));
}

void TestTracerFactory::RegisterThis() {
  TracerFactory::RegisterGlobalFactory(shared_from_this());
}

void TestTracerFactory::Unregister() {
  TracerFactory::RegisterGlobalFactory(nullptr);
}

std::shared_ptr<Tracer> TestTracerFactory::GetTracer(
    absl::string_view tracer_name) const {
  absl::MutexLock lock(mutex_);
  auto& tracer = tracers_[tracer_name];
  if (tracer == nullptr) {
    tracer = std::make_shared<TestTracer>();
  }
  return tracer;
}

std::shared_ptr<TestTracer> TestTracerFactory::GetTestTracer(
    absl::string_view tracer_name) const {
  absl::MutexLock lock(mutex_);
  auto it = tracers_.find(tracer_name);
  if (it != tracers_.end()) {
    return it->second;
  }
  return nullptr;
}

std::vector<std::shared_ptr<TestRecordedSpan>> TestTracerFactory::GetAllSpans()
    const {
  std::vector<std::shared_ptr<TestTracer>> tracers_copy;
  {
    absl::MutexLock lock(mutex_);
    tracers_copy.reserve(tracers_.size());
    for (const auto& [name, tracer] : tracers_) {
      tracers_copy.push_back(tracer);
    }
  }

  std::vector<std::shared_ptr<TestRecordedSpan>> all_spans;
  for (const auto& tracer : tracers_copy) {
    auto spans = tracer->GetSpans();
    all_spans.insert(all_spans.end(), spans.begin(), spans.end());
  }
  return all_spans;
}

void TestTracerFactory::Clear() {
  absl::MutexLock lock(mutex_);
  tracers_.clear();
}

}  // namespace credentio_tracing