// 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.
//

#ifndef THIRD_PARTY_CREDENTIO_TRACING_TEST_TRACING_H_
#define THIRD_PARTY_CREDENTIO_TRACING_TEST_TRACING_H_

#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.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 {

// Simple struct representing a recorded point-in-time event within a span.
struct TestRecordedEvent {
  std::string name;
  absl::Time time = absl::Now();

  bool operator==(const TestRecordedEvent& other) const {
    return name == other.name;
  }
};

// Simple struct holding recorded span data for verification in tests.
class TestRecordedSpan {
 public:
  TestRecordedSpan() = default;

  std::string name_;
  std::string trace_id_;
  std::string span_id_;
  std::string parent_span_id_;
  const Span* parent_ = nullptr;
  absl::Time start_time_ = absl::Now();

  // Set or update an attribute key-value pair.
  void SetAttribute(absl::string_view key, AttributeValue value);

  // Record an event.
  void AddEvent(absl::string_view name);

  // Set status.
  void SetStatus(const absl::Status& status);

  // Mark span as ended.
  void End();

  // Returns true if the span has ended.
  bool ended() const;

  // Returns the recorded status.
  absl::Status status() const;

  // Returns the recorded end time.
  absl::Time end_time() const;

  // Returns a copy of recorded attributes.
  std::vector<std::pair<std::string, AttributeValue>> attributes() const;

  // Returns a copy of recorded events.
  std::vector<TestRecordedEvent> events() const;

  // Returns optional copy of attribute value for specified key.
  std::optional<AttributeValue> GetAttribute(absl::string_view key) const;

  // Returns the attribute value converted to type T if found and matching, or
  // std::nullopt.
  template <typename T>
  std::optional<T> GetAttributeAs(absl::string_view key) const {
    std::optional<AttributeValue> val = GetAttribute(key);
    if (val && std::holds_alternative<T>(*val)) {
      return std::get<T>(*val);
    }
    return std::nullopt;
  }

 private:
  mutable absl::Mutex mutex_;
  std::vector<std::pair<std::string, AttributeValue>> attributes_
      ABSL_GUARDED_BY(mutex_);
  std::vector<TestRecordedEvent> events_ ABSL_GUARDED_BY(mutex_);
  absl::Status status_ ABSL_GUARDED_BY(mutex_) = absl::OkStatus();
  bool ended_ ABSL_GUARDED_BY(mutex_) = false;
  absl::Time end_time_ ABSL_GUARDED_BY(mutex_) = absl::InfinitePast();
};

// Test implementation of Span that records operations to a TestRecordedSpan
// struct.
class TestSpan : public Span {
 public:
  explicit TestSpan(std::shared_ptr<TestRecordedSpan> recorded_span);

  void SetAttribute(absl::string_view key, AttributeValue value) override;
  void AddEvent(absl::string_view name) override;
  void SetStatus(const absl::Status& status) override;
  void End() override;

  // Returns the underlying TestRecordedSpan struct pointer.
  std::shared_ptr<TestRecordedSpan> recorded_span() const;

 private:
  const std::shared_ptr<TestRecordedSpan> recorded_span_;
};

// Test implementation of Tracer that records created spans.
class TestTracer : public Tracer {
 public:
  using Tracer::StartSpan;

  TestTracer() = default;

  std::unique_ptr<Span> StartSpanImpl(
      absl::string_view span_name,
      const StartSpanOptions& options) const override;

  // Returns all recorded spans.
  std::vector<std::shared_ptr<TestRecordedSpan>> GetSpans() const;

  // Returns the first span matching the given name, or nullptr if not found.
  std::shared_ptr<TestRecordedSpan> GetSpan(absl::string_view name) const;

  // Returns all spans matching the given name.
  std::vector<std::shared_ptr<TestRecordedSpan>> GetSpansByName(
      absl::string_view name) const;

  // Clears all recorded spans.
  void Clear();

 private:
  mutable absl::Mutex mutex_;
  mutable std::vector<std::shared_ptr<TestRecordedSpan>> spans_
      ABSL_GUARDED_BY(mutex_);
};

// Test implementation of TracerFactory that manages TestTracer instances.
class TestTracerFactory
    : public TracerFactory,
      public std::enable_shared_from_this<TestTracerFactory> {
 public:
  TestTracerFactory() = default;

  // Creates a new TestTracerFactory instance, registers it as the global
  // factory, and returns it.
  static std::shared_ptr<TestTracerFactory> Register();

  // Registers the specified factory instance as the global factory.
  static void Register(std::shared_ptr<TestTracerFactory> factory);

  // Registers this TestTracerFactory instance as the global factory.
  void RegisterThis();

  // Unregisters the global TracerFactory instance, resetting to default.
  static void Unregister();

  std::shared_ptr<Tracer> GetTracer(
      absl::string_view tracer_name) const override;

  // Returns the TestTracer associated with tracer_name, or nullptr if not
  // found.
  std::shared_ptr<TestTracer> GetTestTracer(
      absl::string_view tracer_name) const;

  // Returns all recorded spans across all managed tracers.
  std::vector<std::shared_ptr<TestRecordedSpan>> GetAllSpans() const;

  // Clears all registered tracers and their recorded spans.
  void Clear();

 private:
  mutable absl::Mutex mutex_;
  mutable absl::flat_hash_map<std::string, std::shared_ptr<TestTracer>> tracers_
      ABSL_GUARDED_BY(mutex_);
};

}  // namespace credentio_tracing

#endif  // THIRD_PARTY_CREDENTIO_TRACING_TEST_TRACING_H_