//============================================================================//
//                                                                            //
//                       artichoke programming language                       //
//                                                                            //
//      Copyright (C) 2025 Erick Saul Guzman Ramos, whoami.artichoke.dev      //
//                                                                            //
//                                                                            //
//  This program is free software: you can redistribute it and/or modify      //
//  it under the terms of the GNU Affero General Public License as published  //
//  by the Free Software Foundation, either version 3 of the License, or      //
//  (at your option) any later version.                                       //
//                                                                            //
//  This program is distributed in the hope that it will be useful,           //
//  but WITHOUT ANY WARRANTY; without even the implied warranty of            //
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the             //
//  GNU Affero General Public License for more details.                       //
//                                                                            //
//  You should have received a copy of the GNU Affero General Public License  //
//  along with this program.  If not, see <https://www.gnu.org/licenses/>.    //
//                                                                            //
//============================================================================//

#pragma once

#include <utility>
#include <variant>
#include <iterator>
#include <coroutine>
#include <exception>
#include <type_traits>

#include <artichoke/Coroutine/Utils.hpp>
#include <artichoke/Util/Expected.hpp>

namespace arti::lang {

  template <typename T>
  requires (not std::is_void_v<T>)
  struct Generator;

  template <typename T>
  requires (not std::is_void_v<T>)
  struct GeneratorPromise {
    using value_type = std::remove_cvref_t<T>;
    using reference_type = value_type &;
    using pointer_type = value_type *;

    using ValueType = value_type;
    using ReferenceType = reference_type;
    using PointerType = pointer_type;

    GeneratorPromise() = default;

    Generator<T> get_return_object();

    std::suspend_always initial_suspend() const noexcept {
      return {};
    }

    std::suspend_always final_suspend() const noexcept {
      return {};
    }

    void unhandled_exception() {
      exception = std::current_exception();
    }

    void rethrow_if_exception() {
      if (exception) {
        std::rethrow_exception(exception);
      }
    }

    void return_value(const ValueType &val) noexcept {
      value = val;
    }

    void return_value(ValueType &&val) noexcept {
      value = std::forward<ValueType &&>(val);
    }

    std::suspend_always yield_value(const ValueType &val) noexcept {
      value = val;
      return {};
    }

    std::suspend_always yield_value(ValueType &&val) noexcept {
      value = std::forward<ValueType &&>(val);
      return {};
    }

    decltype(auto) await_transform(auto &&awaitable) = delete;

    std::exception_ptr exception;
    std::variant<PointerType, ValueType> value;
  };

  struct GeneratorSentinel { };

  template <typename T>
  requires (not std::is_void_v<T>)
  struct GeneratorIterator {
    using iterator_category = std::input_iterator_tag;
    using difference_type = std::ptrdiff_t;

    using promise_type = GeneratorPromise<T>;

    using value_type = typename promise_type::value_type;
    using reference_type = value_type &;
    using pointer_type = value_type *;

    using handle_type = std::coroutine_handle<promise_type>;

    using IteratorCategory = iterator_category;
    using DifferenceType = difference_type;
    using PromiseType = promise_type;
    using ValueType = value_type;
    using ReferenceType = reference_type;
    using PointerType = pointer_type;
    using HandleType = handle_type;

    GeneratorIterator() noexcept = default;
    ~GeneratorIterator() noexcept = default;

    GeneratorIterator(HandleType handle) noexcept
      : handle(handle) { }

    GeneratorIterator(GeneratorIterator &&rhs) noexcept
      : handle(std::exchange(rhs.handle, nullptr)) { }

    GeneratorIterator &operator=(GeneratorIterator &&rhs) noexcept {
      handle = std::exchange(rhs.handle, nullptr);
      return *this;
    }

    GeneratorIterator(const GeneratorIterator &) noexcept = delete;
    GeneratorIterator &operator=(const GeneratorIterator &) noexcept = delete;

    GeneratorIterator &operator++() {
      handle.resume();

      if (handle.done()) {
        handle.promise().rethrow_if_exception();
      }

      return *this;
    }

    void operator++(int) {
      (void) this->operator++();
    }

    ReferenceType operator*() const noexcept {
      auto &ref = handle.promise().value;

      if (std::holds_alternative<PointerType>(ref)) {
        return *std::get<PointerType>(ref);
      }
      else {
        return std::get<ValueType>(ref);
      }
    }

    PointerType operator->() const noexcept {
      auto &ref = handle.promise().value;

      if (std::holds_alternative<PointerType>(ref)) {
        return std::get<PointerType>(ref);
      }
      else {
        return &std::get<ValueType>(ref);
      }
    }

    friend bool operator==(const GeneratorIterator &it, GeneratorSentinel) {
      return (! it.handle || it.handle.done());
    }

    friend bool operator!=(const GeneratorIterator &it, GeneratorSentinel s) {
      return ! (it == s);
    }

    friend bool operator==(GeneratorSentinel s, const GeneratorIterator &it) {
      return it == s;
    }

    friend bool operator!=(GeneratorSentinel s, const GeneratorIterator &it) {
      return it != s;
    }

    HandleType handle;
  };

  template <typename T>
  requires (not std::is_void_v<T>)
  struct [[nodiscard]] Generator {
    using promise_type = GeneratorPromise<T>;
    using handle_type = std::coroutine_handle<promise_type>;

    using iterator_type = GeneratorIterator<T>;
    using sentinel_type = GeneratorSentinel;

    using PromiseType = promise_type;
    using HandleType = handle_type;
    using IteratorType = iterator_type;
    using SentinelType = sentinel_type;

    Generator() noexcept = default;

    ~Generator() noexcept {
      if (handle) {
        handle.destroy();
      }
    }

    explicit Generator(HandleType handle) noexcept
      : handle(handle) { }

    Generator(Generator &&rhs) noexcept
      : handle(std::exchange(rhs.handle, nullptr)) { }

    Generator &operator=(Generator &&rhs) noexcept {
      this->~Generator();
      handle = std::exchange(rhs.handle, nullptr);
      return *this;
    }

    Generator(const Generator &) noexcept = delete;
    Generator &operator=(const Generator &) noexcept = delete;

    bool finished() const noexcept {
      return (not handle || (handle.done() and !handle.promise().exception));
    }

    void next() {
      if (not finished()) {
        handle.resume();

        if (handle.done()) {
          handle.promise().rethrow_if_exception();
        }
      }
    }

    PromiseType::ReferenceType get() {
      using PointerType = PromiseType::PointerType;
      using ValueType = PromiseType::ValueType;

      auto &ref = handle.promise().value;

      if (std::holds_alternative<PointerType>(ref)) {
        return *std::get<PointerType>(ref);
      }
      else {
        return std::get<ValueType>(ref);
      }
    }

    IteratorType begin() {
      if (handle) {
        handle.resume();

        if (handle.done()) {
          handle.promise().rethrow_if_exception();
        }
      }

      return { handle };
    }

    SentinelType end() const noexcept {
      return {};
    }

    void swap(Generator &other) noexcept {
      std::swap(handle, other.handle);
    }

   private:
    HandleType handle;
  };

  template <typename T>
  requires (not std::is_void_v<T>)
  Generator<T> GeneratorPromise<T>::get_return_object() {
    using GeneratorType = Generator<T>;
    using HandleType = typename GeneratorType::HandleType;

    return GeneratorType{ HandleType::from_promise(*this) };
  }

} // namespace arti::lang