#pragma once
#include <atomic>
#include <functional>
#include <mutex>
#include <optional>
#include <string>
#include <unordered_map>
#include <chrono>
#include <taskflow/taskflow.hpp>

namespace lsp::scheduler
{
    class AsyncExecutor
    {
    public:
        using TaskFunc = std::function<std::optional<std::string>()>;
        using Callback = std::function<void(const std::string&)>;
        using Clock = std::chrono::steady_clock;
        using TimePoint = Clock::time_point;

        explicit AsyncExecutor(size_t concurrency = std::thread::hardware_concurrency());
        ~AsyncExecutor();

        // Delete copy operations
        AsyncExecutor(const AsyncExecutor&) = delete;
        AsyncExecutor& operator=(const AsyncExecutor&) = delete;

        // Public interface
        void Submit(const std::string& task_id, TaskFunc task, Callback callback = nullptr);
        bool Cancel(const std::string& task_id);
        void WaitAll();

        // Status and statistics
        size_t GetRunningTaskCount() const;
        void LogStatus() const;

        struct Statistics
        {
            size_t running = 0;
            uint64_t submitted = 0;
            uint64_t completed = 0;
            uint64_t failed = 0;
            uint64_t cancelled = 0;
        };

        Statistics GetStatistics() const;

    private:
        struct TaskContext
        {
            std::atomic<bool> cancelled { false };
            Callback callback {nullptr};
            TimePoint start_time;

            TaskContext() :
                start_time(Clock::now()) {}
        };

        struct StatsManager
        {
            std::atomic<uint64_t> submitted{0};
            std::atomic<uint64_t> completed{0};
            std::atomic<uint64_t> failed{0};
            std::atomic<uint64_t> cancelled{0};

            void IncrementSubmitted() { submitted++; }
            void IncrementCompleted() { completed++; }
            void IncrementFailed() { failed++; }
            void IncrementCancelled() { cancelled++; }

            Statistics GetSnapshot(size_t running_count) const
            {
                return {
                    .running = running_count,
                    .submitted = submitted.load(),
                    .completed = completed.load(),
                    .failed = failed.load(),
                    .cancelled = cancelled.load()
                };
            }
        };

        using TaskContextPtr = std::shared_ptr<TaskContext>;

        // Task lifecycle management
        void ExecuteTask(const std::string& task_id, TaskFunc task, TaskContextPtr context);
        bool RegisterTask(const std::string& task_id, TaskContextPtr context);
        void UnregisterTask(const std::string& task_id);
        bool MarkTaskCancelled(const std::string& task_id);

        // Logging helpers
        void LogTaskStart(const std::string& task_id, bool replaced_existing) const;
        void LogTaskCompletion(const std::string& task_id, TimePoint start_time, bool success) const;
        std::chrono::milliseconds GetElapsedTime(TimePoint start_time) const;

    private:
        // Core components
        tf::Executor executor_;
        mutable std::mutex mutex_;
        std::unordered_map<std::string, TaskContextPtr> running_tasks_;
        StatsManager stats_;

        // Configuration
        static constexpr size_t kStatusLogInterval = 10;
    };
}
