import { EmbeddingService, SupabaseProjectStore, SupabaseStore } from "@imemio/sdk";
import type { SupabaseClient } from "@supabase/supabase-js";
import { TRPCError } from "@trpc/server";
import { z } from "zod";
import { env } from "~/env";
import { createClient } from "~/lib/supabase/server";
import { createTRPCRouter, protectedProcedure } from "~/server/api/trpc";

function getEmbeddingService(): EmbeddingService | null {
	const openaiApiKey = env.OPENAI_API_KEY;

	if (!openaiApiKey) {
		return null;
	}

	return new EmbeddingService({ apiKey: openaiApiKey });
}

const defaultProjectName = "default";

async function getOrCreateDefaultProject(
	supabaseClient: SupabaseClient,
	userId: string,
): Promise<string> {
	const projectStore = new SupabaseProjectStore(supabaseClient, userId);
	const projectsResult = await projectStore.list();

	if (projectsResult.success && projectsResult.value.length > 0) {
		const existingDefaultProject = projectsResult.value.find(
			(project) => project.name === defaultProjectName,
		);

		if (existingDefaultProject) {
			return existingDefaultProject.id;
		}

		const firstProject = projectsResult.value[0];

		if (firstProject) {
			return firstProject.id;
		}
	}

	const createResult = await projectStore.create({ name: defaultProjectName });

	if (!createResult.success) {
		throw new Error("Failed to create default project");
	}

	return createResult.value.id;
}

export const memoryRouter = createTRPCRouter({
	list: protectedProcedure.query(async ({ ctx: context }) => {
		const supabaseClient = await createClient();
		const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
		const memories = await memoryStore.list();

		return memories;
	}),

	listWithProjects: protectedProcedure
		.input(
			z.object({
				projectId: z.string().optional(),
				sortOrder: z.enum(["asc", "desc"]).optional().default("desc"),
			}),
		)
		.query(async ({ ctx: context, input }) => {
			const supabaseClient = await createClient();
			const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
			const projectStore = new SupabaseProjectStore(
				supabaseClient,
				context.user.id,
			);
			const memories = await memoryStore.list(
				input.projectId ? { projectId: input.projectId } : undefined,
			);
			const projectsResult = await projectStore.list();
			const projects = projectsResult.success ? projectsResult.value : [];
			const projectMap = new Map(projects.map((project) => [project.id, project]));
			const memoriesWithProjects = memories.map((memory) => ({
				...memory,
				project: projectMap.get(memory.projectId) ?? null,
			}));
			const sortedMemories = memoriesWithProjects.sort((memoryA, memoryB) => {
				const comparison =
					new Date(memoryA.createdAt).getTime() -
					new Date(memoryB.createdAt).getTime();

				return input.sortOrder === "desc" ? -comparison : comparison;
			});

			return { memories: sortedMemories, projects };
		}),

	create: protectedProcedure
		.input(
			z.object({
				content: z.string().min(1),
				projectId: z.string().optional(),
			}),
		)
		.mutation(async ({ ctx: context, input }) => {
			const supabaseClient = await createClient();
			const projectId =
				input.projectId ??
				(await getOrCreateDefaultProject(supabaseClient, context.user.id));
			const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
			const memory = await memoryStore.create({
				content: input.content,
				projectId,
			});

			return memory;
		}),

	update: protectedProcedure
		.input(
			z.object({
				id: z.string(),
				content: z.string().min(1).optional(),
				projectId: z.string().optional(),
			}),
		)
		.mutation(async ({ ctx: context, input }) => {
			const supabaseClient = await createClient();
			const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
			const updateData: { content?: string; projectId?: string } = {};

			if (input.content !== undefined) {
				updateData.content = input.content;
			}

			if (input.projectId !== undefined) {
				updateData.projectId = input.projectId;
			}

			const updateResult = await memoryStore.update(input.id, updateData);

			if (!updateResult.success) {
				throw new TRPCError({
					code: "NOT_FOUND",
					message: `Memory not found: ${input.id}`,
				});
			}

			return updateResult.value;
		}),

	delete: protectedProcedure
		.input(z.object({ id: z.string() }))
		.mutation(async ({ ctx: context, input }) => {
			const supabaseClient = await createClient();
			const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
			const deleteResult = await memoryStore.delete(input.id);

			if (!deleteResult.success) {
				throw new TRPCError({
					code: "NOT_FOUND",
					message: `Memory not found: ${input.id}`,
				});
			}

			return { success: true };
		}),

	search: protectedProcedure
		.input(
			z.object({
				query: z.string().min(1),
				projectId: z.string().optional(),
				limit: z.number().optional().default(20),
			}),
		)
		.query(async ({ ctx: context, input }) => {
			const embeddingService = getEmbeddingService();

			if (!embeddingService) {
				throw new TRPCError({
					code: "PRECONDITION_FAILED",
					message: "Search requires OPENAI_API_KEY to be configured.",
				});
			}

			const supabaseClient = await createClient();
			const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
			const projectStore = new SupabaseProjectStore(
				supabaseClient,
				context.user.id,
			);
			const queryEmbedding = await embeddingService.generate(input.query);
			const searchResults = await memoryStore.search(queryEmbedding, {
				projectId: input.projectId,
				limit: input.limit,
				threshold: 0.3,
			});
			const projectsResult = await projectStore.list();
			const projects = projectsResult.success ? projectsResult.value : [];
			const projectMap = new Map(projects.map((project) => [project.id, project]));
			const resultsWithProjects = searchResults.map((result) => ({
				...result,
				project: projectMap.get(result.projectId) ?? null,
			}));

			return { results: resultsWithProjects, projects };
		}),

	embeddingInfo: protectedProcedure.query(() => {
		const embeddingService = getEmbeddingService();

		return { available: embeddingService !== null };
	}),

	reembed: protectedProcedure.mutation(async ({ ctx: context }) => {
		const embeddingService = getEmbeddingService();

		if (!embeddingService) {
			throw new TRPCError({
				code: "PRECONDITION_FAILED",
				message: "Search is not configured on this server.",
			});
		}

		const supabaseClient = await createClient();
		const memoryStore = new SupabaseStore(supabaseClient, context.user.id);
		const memories = await memoryStore.list();
		let updated = 0;

		for (const memory of memories) {
			const embedding = await embeddingService.generate(memory.content);
			const { error } = await supabaseClient
				.from("memories")
				.update({
					embedding: JSON.stringify(embedding),
					embedding_dimensions: embeddingService.dimensions,
				})
				.eq("id", memory.id)
				.eq("user_id", context.user.id);

			if (!error) {
				updated++;
			}
		}

		return { total: memories.length, updated };
	}),
});