#include "chatstore.hpp"

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

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

#include <algorithm>

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;
}

// A null QString binds as SQL NULL, which violates the NOT NULL columns;
// DEFAULT only applies to omitted columns, not explicit NULLs.
QString sqlText(const QString& value) {
	if (value.isNull())
		return QStringLiteral("");
	return value;
}

// Plain data for one session's messages, fetched on a worker thread
// and turned into the QObject tree on the GUI thread. Messages are
// ordered as the model displays them (newest first).
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 MessageRow {
	bool user = false;
	qint64 timestamp = 0;
	QVector<GenerationRow> generations;
};

} // namespace

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

ChatStore::~ChatStore() {
	if (m_connectionName.isEmpty())
		return;
	QSqlDatabase db = QSqlDatabase::database(m_connectionName, false);
	db.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 db =
		QSqlDatabase::addDatabase(QStringLiteral("QSQLITE"), m_connectionName);
	db.setDatabaseName(m_dbPath);
	if (!db.open()) {
		qWarning() << "ChatStore: failed to open database" << m_dbPath << ":"
		           << db.lastError().text();
		return;
	}
	{
		QSqlQuery pragma(db);
		pragma.exec(QStringLiteral("PRAGMA foreign_keys = ON"));
	}
	{
		QSqlQuery query(db);
		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(db);
		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"
				"  timestamp INTEGER NOT NULL\n"
				")"));
	}
	{
		QSqlQuery query(db);
		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(db);
		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(db);
		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)"));
	}
}

int ChatStore::count() const {
	return static_cast<int>(m_sessions.size());
}

QVariantList ChatStore::values() const {
	QVariantList vals;
	vals.reserve(m_sessions.size());
	for (const auto* session : m_sessions)
		vals.append(QVariant::fromValue(session));
	return vals;
}

ChatSession* ChatStore::at(int index) const {
	if (index < 0 || index >= m_sessions.size())
		return nullptr;
	return m_sessions.at(index);
}

ChatSession* ChatStore::insert(int index) {
	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	const QString id = QUuid::createUuid().toString();
	{
		QSqlQuery query(db());
		query.prepare(
			"INSERT INTO sessions (id, title, created_at, updated_at) "
			"VALUES (:id, '', :created_at, :updated_at)");
		query.bindValue(":id", id);
		query.bindValue(":created_at", now);
		query.bindValue(":updated_at", now);
		if (!query.exec())
			qWarning() << "ChatStore: failed to insert session" << id << ":"
			           << query.lastError().text();
	}
	auto* session = new ChatSession(id, this);
	session->setMeta(QString(), now, now, 0);
	const int pos = index >= 0 && index <= m_sessions.size() ? index : 0;
	m_sessions.insert(pos, session);
	Q_EMIT countChanged();
	Q_EMIT valuesChanged();
	return session;
}

void ChatStore::remove(int index) {
	removeSession(at(index));
}

void ChatStore::remove(ChatSession* chat) {
	removeSession(chat);
}

void ChatStore::removeSession(ChatSession* session) {
	if (!session || !m_sessions.contains(session))
		return;
	const QList<ChatSession*> before = m_sessions;
	Q_EMIT sessionRemoved(session);
	{
		QSqlQuery query(db());
		query.prepare("DELETE FROM sessions WHERE id = :id");
		query.bindValue(":id", session->id());
		query.exec();
	}
	m_sessions.removeOne(session);
	session->deleteLater();
	notify(before);
}

void ChatStore::move(int from, int to) {
	if (from < 0 || from >= m_sessions.size() || to < 0 ||
	    to >= m_sessions.size() || from == to)
		return;
	m_sessions.move(from, to);
	Q_EMIT valuesChanged();
}

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

ChatSession* ChatStore::sessionById(const QString& id) {
	for (auto* session : m_sessions)
		if (session->id() == id)
			return session;
	return nullptr;
}

void ChatStore::setLlmClient(LlmClient* client) {
	m_llmClient = client;
}

void ChatStore::persist(ChatSession* session) {
	if (!session || !m_sessions.contains(session))
		return;
	if (!session->isLoaded()) {
		// Saving now would persist an incomplete model and wipe the
		// stored history; run it again when the load lands.
		m_pendingPersists.insert(session);
		return;
	}
	session->setUpdatedAt(QDateTime::currentMSecsSinceEpoch());
	if (!saveSession(session))
		return;
	sortAndNotify();
}

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

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, updated_at = :updated_at "
			"WHERE id = :id");
		query.bindValue(":title", sqlText(session->title()));
		query.bindValue(":updated_at", session->updatedAtMs());
		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, timestamp) "
			"VALUES (:id, :role, :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)");
		// The model holds messages most recent first; the database keeps
		// natural rowid order, so iterate from the oldest row up.
		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(":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()));
					segmentInsert.bindValue(":text", sqlText(segment->text()));
					segmentInsert.bindValue(":name", sqlText(segment->name()));
					segmentInsert.bindValue(
						":tool_call_id", sqlText(segment->toolCallId()));
					segmentInsert.bindValue(
						":arguments", sqlText(segment->arguments()));
					segmentInsert.bindValue(":result", sqlText(segment->result()));
					segmentInsert.bindValue(
						":status", static_cast<int>(segment->status()));
					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 || !handle.commit()) {
		qWarning() << "ChatStore: saveSession" << id << "commit failed, rolling back";
		handle.rollback();
		ok = false;
		qWarning() << "ChatStore: failed to save session" << session->id() << ":"
		           << handle.lastError().text();
	}
	return ok;
}
void ChatStore::loadMessagesInto(ChatSession* session) {
	if (!session)
		return;
	const QString sessionId = session->id();
	const QString path = m_dbPath;

	// SQL on a worker thread (its own connection; QSqlDatabase objects
	// are thread-affine). Rows come back as plain data.
	QThreadPool::globalInstance()->start(
		[store = QPointer<ChatStore>(this),
		 session = QPointer<ChatSession>(session), sessionId, path]() {
			QList<MessageRow> rows;
			const QString connName = QUuid::createUuid().toString();
			{
				QSqlDatabase db = QSqlDatabase::addDatabase(
					QStringLiteral("QSQLITE"), connName);
				db.setDatabaseName(path);
				if (db.open()) {
					// Tolerate the GUI thread writing while we read.
					QSqlQuery busy(db);
					busy.exec(QStringLiteral("PRAGMA busy_timeout = 3000"));
					// Newest first so the model receives rows in
					// display order.
					QSqlQuery query(db);
					query.prepare(
						"SELECT id, role, 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.timestamp = query.value(2).toLongLong();
							QSqlQuery generationQuery(db);
							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(db);
									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();
							}
							rows.append(message);
						}
					}
					db.close();
				} else {
					qWarning()
						<< "ChatStore: failed to open database for load:"
						<< db.lastError().text();
				}
			}
			// Remove only once every QSqlDatabase copy and query is gone;
			// while any reference is alive Qt refuses the removal and the
			// connection is left dangling in a broken state.
			QSqlDatabase::removeDatabase(connName);

			// Build the object tree on the GUI thread. Deliver through
			// the app instance (never destroyed) and re-check the
			// pointers there: posting to `store` from the pool thread
			// would race with its destruction.
			QMetaObject::invokeMethod(
				QCoreApplication::instance(),
				[store, session, rows = std::move(rows)]() mutable {
					ChatStore* st = store;
					ChatSession* s = session;
					if (!st || !s)
						return;

					// Rows fetched from disk; newest first.
					auto* model = s->model();
					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);
						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);
							for (const SegmentRow& segmentRow :
							     generationRow.segments) {
								auto* segment = new LlmSegment(
									segmentTypeFromName(segmentRow.type),
									segmentRow.timestamp,
									generation);
								segment->setText(segmentRow.text);
								segment->setName(segmentRow.name);
								segment->setToolCallId(segmentRow.toolCallId);
								segment->appendArguments(segmentRow.arguments);
								segment->setResult(segmentRow.result);
								segment->setStatus(
									static_cast<LlmSegment::Status>(
										segmentRow.status));
								segment->restore(segmentRow.elapsedMs);
								generation->addSegment(segment);
							}
							if (generationRow.active)
								activeIndex = i;
						}
						message->setActiveGeneration(activeIndex);
						messages.append(message);
					}
					// Rows added live while the load was in flight are
					// newer than anything on disk; keep them in front.
					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);

					// Mark loaded only once the model holds both the
					// fetched history and the rows added live while the
					// load ran, so a deferred startGeneration (triggered
					// by loaded()) builds its context from the complete
					// conversation.
					s->markLoaded();

					if (s->takeClearPending()) {
						// Cleared while the load was in flight; drop
						// everything now that the model is populated.
						s->clear();
					} else if (st->m_pendingPersists.remove(s)) {
						st->persist(s);
					}
				},
				Qt::QueuedConnection);
		});
}

void ChatStore::load() {
	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(query.value(0).toString(), this);
		session->setMeta(
			query.value(1).toString(),
			query.value(3).toLongLong(),
			query.value(4).toLongLong(),
			query.value(6).toInt());
		session->setIcon(query.value(2).toString());
		session->setPinned(query.value(5).toInt() != 0);
		m_sessions.append(session);
	}
	sortAndNotify();
}

void ChatStore::sortAndNotify() {
	const QList<ChatSession*> before = m_sessions;
	std::stable_sort(
		m_sessions.begin(),
		m_sessions.end(),
		[](const ChatSession* a, const ChatSession* b) {
			if (a->pinned() != b->pinned())
				return a->pinned() > b->pinned();
			return a->updatedAtMs() > b->updatedAtMs();
		});
	notify(before);
}

void ChatStore::notify(const QList<ChatSession*>& before) {
	if (before.size() != m_sessions.size())
		Q_EMIT countChanged();
	bool same = before.size() == m_sessions.size();
	for (int i = 0; same && i < m_sessions.size(); ++i)
		if (before.at(i) != m_sessions.at(i))
			same = false;
	if (!same)
		Q_EMIT valuesChanged();
}

} // namespace ZShell::llm