#include <spdlog/spdlog.h>
#include "./async_executor.hpp"

namespace lsp::scheduler
{
    namespace
    {
        constexpr const char* kLogTag = "AsyncExecutor";

        // Helper for formatting durations
        template<typename Duration>
        std::string FormatDuration(Duration duration)
        {
            auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
            if (ms < 1000)
                return std::to_string(ms) + "ms";
            else
                return std::to_string(ms / 1000.0) + "s";
        }
    }

    // ============= AsyncExecutor Implementation =============
    AsyncExecutor::AsyncExecutor(size_t concurrency) :
        executor_(concurrency)
    {
        spdlog::info("{}: Initialized with {} threads (hardware: {})", kLogTag, concurrency, std::thread::hardware_concurrency());
    }

    AsyncExecutor::~AsyncExecutor()
    {
        spdlog::info("{}: Shutting down...", kLogTag);
        LogStatus();
        WaitAll();
        spdlog::info("{}: Destroyed", kLogTag);
    }

    void AsyncExecutor::Submit(const std::string& task_id, TaskFunc task, Callback callback)
    {
        auto submit_start = Clock::now();
        spdlog::debug("[{}] Submitting task", task_id);

        auto context = std::make_shared<TaskContext>();
        context->callback = callback;
        bool replaced_existing = RegisterTask(task_id, context);

        LogTaskStart(task_id, replaced_existing);

        try
        {
            executor_.async([this, task_id, task = std::move(task), context]() {
                ExecuteTask(task_id, std::move(task), context);
            });

            auto duration = GetElapsedTime(submit_start);
            spdlog::debug("[{}] Submitted in {}", task_id, FormatDuration(duration));
        }
        catch (const std::exception& e)
        {
            spdlog::error("[{}] Failed to submit: {}", task_id, e.what());
            UnregisterTask(task_id);
            stats_.IncrementFailed();
            throw;
        }
    }

    bool AsyncExecutor::Cancel(const std::string& task_id)
    {
        bool cancelled = MarkTaskCancelled(task_id);

        if (cancelled)
            spdlog::info("[{}] Cancelled successfully", task_id);
        else
            spdlog::debug("[{}] Cancel requested but task not found", task_id);

        return cancelled;
    }

    void AsyncExecutor::WaitAll()
    {
        spdlog::info("{}: Waiting for all tasks...", kLogTag);

        auto wait_start = Clock::now();
        executor_.wait_for_all();
        auto duration = GetElapsedTime(wait_start);

        spdlog::info("{}: All tasks completed in {}", kLogTag, FormatDuration(duration));
        LogStatus();
    }

    size_t AsyncExecutor::GetRunningTaskCount() const
    {
        std::lock_guard lock(mutex_);
        return running_tasks_.size();
    }

    AsyncExecutor::Statistics AsyncExecutor::GetStatistics() const
    {
        std::lock_guard lock(mutex_);
        return stats_.GetSnapshot(running_tasks_.size());
    }

    void AsyncExecutor::ExecuteTask(const std::string& task_id, TaskFunc task, TaskContextPtr context)
    {
        auto task_start = Clock::now();
        spdlog::info("[{}] Task execution started", task_id);

        try
        {
            // Check if already cancelled
            if (context->cancelled.load())
            {
                spdlog::debug("[{}] Task was pre-cancelled", task_id);
                stats_.IncrementCancelled();
                LogTaskCompletion(task_id, task_start, false);
                UnregisterTask(task_id);
                return;
            }

            // Execute the task
            spdlog::debug("[{}] Executing task function", task_id);
            auto result = task();

            // Check if cancelled during execution
            if (context->cancelled.load())
            {
                spdlog::debug("[{}] Task cancelled during execution", task_id);
                stats_.IncrementCancelled();
            }
            else if (result)
            {
                spdlog::trace("[{}] Task produced {} bytes result", task_id, result->size());
                if (context->callback && result.has_value())
                    context->callback(result.value());
                stats_.IncrementCompleted();
            }
            else
            {
                spdlog::warn("[{}] Task completed with no result", task_id);
                stats_.IncrementCompleted();
            }

            LogTaskCompletion(task_id, task_start, true);
        }
        catch (const std::exception& e)
        {
            spdlog::error("[{}] Task failed: {}", task_id, e.what());
            stats_.IncrementFailed();
            LogTaskCompletion(task_id, task_start, false);
        }
        catch (...)
        {
            spdlog::error("[{}] Task failed with unknown exception", task_id);
            stats_.IncrementFailed();
            LogTaskCompletion(task_id, task_start, false);
        }

        UnregisterTask(task_id);

        // Periodic status logging
        if (stats_.completed.load() % kStatusLogInterval == 0)
            LogStatus();
    }

    bool AsyncExecutor::RegisterTask(const std::string& task_id, TaskContextPtr context)
    {
        std::lock_guard lock(mutex_);

        bool replaced = false;
        auto it = running_tasks_.find(task_id);

        if (it != running_tasks_.end())
        {
            spdlog::warn("[{}] Cancelling existing task", task_id);
            it->second->cancelled.store(true);
            stats_.IncrementCancelled();
            replaced = true;
        }

        running_tasks_[task_id] = context;
        stats_.IncrementSubmitted();

        spdlog::debug("[{}] Registered (running: {}, total: {})", task_id, running_tasks_.size(), stats_.submitted.load());
        return replaced;
    }

    void AsyncExecutor::UnregisterTask(const std::string& task_id)
    {
        std::lock_guard lock(mutex_);

        auto it = running_tasks_.find(task_id);
        if (it != running_tasks_.end())
        {
            auto lifetime = GetElapsedTime(it->second->start_time);
            spdlog::trace("[{}] Task lifetime: {}", task_id, FormatDuration(lifetime));

            running_tasks_.erase(it);
            spdlog::debug("[{}] Unregistered (remaining: {})", task_id, running_tasks_.size());
        }
        else
        {
            spdlog::warn("[{}] Attempted to unregister non-existent task", task_id);
        }
    }

    bool AsyncExecutor::MarkTaskCancelled(const std::string& task_id)
    {
        std::lock_guard lock(mutex_);

        auto it = running_tasks_.find(task_id);
        if (it != running_tasks_.end())
        {
            it->second->cancelled.store(true);
            return true;
        }

        return false;
    }

    void AsyncExecutor::LogTaskStart(const std::string& task_id, bool replaced_existing) const
    {
        auto stats = GetStatistics();

        if (replaced_existing)
            spdlog::info("[{}] Submitted (replaced existing) - Running: {}, Total: {}", task_id, stats.running, stats.submitted);
        else
            spdlog::debug("[{}] Submitted - Running: {}, Total: {}", task_id, stats.running, stats.submitted);
    }

    void AsyncExecutor::LogTaskCompletion(const std::string& task_id, TimePoint start_time, bool success) const
    {
        auto duration = GetElapsedTime(start_time);

        if (success)
            spdlog::info("[{}] Completed in {}", task_id, FormatDuration(duration));
        else
            spdlog::warn("[{}] Failed/cancelled after {}", task_id, FormatDuration(duration));
    }

    std::chrono::milliseconds AsyncExecutor::GetElapsedTime(TimePoint start_time) const
    {
        return std::chrono::duration_cast<std::chrono::milliseconds>(Clock::now() - start_time);
    }

    void AsyncExecutor::LogStatus() const
    {
        // if (spdlog::get_level() > spdlog::level::debug)
        //     return;
        // 一次性获取所有需要的数据
        std::vector<std::pair<std::string, std::pair<TimePoint, bool>>> active_tasks;
        Statistics stats;

        {
            std::lock_guard lock(mutex_);
            stats = stats_.GetSnapshot(running_tasks_.size());

            for (const auto& [id, context] : running_tasks_)
                active_tasks.emplace_back(id, std::make_pair(context->start_time, context->cancelled.load()));
        }

        auto total = stats.submitted;
        std::ostringstream status_output;

        status_output << "\n";
        status_output << "    ╭────────────────────────────────────────────╮\n";
        status_output << "    │        🔔  AsyncExecutor Dashboard         │\n";
        status_output << "    ├────────────────────────────────────────────┤\n";
        status_output << "    │                                            │\n";
        status_output << "    │   📊 SYSTEM METRICS                        │\n";
        status_output << "    │   ──────────────────────────────────────   │\n";
        status_output << fmt::format("    │   📤 Submitted   ▶ {:>6} total            │\n", stats.submitted);
        status_output << fmt::format("    │   🚀 Running     ▶ {:>6} tasks ({:>5.1f}%)   │\n", stats.running, total > 0 ? (stats.running * 100.0 / total) : 0);
        status_output << fmt::format("    │   ✅ Completed   ▶ {:>6} tasks ({:>5.1f}%)   │\n", stats.completed, total > 0 ? (stats.completed * 100.0 / total) : 0);
        status_output << fmt::format("    │   ❌ Failed      ▶ {:>6} tasks ({:>5.1f}%)   │\n", stats.failed, total > 0 ? (stats.failed * 100.0 / total) : 0);
        status_output << fmt::format("    │   🥏 Cancelled   ▶ {:>6} tasks ({:>5.1f}%)   │\n", stats.cancelled, total > 0 ? (stats.cancelled * 100.0 / total) : 0);
        status_output << "    │                                            │\n";
        if (stats.running > 0)
        {
            status_output << "    │   🎯 ACTIVE TASKS                          │\n";
            status_output << "    │   ──────────────────────────────────────   │\n";
            for (const auto& [id, info] : active_tasks)
            {
                auto [start_time, is_cancelled] = info;
                auto runtime = GetElapsedTime(start_time);
                auto status_icon = is_cancelled ? "🔴" : "🟢";
                auto short_id = id.length() > 20 ? id.substr(0, 20) + "…" : id;
                status_output << fmt::format("    │   {} │ {:22} │ ⏱️ {:<8}│\n", status_icon, short_id, FormatDuration(runtime));
            }
        }
        status_output << "    ╰────────────────────────────────────────────╯";
        spdlog::debug("{}", status_output.str());
    }
}