#include "chatstore.hpp"

#include "attachment.hpp"
#include "llmclient.hpp"
#include "markdownblock.hpp"
#include "message.hpp"
#include "segment.hpp"

#include <QCoreApplication>
#include <QDateTime>
#include <QDebug>
#include <QDir>
#include <QFileInfo>
#include <QPointer>
#include <QStandardPaths>
#include <QSqlError>
#include <QSqlQuery>
#include <QThreadPool>
#include <QVector>
#include <QUuid>

namespace ZShell::llm {

namespace {

QString segmentTypeName(LlmSegment::Type type) {
	switch (type) {
	case LlmSegment::Type::Reasoning:
		return QStringLiteral("reasoning");
	case LlmSegment::Type::ToolCall:
		return QStringLiteral("tool_call");
	case LlmSegment::Type::Content:
		return QStringLiteral("content");
	}
	return QStringLiteral("reasoning");
}

LlmSegment::Type segmentTypeFromName(const QString& name) {
	if (name == QLatin1String("tool_call")) return LlmSegment::Type::ToolCall;
	if (name == QLatin1String("content")) return LlmSegment::Type::Content;
	return LlmSegment::Type::Reasoning;
}

QString sqlText(const QString& value) {
	if (value.isNull()) return QStringLiteral("");
	return value;
}

struct SegmentRow {
	QString type;
	QString text;
	QString name;
	QString toolCallId;
	QString arguments;
	QString result;
	int status = 0;
	qint64 elapsedMs = 0;
	qint64 timestamp = 0;
};

struct GenerationRow {
	qint64 timestamp = 0;
	bool active = false;
	QVector<SegmentRow> segments;
};

struct AttachmentRow {
	QString name;
	QString path;
	int kind = 1;
	QString mime;
	qint64 size = 0;
};

struct MessageRow {
	bool user = false;
	qint64 timestamp = 0;
	QString content;
	QVector<GenerationRow> generations;
	QVector<AttachmentRow> attachments;
};

} // namespace

ChatStore::ChatStore(LlmClient* client, QObject* parent)
	: QAbstractListModel(parent)
	, m_client(client)
	, m_connectionName(QUuid::createUuid().toString()) {
	openDb();
	load();
}

ChatStore::~ChatStore() {
	if (m_connectionName.isEmpty()) return;
	QSqlDatabase handle = QSqlDatabase::database(m_connectionName, false);
	handle.close();
	QSqlDatabase::removeDatabase(m_connectionName);
}

QSqlDatabase ChatStore::db() const {
	return QSqlDatabase::database(m_connectionName);
}

void ChatStore::openDb() {
	m_dbPath =
		QStandardPaths::writableLocation(QStandardPaths::GenericStateLocation) +
		QStringLiteral("/zshell/chats.sqlite");
	QDir().mkpath(QFileInfo(m_dbPath).absolutePath());

	QSqlDatabase handle =
		QSqlDatabase::addDatabase(QStringLiteral("QSQLITE"), m_connectionName);
	handle.setDatabaseName(m_dbPath);
	if (!handle.open()) {
		qWarning() << "ChatStore: failed to open database" << m_dbPath << ":"
				   << handle.lastError().text();
		return;
	}
	{
		QSqlQuery pragma(handle);
		pragma.exec(QStringLiteral("PRAGMA foreign_keys = ON"));
	}
	{
		QSqlQuery query(handle);
		query.exec(QStringLiteral(
			"CREATE TABLE IF NOT EXISTS sessions (\n"
			"  id TEXT PRIMARY KEY,\n"
			"  title TEXT NOT NULL DEFAULT '',\n"
			"  icon TEXT NOT NULL DEFAULT '',\n"
			"  created_at INTEGER NOT NULL,\n"
			"  updated_at INTEGER NOT NULL,\n"
			"  pinned INTEGER NOT NULL DEFAULT 0\n"
			")"));
	}
	{
		QSqlQuery query(handle);
		query.exec(QStringLiteral(
			"CREATE TABLE IF NOT EXISTS messages (\n"
			"  id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
			"  session_id TEXT NOT NULL REFERENCES sessions (id) "
			"ON DELETE CASCADE,\n"
			"  role TEXT NOT NULL,\n"
			"  content TEXT NOT NULL DEFAULT '',\n"
			"  timestamp INTEGER NOT NULL\n"
			")"));
		if (!query.exec(QStringLiteral(
				"ALTER TABLE messages ADD COLUMN content TEXT "
				"NOT NULL DEFAULT ''"))) {
		}
	}
	{
		QSqlQuery query(handle);
		query.exec(QStringLiteral(
			"CREATE TABLE IF NOT EXISTS generations (\n"
			"  id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
			"  message_id INTEGER NOT NULL REFERENCES messages "
			"(id) ON DELETE CASCADE,\n"
			"  timestamp INTEGER NOT NULL,\n"
			"  is_active INTEGER NOT NULL DEFAULT 1\n"
			")"));
	}
	{
		QSqlQuery query(handle);
		query.exec(QStringLiteral(
			"CREATE TABLE IF NOT EXISTS segments (\n"
			"  id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
			"  generation_id INTEGER NOT NULL REFERENCES generations "
			"(id) ON DELETE CASCADE,\n"
			"  type TEXT NOT NULL,\n"
			"  text TEXT NOT NULL DEFAULT '',\n"
			"  name TEXT NOT NULL DEFAULT '',\n"
			"  tool_call_id TEXT NOT NULL DEFAULT '',\n"
			"  arguments TEXT NOT NULL DEFAULT '',\n"
			"  result TEXT NOT NULL DEFAULT '',\n"
			"  status INTEGER NOT NULL DEFAULT 0,\n"
			"  elapsed_ms INTEGER NOT NULL DEFAULT 0,\n"
			"  timestamp INTEGER NOT NULL\n"
			")"));
	}
	{
		QSqlQuery query(handle);
		query.exec(QStringLiteral(
			"CREATE TABLE IF NOT EXISTS attachments (\n"
			"  id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
			"  message_id INTEGER NOT NULL REFERENCES messages "
			"(id) ON DELETE CASCADE,\n"
			"  name TEXT NOT NULL,\n"
			"  path TEXT NOT NULL,\n"
			"  kind INTEGER NOT NULL DEFAULT 1,\n"
			"  mime TEXT NOT NULL DEFAULT '',\n"
			"  size INTEGER NOT NULL DEFAULT 0\n"
			")"));
	}
	{
		QSqlQuery query(handle);
		query.exec(QStringLiteral(
			"CREATE INDEX IF NOT EXISTS idx_messages_session "
			"ON messages (session_id)"));
		query.exec(QStringLiteral(
			"CREATE INDEX IF NOT EXISTS idx_generations_message "
			"ON generations (message_id)"));
		query.exec(QStringLiteral(
			"CREATE INDEX IF NOT EXISTS idx_segments_generation "
			"ON segments (generation_id)"));
		query.exec(QStringLiteral(
			"CREATE INDEX IF NOT EXISTS idx_attachments_message "
			"ON attachments (message_id)"));
	}
}

int ChatStore::rowCount(const QModelIndex& parent) const {
	if (parent.isValid()) return 0;
	return static_cast<int>(m_sessions.size());
}

QVariant ChatStore::data(const QModelIndex& index, int role) const {
	if (role != Qt::UserRole || !index.isValid() || index.row() < 0 ||
		index.row() >= m_sessions.size())
		return {};
	return QVariant::fromValue(m_sessions.at(index.row()));
}

QHash<int, QByteArray> ChatStore::roleNames() const {
	return {{Qt::UserRole, "modelData"}};
}

int ChatStore::rowOf(const ChatSession* session) const {
	for (int i = 0; i < m_sessions.size(); ++i)
		if (m_sessions.at(i) == session) return i;
	return -1;
}

bool ChatStore::firstAfter(const ChatSession* a, const ChatSession* b) {
	if (a->pinned() != b->pinned()) return a->pinned();
	return a->updatedAtMs() > b->updatedAtMs();
}

int ChatStore::sortedRow(const ChatSession* session) const {
	int row = 0;
	while (row < m_sessions.size() && !firstAfter(session, m_sessions.at(row)))
		++row;
	return row;
}

void ChatStore::insertRow(ChatSession* session) {
	const int row = sortedRow(session);
	beginInsertRows(QModelIndex(), row, row);
	m_sessions.insert(row, session);
	endInsertRows();
}

void ChatStore::moveRow(int from, int to) {
	if (from < 0 || from >= m_sessions.size() || to < 0 ||
		to > m_sessions.size() || from == to || to == from + 1)
		return;
	if (!beginMoveRows(QModelIndex(), from, from, QModelIndex(), to)) return;
	m_sessions.move(from, to > from ? to - 1 : to);
	endMoveRows();
}

void ChatStore::reposition(ChatSession* session) {
	if (m_loading) return;
	const int row = rowOf(session);
	if (row < 0) return;
	const int target = sortedRow(session);
	if (target == row) return;
	moveRow(row, target);
}

ChatSession* ChatStore::draft() {
	if (m_draft) return m_draft;
	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	m_draft = new ChatSession(this, QUuid::createUuid().toString(), this);
	m_draft->hydrate(QString(), QString(), now, now, 0, false);
	return m_draft;
}

void ChatStore::commitDraft(ChatSession* session) {
	if (!session || session->isCommitted()) return;
	if (m_draft == session) m_draft = nullptr;

	if (!writeSessionRow(session)) return;
	session->setCommitted(true);
	insertRow(session);
}

bool ChatStore::writeSessionRow(ChatSession* session) {
	QSqlQuery query(db());
	query.prepare(
		"INSERT INTO sessions (id, title, icon, created_at, updated_at, "
		"pinned) VALUES (:id, :title, :icon, :created_at, :updated_at, "
		":pinned)");
	query.bindValue(":id", session->id());
	query.bindValue(":title", sqlText(session->title()));
	query.bindValue(":icon", sqlText(session->icon()));
	query.bindValue(":created_at", session->createdAtMs());
	query.bindValue(":updated_at", session->updatedAtMs());
	query.bindValue(":pinned", session->pinned() ? 1 : 0);
	if (!query.exec()) {
		qWarning() << "ChatStore: failed to insert session" << session->id()
				   << ":" << query.lastError().text();
		return false;
	}
	return true;
}

void ChatStore::removeRow(ChatSession* session) {
	if (!session) return;
	const int row = rowOf(session);

	Q_EMIT sessionRemoved(session);

	{
		QSqlQuery query(db());
		query.prepare("DELETE FROM sessions WHERE id = :id");
		query.bindValue(":id", session->id());
		query.exec();
	}

	if (row >= 0) {
		beginRemoveRows(QModelIndex(), row, row);
		m_sessions.removeAt(row);
		endRemoveRows();
	}
	session->deleteLater();
}

void ChatStore::remove(ChatSession* chat) {
	if (!chat) return;
	m_pendingPersists.remove(chat);
	m_pendingClears.remove(chat);
	if (m_draft == chat) m_draft = nullptr;

	if (!chat->isCommitted()) {
		Q_EMIT sessionRemoved(chat);
		chat->deleteLater();
		return;
	}

	removeRow(chat);
}

void ChatStore::clear() {
	const QList<ChatSession*> sessions = m_sessions;
	for (ChatSession* session : sessions)
		remove(session);
}

void ChatStore::requestClear(ChatSession* session) {
	if (!session) return;
	m_pendingClears.insert(session);
	if (session->isLoaded() && m_client->streamingSession() != session)
		applyPendingClear(session);
}

void ChatStore::applyPendingClear(ChatSession* session) {
	if (!session || !m_pendingClears.remove(session)) return;
	if (!session->isLoaded()) return;
	session->clearMessages();
	persist(session);
}

void ChatStore::persist(ChatSession* session) {
	if (!session || !m_sessions.contains(session)) return;
	if (!session->isLoaded()) {
		m_pendingPersists.insert(session);
		return;
	}
	session->setUpdatedAt(QDateTime::currentMSecsSinceEpoch());
	if (!saveSession(session)) return;
	reposition(session);
}

void ChatStore::saveFlags(ChatSession* session) {
	if (!session || !m_sessions.contains(session)) return;
	QSqlQuery query(db());
	query.prepare(
		"UPDATE sessions SET title = :title, icon = :icon, "
		"pinned = :pinned WHERE id = :id");
	query.bindValue(":title", sqlText(session->title()));
	query.bindValue(":icon", sqlText(session->icon()));
	query.bindValue(":pinned", session->pinned() ? 1 : 0);
	query.bindValue(":id", session->id());
	if (!query.exec())
		qWarning() << "ChatStore: failed to save flags for" << session->id()
				   << ":" << query.lastError().text();
	reposition(session);
}

bool ChatStore::saveSession(ChatSession* session) {
	const QString id = session->id();
	QSqlDatabase handle = db();
	if (!handle.transaction()) {
		qWarning() << "ChatStore: failed to begin transaction:"
				   << handle.lastError().text();
		return false;
	}
	bool ok = true;
	{
		QSqlQuery query(handle);
		query.prepare(
			"UPDATE sessions SET title = :title, icon = :icon, "
			"updated_at = :updated_at, pinned = :pinned WHERE id = :id");
		query.bindValue(":title", sqlText(session->title()));
		query.bindValue(":icon", sqlText(session->icon()));
		query.bindValue(":updated_at", session->updatedAtMs());
		query.bindValue(":pinned", session->pinned() ? 1 : 0);
		query.bindValue(":id", session->id());
		ok = query.exec();
	}
	if (ok) {
		QSqlQuery query(handle);
		query.prepare("DELETE FROM messages WHERE session_id = :id");
		query.bindValue(":id", session->id());
		ok = query.exec();
	}
	if (ok) {
		QSqlQuery messageInsert(handle);
		ok = messageInsert.prepare(
			"INSERT INTO messages (session_id, role, content, timestamp) "
			"VALUES (:id, :role, :content, :timestamp)");
		QSqlQuery generationInsert(handle);
		ok = ok &&
			 generationInsert.prepare(
				 "INSERT INTO generations (message_id, timestamp, is_active) "
				 "VALUES (:mid, :timestamp, :is_active)");
		QSqlQuery segmentInsert(handle);
		ok = ok &&
			 segmentInsert.prepare(
				 "INSERT INTO segments (generation_id, type, text, name, "
				 "tool_call_id, arguments, result, status, elapsed_ms, "
				 "timestamp) VALUES (:gid, :type, :text, :name, "
				 ":tool_call_id, :arguments, :result, :status, :elapsed_ms, "
				 ":timestamp)");
		QSqlQuery attachmentInsert(handle);
		ok = ok && attachmentInsert.prepare(
					   "INSERT INTO attachments (message_id, name, path, kind, "
					   "mime, size) VALUES (:mid, :name, :path, :kind, :mime, "
					   ":size)");
		const auto* model = session->messagesModel();
		for (int row = model->rowCount() - 1; ok && row >= 0; --row) {
			const auto* message = model->at(row);
			messageInsert.bindValue(":id", session->id());
			messageInsert.bindValue(
				":role",
				message->role() == ChatMessage::Role::User
					? QStringLiteral("user")
					: QStringLiteral("assistant"));
			messageInsert.bindValue(":content", sqlText(message->content()));
			messageInsert.bindValue(":timestamp", message->timestamp());
			if (!messageInsert.exec()) {
				ok = false;
				qWarning() << "ChatStore: saveSession" << id
						   << "message insert failed:"
						   << messageInsert.lastError().text();
				break;
			}
			const int messageId = messageInsert.lastInsertId().toInt();
			for (int i = 0; ok && i < message->generationCount(); ++i) {
				const auto* generation = message->generation(i);
				generationInsert.bindValue(":mid", messageId);
				generationInsert.bindValue(
					":timestamp", generation->timestamp());
				generationInsert.bindValue(
					":is_active",
					i == message->activeGenerationIndex() ? 1 : 0);
				if (!generationInsert.exec()) {
					ok = false;
					qWarning() << "ChatStore: saveSession" << id
							   << "generation insert failed:"
							   << generationInsert.lastError().text();
					break;
				}
				const int generationId =
					generationInsert.lastInsertId().toInt();
				for (const auto* segment : generation->segments()) {
					segmentInsert.bindValue(":gid", generationId);
					segmentInsert.bindValue(
						":type", segmentTypeName(segment->type()));
					const auto* text =
						qobject_cast<const LlmTextSegment*>(segment);
					segmentInsert.bindValue(
						":text", sqlText(text ? text->text() : QString()));
					const auto* call =
						qobject_cast<const LlmToolCall*>(segment);
					segmentInsert.bindValue(
						":name", sqlText(call ? call->name() : QString()));
					segmentInsert.bindValue(
						":tool_call_id",
						sqlText(call ? call->toolCallId() : QString()));
					segmentInsert.bindValue(
						":arguments",
						sqlText(call ? call->arguments() : QString()));
					segmentInsert.bindValue(
						":result", sqlText(call ? call->result() : QString()));
					segmentInsert.bindValue(
						":status", call ? static_cast<int>(call->status()) : 0);
					segmentInsert.bindValue(":elapsed_ms", segment->elapsedMs());
					segmentInsert.bindValue(":timestamp", segment->timestamp());
					if (!segmentInsert.exec()) {
						ok = false;
						qWarning() << "ChatStore: saveSession" << id
								   << "segment insert failed:"
								   << segmentInsert.lastError().text();
						break;
					}
				}
			}
			if (ok)
				for (const auto* attachment : message->attachments()) {
					attachmentInsert.bindValue(":mid", messageId);
					attachmentInsert.bindValue(
						":name", sqlText(attachment->name()));
					attachmentInsert.bindValue(
						":path", sqlText(attachment->path()));
					attachmentInsert.bindValue(
						":kind", static_cast<int>(attachment->kind()));
					attachmentInsert.bindValue(
						":mime", sqlText(attachment->mimeType()));
					attachmentInsert.bindValue(":size", attachment->size());
					if (!attachmentInsert.exec()) {
						ok = false;
						qWarning() << "ChatStore: saveSession" << id
								   << "attachment insert failed:"
								   << attachmentInsert.lastError().text();
						break;
					}
				}
		}
	}
	if (!ok || !handle.commit()) {
		qWarning() << "ChatStore: saveSession" << id
				   << "commit failed, rolling back";
		handle.rollback();
		qWarning() << "ChatStore: failed to save session" << session->id()
				   << ":" << handle.lastError().text();
		return false;
	}
	return true;
}

void ChatStore::loadMessagesInto(ChatSession* session) {
	if (!session) return;
	const QString sessionId = session->id();
	const QString path = m_dbPath;

	QThreadPool::globalInstance()->start([store = QPointer<ChatStore>(this),
										  session =
											  QPointer<ChatSession>(session),
										  sessionId,
										  path]() {
		QList<MessageRow> rows;
		const QString connName = QUuid::createUuid().toString();
		{
			QSqlDatabase handle =
				QSqlDatabase::addDatabase(QStringLiteral("QSQLITE"), connName);
			handle.setDatabaseName(path);
			if (handle.open()) {
				QSqlQuery busy(handle);
				busy.exec(QStringLiteral("PRAGMA busy_timeout = 3000"));
				QSqlQuery query(handle);
				query.prepare(
					"SELECT id, role, content, timestamp FROM messages "
					"WHERE session_id = :id ORDER BY rowid DESC");
				query.bindValue(":id", sessionId);
				if (!query.exec()) {
					qWarning() << "ChatStore: failed to load messages for"
							   << sessionId << ":" << query.lastError().text();
				} else {
					while (query.next()) {
						const int messageId = query.value(0).toInt();
						MessageRow message;
						message.user = query.value(1).toString() ==
									   QLatin1String("user");
						message.content = query.value(2).toString();
						message.timestamp = query.value(3).toLongLong();
						QSqlQuery generationQuery(handle);
						generationQuery.prepare(
							"SELECT id, timestamp, is_active FROM "
							"generations WHERE message_id = :mid "
							"ORDER BY rowid");
						generationQuery.bindValue(":mid", messageId);
						if (generationQuery.exec()) {
							while (generationQuery.next()) {
								GenerationRow generation;
								generation.timestamp =
									generationQuery.value(1).toLongLong();
								generation.active =
									generationQuery.value(2).toInt() != 0;
								QSqlQuery segmentQuery(handle);
								segmentQuery.prepare(
									"SELECT type, text, name, tool_call_id, "
									"arguments, result, status, elapsed_ms, "
									"timestamp FROM segments WHERE "
									"generation_id = :gid ORDER BY rowid");
								segmentQuery.bindValue(
									":gid", generationQuery.value(0).toInt());
								if (segmentQuery.exec()) {
									while (segmentQuery.next()) {
										SegmentRow segment;
										segment.type =
											segmentQuery.value(0).toString();
										segment.text =
											segmentQuery.value(1).toString();
										segment.name =
											segmentQuery.value(2).toString();
										segment.toolCallId =
											segmentQuery.value(3).toString();
										segment.arguments =
											segmentQuery.value(4).toString();
										segment.result =
											segmentQuery.value(5).toString();
										segment.status =
											segmentQuery.value(6).toInt();
										segment.elapsedMs =
											segmentQuery.value(7).toLongLong();
										segment.timestamp =
											segmentQuery.value(8).toLongLong();
										generation.segments.append(segment);
									}
								} else {
									qWarning()
										<< "ChatStore: failed to load "
										   "segments for generation"
										<< generationQuery.value(0).toInt()
										<< ":"
										<< segmentQuery.lastError().text();
								}
								message.generations.append(generation);
							}
						} else {
							qWarning()
								<< "ChatStore: failed to load generations "
								   "for message"
								<< messageId << ":"
								<< generationQuery.lastError().text();
						}
						QSqlQuery attachmentQuery(handle);
						attachmentQuery.prepare(
							"SELECT name, path, kind, mime, size FROM "
							"attachments WHERE message_id = :mid "
							"ORDER BY rowid");
						attachmentQuery.bindValue(":mid", messageId);
						if (attachmentQuery.exec()) {
							while (attachmentQuery.next()) {
								AttachmentRow attachment;
								attachment.name =
									attachmentQuery.value(0).toString();
								attachment.path =
									attachmentQuery.value(1).toString();
								attachment.kind =
									attachmentQuery.value(2).toInt();
								attachment.mime =
									attachmentQuery.value(3).toString();
								attachment.size =
									attachmentQuery.value(4).toLongLong();
								message.attachments.append(attachment);
							}
						} else {
							qWarning()
								<< "ChatStore: failed to load attachments "
								   "for message"
								<< messageId << ":"
								<< attachmentQuery.lastError().text();
						}
						rows.append(message);
					}
				}
				handle.close();
			} else {
				qWarning() << "ChatStore: failed to open database for load:"
						   << handle.lastError().text();
			}
		}
		QSqlDatabase::removeDatabase(connName);

		QMetaObject::invokeMethod(
			QCoreApplication::instance(),
			[store, session, rows = std::move(rows)]() mutable {
				ChatStore* st = store;
				ChatSession* s = session;
				if (!st || !s) return;

				auto* model = s->messagesModel();
				if (!model) return;
				QList<ChatMessage*> messages;
				for (const MessageRow& row : rows) {
					auto* message = model->createMessage(
						row.user ? ChatMessage::Role::User
								 : ChatMessage::Role::Assistant,
						row.timestamp);

					if (!row.attachments.isEmpty()) {
						QList<Attachment*> attachments;
						for (const AttachmentRow& attachmentRow :
							 row.attachments) {
							attachments.append(new Attachment(
								attachmentRow.name,
								attachmentRow.path,
								attachmentRow.kind == 0
									? Attachment::Kind::Image
									: Attachment::Kind::File,
								attachmentRow.mime,
								attachmentRow.size,
								message));
						}
						message->setAttachments(attachments);
					}

					if (row.user) {
						QString content = row.content;
						if (content.isEmpty() && !row.generations.isEmpty()) {
							for (const SegmentRow& segment :
								 row.generations.first().segments)
								if (segmentTypeFromName(segment.type) ==
									LlmSegment::Type::Content) {
									content = segment.text;
									break;
								}
						}
						if (!content.isEmpty()) message->setContent(content);
						messages.append(message);
						continue;
					}

					int activeIndex = 0;
					for (int i = 0; i < row.generations.size(); ++i) {
						const GenerationRow& generationRow =
							row.generations.at(i);
						auto* generation =
							message->addGeneration(generationRow.timestamp);
						if (!generation) continue;
						for (const SegmentRow& segmentRow :
							 generationRow.segments) {
							const LlmSegment::Type type =
								segmentTypeFromName(segmentRow.type);
							LlmSegment* segment = nullptr;
							if (type == LlmSegment::Type::ToolCall) {
								auto* call = new LlmToolCall(
									segmentRow.timestamp, generation);
								call->setName(segmentRow.name);
								call->setToolCallId(segmentRow.toolCallId);
								call->appendArguments(segmentRow.arguments);
								call->setResult(segmentRow.result);
								call->setStatus(
									static_cast<LlmToolCall::Status>(
										segmentRow.status));
								call->restore(segmentRow.elapsedMs);
								segment = call;
							} else {
								auto* text = new LlmTextSegment(
									type, segmentRow.timestamp, generation);
								text->setText(segmentRow.text);
								text->restore(segmentRow.elapsedMs);
								segment = text;
							}
							generation->addSegment(segment);
						}
						if (generationRow.active) activeIndex = i;
					}
					message->setActiveGeneration(activeIndex);
					messages.append(message);
				}
				if (model->rowCount() > 0) {
					QList<ChatMessage*> live = messages;
					for (int r = 0; r < model->rowCount(); ++r)
						live.prepend(model->at(r));
					messages = live;
				}
				if (!messages.isEmpty() || model->rowCount() > 0)
					s->adoptMessages(messages);

				s->markLoaded();

				st->applyPendingClear(s);
				if (st->m_pendingPersists.remove(s)) st->persist(s);
			},
			Qt::QueuedConnection);
	});
}

void ChatStore::load() {
	m_loading = true;
	QSqlQuery query(db());
	query.exec(
		"SELECT s.id, s.title, s.icon, s.created_at, s.updated_at, s.pinned, "
		"(SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) AS cnt "
		"FROM sessions s ORDER BY s.pinned DESC, s.updated_at DESC");
	while (query.next()) {
		auto* session = new ChatSession(this, query.value(0).toString(), this);
		session->hydrate(
			query.value(1).toString(),
			query.value(2).toString(),
			query.value(3).toLongLong(),
			query.value(4).toLongLong(),
			query.value(6).toInt(),
			query.value(5).toInt() != 0);
		m_sessions.append(session);
	}
	m_loading = false;

	beginResetModel();
	endResetModel();
}

} // namespace ZShell::llm