#ifndef LLSH_MESSAGE_H
#define LLSH_MESSAGE_H

#include <nlohmann/json.hpp>

#include <string>
#include <variant>
#include <vector>

/* ------------------------------------------------------------------ */
/* Message model (mirrors the rig message types llsh uses)            */
/*                                                                     */
/* message         = system_message    { text }                       */
/*                   | user_message    { parts: [user_part] }         */
/*                   | assistant_message { parts: [assistant_part] }  */
/* user_part       = user_text_part        { text }                   */
/*                   | user_tool_result_part { call, name,           */
/*                                               content: [trc] }     */
/* trc             = trc_text { text } | trc_json { value } | trc_image */
/* assistant_part  = asst_text_part      { text }                     */
/*                   | asst_tool_call_part { name, args }             */
/*                   | asst_reasoning_part { text }                   */
/*                   | asst_image_part                                      */
/*                                                                     */
/* JSON shape (recursive, nlohmann ADL below): each type serializes   */
/* to exactly this wire format and the `role`/`type` discriminators   */
/* keep old session files loadable.                                    */
/* ------------------------------------------------------------------ */

/* -- tool result content -- */

struct trc_text {
    std::string text;
};
struct trc_json {
    std::string value;
};
struct trc_image {
};
using tool_result_content = std::variant<trc_text, trc_json, trc_image>;

/* -- user parts -- */

struct user_text_part {
    std::string text;
};
struct user_tool_result_part {
    std::string call; /* tool-call id */
    std::string name; /* tool name */
    std::vector<tool_result_content> content;
};
using user_part = std::variant<user_text_part, user_tool_result_part>;

/* -- assistant parts -- */

struct asst_text_part {
    std::string text;
};
struct asst_tool_call_part {
    std::string name;
    std::string args; /* stringified arguments */
};
struct asst_reasoning_part {
    std::string text;
};
struct asst_image_part {
};
using assistant_part =
    std::variant<asst_text_part, asst_tool_call_part, asst_reasoning_part,
                 asst_image_part>;

/* -- messages -- */

struct system_message {
    std::string text;
};
struct user_message {
    std::vector<user_part> parts;
};
struct assistant_message {
    std::vector<assistant_part> parts;
};

struct message {
    std::variant<system_message, user_message, assistant_message> body;

    bool is_system() const {
        return std::holds_alternative<system_message>(body);
    }
    bool is_user() const {
        return std::holds_alternative<user_message>(body);
    }
    bool is_assistant() const {
        return std::holds_alternative<assistant_message>(body);
    }

    system_message &system() { return std::get<system_message>(body); }
    const system_message &system() const {
        return std::get<system_message>(body);
    }
    user_message &user() { return std::get<user_message>(body); }
    const user_message &user() const {
        return std::get<user_message>(body);
    }
    assistant_message &assistant() {
        return std::get<assistant_message>(body);
    }
    const assistant_message &assistant() const {
        return std::get<assistant_message>(body);
    }
};

/* Constructors (fresh message, content copied). */
inline message msg_system(const char *text) {
    system_message m;
    m.text = text ? text : "";
    return message{std::move(m)};
}
inline message msg_user_text(const char *text) {
    user_message m;
    user_text_part p;
    p.text = text ? text : "";
    m.parts.push_back(std::move(p));
    return message{std::move(m)};
}
inline message msg_user_tool_result(const char *call, const char *name,
                                    const char *text) {
    user_message m;
    user_tool_result_part p;
    p.call = call ? call : "";
    p.name = name ? name : "";
    trc_text t;
    t.text = text ? text : "";
    p.content.push_back(std::move(t));
    m.parts.push_back(std::move(p));
    return message{std::move(m)};
}
inline message msg_assistant_text(const char *text) {
    assistant_message m;
    asst_text_part p;
    p.text = text ? text : "";
    m.parts.push_back(std::move(p));
    return message{std::move(m)};
}
inline message msg_assistant_tool_call(const char *name, const char *args) {
    assistant_message m;
    asst_tool_call_part p;
    p.name = name ? name : "";
    p.args = args ? args : "";
    m.parts.push_back(std::move(p));
    return message{std::move(m)};
}

/* Append parts to an existing message. */
inline message &msg_push_user_text(message &m, const char *text) {
    if (!m.is_user())
        m.body = user_message{};
    user_text_part p;
    p.text = text ? text : "";
    m.user().parts.push_back(std::move(p));
    return m;
}
inline message &msg_push_user_tool_result(message &m, const char *call,
                                          const char *name, const char *text) {
    if (!m.is_user())
        m.body = user_message{};
    user_tool_result_part p;
    p.call = call ? call : "";
    p.name = name ? name : "";
    trc_text t;
    t.text = text ? text : "";
    p.content.push_back(std::move(t));
    m.user().parts.push_back(std::move(p));
    return m;
}
inline message &msg_push_assistant_text(message &m, const char *text) {
    if (!m.is_assistant())
        m.body = assistant_message{};
    asst_text_part p;
    p.text = text ? text : "";
    m.assistant().parts.push_back(std::move(p));
    return m;
}
inline message &msg_push_assistant_tool_call(message &m, const char *name,
                                             const char *args) {
    if (!m.is_assistant())
        m.body = assistant_message{};
    asst_tool_call_part p;
    p.name = name ? name : "";
    p.args = args ? args : "";
    m.assistant().parts.push_back(std::move(p));
    return m;
}
inline message &msg_push_assistant_reasoning(message &m, const char *text) {
    if (!m.is_assistant())
        m.body = assistant_message{};
    asst_reasoning_part p;
    p.text = text ? text : "";
    m.assistant().parts.push_back(std::move(p));
    return m;
}

/* ------------------------------------------------------------------ */
/* History (growable message list)                                     */
/* ------------------------------------------------------------------ */

struct history {
    std::vector<message> msgs;
};

/* Destroy the stored messages and the vector buffer; the `history` object
 * (wherever it lives) must not be used afterwards. */
inline void history_free(history *h) { h->msgs.~vector<message>(); }
inline void history_push(history *h, message m) {
    h->msgs.push_back(std::move(m));
}
/* Remove and return element i. */
inline message history_take(history *h, size_t i) {
    message out = std::move(h->msgs[i]);
    h->msgs.erase(h->msgs.begin() + static_cast<long>(i));
    return out;
}
inline size_t history_len(const history *h) { return h->msgs.size(); }
inline const message *history_at(const history *h, size_t i) {
    return &h->msgs[i];
}

/* ------------------------------------------------------------------ */
/* Recursive JSON (nlohmann ADL: to_json / from_json)                 */
/* ------------------------------------------------------------------ */

/* Tolerant string read: missing / null / non-string values become "". */
inline std::string jstring(const nlohmann::json &j, const char *key) {
    if (j.contains(key) && j.at(key).is_string())
        return j.at(key).get<std::string>();
    return std::string();
}

inline void to_json(nlohmann::json &j, const tool_result_content &v) {
    if (const auto *t = std::get_if<trc_text>(&v))
        j = nlohmann::json{{"text", t->text}};
    else if (const auto *s = std::get_if<trc_json>(&v))
        j = nlohmann::json{{"json", s->value}};
    else
        j = nlohmann::json{{"image", "1"}};
}
inline void from_json(const nlohmann::json &j, tool_result_content &v) {
    if (j.contains("text")) {
        trc_text t;
        t.text = jstring(j, "text");
        v = std::move(t);
    } else if (j.contains("json")) {
        trc_json s;
        s.value = jstring(j, "json");
        v = std::move(s);
    } else {
        v = trc_image{};
    }
}

inline void to_json(nlohmann::json &j, const user_part &p) {
    if (const auto *t = std::get_if<user_text_part>(&p))
        j = nlohmann::json{{"type", "text"}, {"text", t->text}};
    else {
        const auto &r = std::get<user_tool_result_part>(p);
        nlohmann::json arr = nlohmann::json::array();
        for (const auto &c : r.content)
            arr.push_back(c);
        j = nlohmann::json{{"type", "tool_result"},
                           {"call", r.call},
                           {"name", r.name},
                           {"content", std::move(arr)}};
    }
}
inline void from_json(const nlohmann::json &j, user_part &p) {
    const std::string type = jstring(j, "type");
    if (type == "text") {
        user_text_part t;
        t.text = jstring(j, "text");
        p = std::move(t);
    } else if (type == "tool_result") {
        user_tool_result_part r;
        r.call = jstring(j, "call");
        r.name = jstring(j, "name");
        if (j.contains("content") && j.at("content").is_array())
            for (const auto &c : j.at("content"))
                r.content.push_back(c.get<tool_result_content>());
        p = std::move(r);
    } else {
        user_text_part t;
        t.text = jstring(j, "text");
        p = std::move(t);
    }
}

inline void to_json(nlohmann::json &j, const assistant_part &p) {
    if (const auto *t = std::get_if<asst_text_part>(&p))
        j = nlohmann::json{{"type", "text"}, {"text", t->text}};
    else if (const auto *c = std::get_if<asst_tool_call_part>(&p))
        j = nlohmann::json{{"type", "tool_call"},
                           {"name", c->name},
                           {"args", c->args}};
    else if (const auto *r = std::get_if<asst_reasoning_part>(&p))
        j = nlohmann::json{{"type", "reasoning"}, {"text", r->text}};
    else
        j = nlohmann::json::object();
}
inline void from_json(const nlohmann::json &j, assistant_part &p) {
    const std::string type = jstring(j, "type");
    if (type == "text") {
        asst_text_part t;
        t.text = jstring(j, "text");
        p = std::move(t);
    } else if (type == "tool_call") {
        asst_tool_call_part c;
        c.name = jstring(j, "name");
        c.args = jstring(j, "args");
        p = std::move(c);
    } else if (type == "reasoning") {
        asst_reasoning_part r;
        r.text = jstring(j, "text");
        p = std::move(r);
    } else {
        p = asst_image_part{};
    }
}

inline void to_json(nlohmann::json &j, const message &m) {
    if (m.is_system())
        j = nlohmann::json{{"role", "system"}, {"text", m.system().text}};
    else if (m.is_user()) {
        nlohmann::json parts = nlohmann::json::array();
        for (const auto &p : m.user().parts)
            parts.push_back(p);
        j = nlohmann::json{{"role", "user"}, {"parts", std::move(parts)}};
    } else {
        nlohmann::json parts = nlohmann::json::array();
        for (const auto &p : m.assistant().parts)
            parts.push_back(p);
        j =
            nlohmann::json{{"role", "assistant"}, {"parts", std::move(parts)}};
    }
}
inline void from_json(const nlohmann::json &j, message &m) {
    const std::string role = jstring(j, "role");
    if (role == "system") {
        system_message s;
        s.text = jstring(j, "text");
        m.body = std::move(s);
    } else if (role == "assistant") {
        assistant_message a;
        if (j.contains("parts") && j.at("parts").is_array())
            for (const auto &p : j.at("parts"))
                if (p.is_object())
                    a.parts.push_back(p.get<assistant_part>());
        m.body = std::move(a);
    } else {
        user_message u;
        if (j.contains("parts") && j.at("parts").is_array())
            for (const auto &p : j.at("parts"))
                if (p.is_object())
                    u.parts.push_back(p.get<user_part>());
        m.body = std::move(u);
    }
}

#endif