import { drizzle } from "drizzle-orm/postgres-js";
import postgres from "postgres";
import { sql } from "drizzle-orm";

const connectionString =
  process.env.DATABASE_URL ??
  "postgresql://sec_cybert:sec_cybert@localhost:5432/sec_cybert";

const client = postgres(connectionString);
const db = drizzle(client);

export async function resetTestData() {
  await db.execute(sql`DELETE FROM human_labels WHERE annotator_id LIKE 'e2e-%'`);
  await db.execute(sql`DELETE FROM adjudications WHERE adjudicator_id LIKE 'e2e-%'`);
  await db.execute(sql`DELETE FROM quiz_sessions WHERE annotator_id LIKE 'e2e-%'`);
  await db.execute(sql`DELETE FROM assignments WHERE annotator_id LIKE 'e2e-%'`);
  await db.execute(sql`DELETE FROM annotators WHERE id LIKE 'e2e-%'`);
}

export async function seedE2EData() {
  // Create E2E test annotator
  await db.execute(sql`
    INSERT INTO annotators (id, display_name, password)
    VALUES ('e2e-tester', 'E2E Tester', 'e2e-tester')
    ON CONFLICT DO NOTHING
  `);

  // Assign 3 paragraphs from the existing data
  await db.execute(sql`
    INSERT INTO assignments (paragraph_id, annotator_id)
    SELECT p.id, 'e2e-tester'
    FROM paragraphs p
    LIMIT 3
    ON CONFLICT DO NOTHING
  `);
}

export async function cleanup() {
  await client.end();
}