#include "sessionmodel.h"
#include "ollie9pclient.h"

SessionModel::SessionModel(Ollie9pClient *client, QObject *parent)
    : QAbstractListModel(parent), m_client(client)
{
    connect(m_client, &Ollie9pClient::sessionsChanged, this, &SessionModel::rebuild);
    connect(m_client, &Ollie9pClient::agentStateChanged, this, &SessionModel::updateAgentState);
    rebuild();
}

int SessionModel::rowCount(const QModelIndex &parent) const
{
    return parent.isValid() ? 0 : m_rows.size();
}

QVariant SessionModel::data(const QModelIndex &index, int role) const
{
    if (!index.isValid() || index.row() < 0 || index.row() >= m_rows.size()) return {};
    const Row &r = m_rows.at(index.row());
    switch (role) {
    case TypeRole: return r.type;
    case IdRole: return r.id;
    case ParentIdRole: return r.parentId;
    case AgentIdRole: return r.agentId;
    case LevelRole: return r.level;
    case ExpandedRole: return r.expanded;
    case ParentExpandedRole: return r.parentExpanded;
    case StateRole: return r.state;
    case ModelRole: return r.model;
    case DisplayNameRole: return r.displayName;
    case BaseNameRole: return r.baseName;
    case ConnectionColorRole: return r.connectionColor;
    case PausedRole: return r.paused;
    default: return {};
    }
}

QHash<int, QByteArray> SessionModel::roleNames() const
{
    return {{TypeRole, "type"}, {IdRole, "id"}, {ParentIdRole, "parentId"},
            {AgentIdRole, "agentId"}, {LevelRole, "level"}, {ExpandedRole, "expanded"},
            {ParentExpandedRole, "parentExpanded"}, {StateRole, "state"}, {ModelRole, "model"},
            {DisplayNameRole, "displayName"}, {BaseNameRole, "baseName"},
            {ConnectionColorRole, "connectionColor"}, {PausedRole, "paused"}};
}

QString SessionModel::connectionColor(const QVariantMap &session)
{
    const bool paused = session.value("paused").toBool();
    const bool connected = session.value("connected").toBool();
    if (paused && connected) return "yellow";
    if (paused) return "gray";
    return connected ? "green" : "red";
}

void SessionModel::appendSession(const QVariantMap &session)
{
    const QString id = session.value("id").toString();
    const QString name = session.value("name").toString().isEmpty() ? id : session.value("name").toString();
    const QVariantList agents = session.value("agents").toList();
    const bool paused = session.value("paused").toBool();
    const bool expanded = m_expanded.value(id, true);
    m_expanded[id] = expanded;

    Row row;
    row.type = "session"; row.id = id; row.expanded = expanded; row.paused = paused;
    row.state = session.value("state").toString(); row.model = session.value("model").toString();
    row.baseName = name;
    row.displayName = agents.isEmpty() ? name : name + " (" + QString::number(agents.size()) + ")";
    row.connectionColor = connectionColor(session);
    m_rows.append(row);

    for (const QVariant &value : agents) {
        const QVariantMap agent = value.toMap();
        Row child;
        child.type = "agent"; child.parentId = id; child.agentId = agent.value("id").toString();
        child.level = 1; child.parentExpanded = expanded;
        child.state = agent.value("state").toString();
        if (child.state.isEmpty()) child.state = "idle";
        child.displayName = agent.value("name").toString(); child.paused = paused;
        m_rows.append(child);
    }
}

void SessionModel::rebuild()
{
    beginResetModel();
    m_rows.clear();
    QSet<QString> present;
    bool haveActive = false;
    bool haveInactive = false;
    for (const QVariant &value : m_client->sessions()) {
        const QVariantMap session = value.toMap();
        const bool paused = session.value("paused").toBool();
        if (paused && !haveInactive) {
            Row header; header.type = "header"; header.displayName = "Inactive"; m_rows.append(header); haveInactive = true;
        } else if (!paused && !haveActive) {
            Row header; header.type = "header"; header.displayName = "Active"; m_rows.append(header); haveActive = true;
        }
        const QString id = session.value("id").toString();
        present.insert(id);
        appendSession(session);
    }
    for (auto it = m_expanded.begin(); it != m_expanded.end();) {
        if (!present.contains(it.key())) it = m_expanded.erase(it); else ++it;
    }
    endResetModel();
    emit countChanged();
}

void SessionModel::updateVisibility()
{
    for (int i = 0; i < m_rows.size(); ++i) {
        if (m_rows[i].type != "session") continue;
        for (int j = i + 1; j < m_rows.size() && m_rows[j].type != "session"; ++j)
            m_rows[j].parentExpanded = m_rows[i].expanded;
    }
    if (!m_rows.isEmpty()) emit dataChanged(index(0), index(m_rows.size() - 1), {ParentExpandedRole, ExpandedRole});
}

void SessionModel::setExpanded(int index, bool expanded)
{
    if (index < 0 || index >= m_rows.size() || m_rows[index].type != "session") return;
    m_rows[index].expanded = expanded;
    m_expanded[m_rows[index].id] = expanded;
    updateVisibility();
}

QVariantMap SessionModel::get(int index) const
{
    QVariantMap result;
    if (index < 0 || index >= m_rows.size()) return result;
    const Row &r = m_rows.at(index);
    result["type"] = r.type; result["id"] = r.id; result["parentId"] = r.parentId;
    result["agentId"] = r.agentId; result["level"] = r.level; result["expanded"] = r.expanded;
    result["parentExpanded"] = r.parentExpanded; result["state"] = r.state;
    result["model"] = r.model; result["displayName"] = r.displayName;
    result["baseName"] = r.baseName; result["connectionColor"] = r.connectionColor;
    result["paused"] = r.paused;
    return result;
}

QString SessionModel::firstAgentId(const QString &sessionId) const
{
    for (const Row &row : m_rows)
        if (row.type == "agent" && row.parentId == sessionId) return row.agentId;
    return {};
}

void SessionModel::updateAgentState(const QString &sessionId, const QString &agentId, const QString &state)
{
    for (int i = 0; i < m_rows.size(); ++i) {
        if (m_rows[i].type == "agent" && m_rows[i].parentId == sessionId && m_rows[i].agentId == agentId) {
            m_rows[i].state = state;
            emit dataChanged(index(i), index(i), {StateRole});
            return;
        }
    }
}