import { randomUUID } from 'node:crypto';
import {
	mapChatMessagesToStoredMessages,
	mapStoredMessagesToChatMessages,
} from '@langchain/core/messages';
import {
	isLangchainMessagesArray,
	type ISessionStorage,
	type LangchainMessage,
	type StoredSession,
} from '@n8n/ai-workflow-builder';
import { Service } from '@n8n/di';
import { DataSource, Repository } from '@n8n/typeorm';

import { WorkflowBuilderSession } from './workflow-builder-session.entity';

@Service()
export class WorkflowBuilderSessionRepository
	extends Repository<WorkflowBuilderSession>
	implements ISessionStorage
{
	constructor(dataSource: DataSource) {
		super(WorkflowBuilderSession, dataSource.manager);
	}

	async getSession(threadId: string): Promise<StoredSession | null> {
		const { workflowId, userId } = this.parseThreadId(threadId);
		const entity = await this.findOne({ where: { workflowId, userId } });

		if (!entity) return null;

		const restoredMessages = mapStoredMessagesToChatMessages(entity.messages);
		const messages: LangchainMessage[] = isLangchainMessagesArray(restoredMessages)
			? restoredMessages
			: [];

		return {
			messages,
			previousSummary: entity.previousSummary ?? undefined,
			updatedAt: entity.updatedAt,
			activeVersionCardId: entity.activeVersionCardId,
			resumeAfterRestoreMessageId: entity.resumeAfterRestoreMessageId,
		};
	}

	async saveSession(threadId: string, data: StoredSession): Promise<void> {
		const { workflowId, userId } = this.parseThreadId(threadId);
		const messages = mapChatMessagesToStoredMessages(data.messages);
		const previousSummary = data.previousSummary ?? null;

		const activeVersionCardId = data.activeVersionCardId ?? null;
		const resumeAfterRestoreMessageId = data.resumeAfterRestoreMessageId ?? null;

		await this.createQueryBuilder()
			.insert()
			.into(WorkflowBuilderSession)
			.values({
				id: randomUUID(),
				workflowId,
				userId,
				messages,
				previousSummary,
				activeVersionCardId,
				resumeAfterRestoreMessageId,
			})
			.orUpdate(
				['messages', 'previousSummary', 'activeVersionCardId', 'resumeAfterRestoreMessageId'],
				['workflowId', 'userId'],
			)
			.execute();
	}

	async deleteSession(threadId: string): Promise<void> {
		const { workflowId, userId } = this.parseThreadId(threadId);
		await this.delete({ workflowId, userId });
	}

	private parseThreadId(threadId: string): { workflowId: string; userId: string } {
		// Format: "workflow-{workflowId}-user-{userId}" with an optional "-code" suffix
		// for the code-builder agent thread variant. Strip the suffix before parsing so
		// the greedy userId capture doesn't swallow it (userId is a uuid column in PG).
		const normalized = threadId.endsWith('-code') ? threadId.slice(0, -'-code'.length) : threadId;
		const match = normalized.match(/^workflow-(.+)-user-(.+)$/);
		if (!match) {
			throw new Error(`Invalid thread ID format: ${threadId}`);
		}
		return { workflowId: match[1], userId: match[2] };
	}
}
