import { pool } from "./db.js";
import type { TextChunk } from "./chunking.js";

export interface NewDocument {
  title: string;
  source?: string;
}

export interface StoredChunk extends TextChunk {
  embedding: number[];
}

export const createDocument = async (
  document: NewDocument,
): Promise<string> => {
  const result = await pool.query<{ id: string }>(
    `INSERT INTO documents (title, source, status)
    VALUES ($1, $2, 'processing')
    RETURNING id`,
    [document.title, document.source ?? null],
  );

  const id = result.rows[0]?.id;

  if (!id) {
    throw new Error("Failed to create document");
  }

  return id;
};

export const storeDocumentChunks = async (
  documentId: string,
  chunks: StoredChunk[],
): Promise<void> => {
  const client = await pool.connect();

  try {
    await client.query("BEGIN");

    for (const chunk of chunks) {
      await client.query(
        `INSERT INTO document_chunks
          (document_id, chunk_index, content, embedding)
        VALUES ($1, $2, $3, $4::vector)`,
        [
          documentId,
          chunk.chunkIndex,
          chunk.content,
          `[${chunk.embedding.join(",")}]`,
        ],
      );
    }

    await client.query(
      `UPDATE documents
      SET status = 'completed'
      WHERE id = $1`,
      [documentId],
    );

    await client.query("COMMIT");
  } catch (error) {
    try {
      await client.query("ROLLBACK");
    } catch {
      // Preserve the original error.
    }

    throw error;
  } finally {
    client.release();
  }
};