#include "output.hpp"
#include <iostream>
#include <mutex>

namespace dropshell
{
    // Mutex to synchronize output
    std::mutex output_mutex;

    // ANSI colour codes
    constexpr const char *GREY = "\033[90m";
    constexpr const char *RESET = "\033[0m";
    constexpr const char *DEBUG_COLOUR = "\033[36m";   // Cyan
    // constexpr const char *INFO_COLOUR = "\033[32m";    // Green
    constexpr const char *INFO_COLOUR = "\033[37m";    // White
    constexpr const char *WARNING_COLOUR = "\033[33m"; // Yellow
    constexpr const char *ERROR_COLOUR = "\033[31m";   // Red

    // Tag and colour for each stream
    struct StreamInfo
    {
        const char *tag;
        const char *colour;
    };

    const StreamInfo stream_infos[] = {
        {"[DBG]", DEBUG_COLOUR},
        {"[INF]", INFO_COLOUR},
        {"[WRN]", WARNING_COLOUR},
        {"[ERR]", ERROR_COLOUR}};

    // Custom streambuf to prefix and colour each line
    class PrefixStreambuf : public std::streambuf
    {
    public:
        PrefixStreambuf(std::ostream &dest, const char *tag, const char *colour)
            : dest_(dest), tag_(tag), colour_(colour), at_line_start_(true) {}

    protected:
        int overflow(int c) override
        {
            std::lock_guard<std::mutex> lock(output_mutex);
            if (c == EOF)
                return !EOF;
            if (at_line_start_ && c != '\n')
            {
                dest_ << GREY << tag_ << RESET << ' ' << colour_;
                at_line_start_ = false;
            }
            dest_.put(static_cast<char>(c));
            if (c == '\n')
            {
                dest_ << RESET;
                at_line_start_ = true;
            }
            return c;
        }

    private:
        std::ostream &dest_;
        const char *tag_;
        const char *colour_;
        bool at_line_start_;
    };

    PrefixStreambuf debug_buf(std::clog, stream_infos[0].tag, stream_infos[0].colour);
    PrefixStreambuf info_buf(std::clog, stream_infos[1].tag, stream_infos[1].colour);
    PrefixStreambuf warning_buf(std::clog, stream_infos[2].tag, stream_infos[2].colour);
    PrefixStreambuf error_buf(std::cerr, stream_infos[3].tag, stream_infos[3].colour);

    std::ostream debug_stream(&debug_buf);
    std::ostream info_stream(&info_buf);
    std::ostream warning_stream(&warning_buf);
    std::ostream error_stream(&error_buf);

    std::ostream &debug = debug_stream;
    std::ostream &info = info_stream;
    std::ostream &warning = warning_stream;
    std::ostream &error = error_stream;

    std::ostream &rawout = std::cout;
    std::ostream &rawerr = std::cerr;

    std::ostream &colourstream(sColour colour)
    {
        switch (colour)
        {
        case sColour::DEBUG:
            return debug_stream;
        case sColour::INFO:
            return info_stream;
        case sColour::WARNING:
            return warning_stream;
        case sColour::ERROR:
            return error_stream;
        default:
            return info_stream;
        }
    }

    void SetColour(sColour colour, std::ostream &os)
    {
        switch (colour)
        {
        case sColour::RESET:
            os << RESET;
            break;
        case sColour::DEBUG:
            os << DEBUG_COLOUR;
            break;
        case sColour::INFO:
            os << INFO_COLOUR;
            break;
        case sColour::WARNING:
            os << WARNING_COLOUR;
            break;
        case sColour::ERROR:
            os << ERROR_COLOUR;
            break;
        }
    }

    SwitchColour::SwitchColour(sColour colour, std::ostream &os) : os_(os), colour_(colour)
    {
        SetColour(colour_, os_);
    }

    SwitchColour::~SwitchColour()
    {
        SetColour(sColour::RESET, os_);
    }

} // namespace dropshell