import { describe, it, expect, beforeAll, afterAll, beforeEach } from "vitest";
import {
  getTestPool,
  setupTestDb,
  cleanTestDb,
  teardownTestDb,
  uuidv4,
} from "./helpers.js";
import pg from "pg";

describe("Jobs", () => {
  let pool: pg.Pool;
  let orgId: string;
  let userId: string;

  beforeAll(async () => {
    pool = await getTestPool();
  });

  afterAll(async () => {
    await teardownTestDb();
  });

  beforeEach(async () => {
    await cleanTestDb(pool);
    const seeded = await setupTestDb();
    orgId = seeded.orgId;
    userId = seeded.userId;
  });

  it("should create a job with targets", async () => {
    const deviceId = uuidv4();
    await pool.query(
      `INSERT INTO devices (id, organization_id, name, device_type, platform, os_name, os_version, architecture, agent_version, status)
       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`,
      [deviceId, orgId, "job-test", "linux_server", "linux", "Ubuntu", "22.04", "amd64", "0.1.0", "online"]
    );

    const { rows: jobRows } = await pool.query(
      `INSERT INTO jobs (organization_id, created_by_user_id, job_type, status, parameters)
       VALUES ($1, $2, 'update_inventory', 'queued', '{}') RETURNING id`,
      [orgId, userId]
    );
    const jobId = jobRows[0].id;

    await pool.query(
      `INSERT INTO job_targets (job_id, device_id, status) VALUES ($1, $2, 'queued')`,
      [jobId, deviceId]
    );

    const { rows: jobs } = await pool.query(
      "SELECT * FROM jobs WHERE id = $1",
      [jobId]
    );
    expect(jobs.length).toBe(1);
    expect(jobs[0].job_type).toBe("update_inventory");
    expect(jobs[0].status).toBe("queued");

    const { rows: targets } = await pool.query(
      "SELECT * FROM job_targets WHERE job_id = $1",
      [jobId]
    );
    expect(targets.length).toBe(1);
    expect(targets[0].device_id).toBe(deviceId);
  });

  it("should transition job through states", async () => {
    const { rows: jobRows } = await pool.query(
      `INSERT INTO jobs (organization_id, created_by_user_id, job_type, status, parameters)
       VALUES ($1, $2, 'update_system_info', 'queued', '{}') RETURNING id`,
      [orgId, userId]
    );
    const jobId = jobRows[0].id;

    await pool.query(
      "UPDATE jobs SET status = 'dispatched', dispatched_at = NOW() WHERE id = $1",
      [jobId]
    );

    const { rows: step1 } = await pool.query("SELECT status FROM jobs WHERE id = $1", [jobId]);
    expect(step1[0].status).toBe("dispatched");

    await pool.query(
      "UPDATE jobs SET status = 'running', started_at = NOW() WHERE id = $1",
      [jobId]
    );

    const { rows: step2 } = await pool.query("SELECT status FROM jobs WHERE id = $1", [jobId]);
    expect(step2[0].status).toBe("running");

    await pool.query(
      "UPDATE jobs SET status = 'succeeded', completed_at = NOW() WHERE id = $1",
      [jobId]
    );

    const { rows: step3 } = await pool.query("SELECT status FROM jobs WHERE id = $1", [jobId]);
    expect(step3[0].status).toBe("succeeded");
  });

  it("should cancel queued jobs", async () => {
    const { rows: jobRows } = await pool.query(
      `INSERT INTO jobs (organization_id, created_by_user_id, job_type, status, parameters)
       VALUES ($1, $2, 'reboot_system', 'queued', '{}') RETURNING id`,
      [orgId, userId]
    );
    const jobId = jobRows[0].id;

    const { rows: updated } = await pool.query(
      `UPDATE jobs SET status = 'cancelled', completed_at = NOW()
       WHERE id = $1 AND status = 'queued' RETURNING id`,
      [jobId]
    );
    expect(updated.length).toBe(1);

    const { rows: step } = await pool.query("SELECT status FROM jobs WHERE id = $1", [jobId]);
    expect(step[0].status).toBe("cancelled");
  });

  it("should not cancel running jobs", async () => {
    const { rows: jobRows } = await pool.query(
      `INSERT INTO jobs (organization_id, created_by_user_id, job_type, status, parameters)
       VALUES ($1, $2, 'update_inventory', 'running', '{}') RETURNING id`,
      [orgId, userId]
    );
    const jobId = jobRows[0].id;

    const { rows: updated } = await pool.query(
      `UPDATE jobs SET status = 'cancelled', completed_at = NOW()
       WHERE id = $1 AND status = 'queued' RETURNING id`,
      [jobId]
    );
    expect(updated.length).toBe(0);

    const { rows: step } = await pool.query("SELECT status FROM jobs WHERE id = $1", [jobId]);
    expect(step[0].status).toBe("running");
  });

  it("should complete job when all targets finish", async () => {
    const deviceId1 = uuidv4();
    const deviceId2 = uuidv4();

    await pool.query(
      `INSERT INTO devices (id, organization_id, name, device_type, platform, os_name, os_version, architecture, agent_version, status)
       VALUES ($1, $2, 'multi-1', 'linux_server', 'linux', 'Ubuntu', '22.04', 'amd64', '0.1.0', 'online')`,
      [deviceId1, orgId]
    );
    await pool.query(
      `INSERT INTO devices (id, organization_id, name, device_type, platform, os_name, os_version, architecture, agent_version, status)
       VALUES ($1, $2, 'multi-2', 'linux_server', 'linux', 'Ubuntu', '22.04', 'amd64', '0.1.0', 'online')`,
      [deviceId2, orgId]
    );

    const { rows: jobRows } = await pool.query(
      `INSERT INTO jobs (organization_id, created_by_user_id, job_type, status, parameters)
       VALUES ($1, $2, 'update_inventory', 'dispatched', '{}') RETURNING id`,
      [orgId, userId]
    );
    const jobId = jobRows[0].id;

    await pool.query(
      `INSERT INTO job_targets (job_id, device_id, status) VALUES ($1, $2, 'dispatched'), ($1, $3, 'dispatched')`,
      [jobId, deviceId1, deviceId2]
    );

    await pool.query(
      "UPDATE job_targets SET status = 'succeeded', completed_at = NOW() WHERE job_id = $1 AND device_id = $2",
      [jobId, deviceId1]
    );

    const { rows: partial } = await pool.query(
      "SELECT status FROM jobs WHERE id = $1",
      [jobId]
    );
    expect(partial[0].status).toBe("dispatched");

    await pool.query(
      "UPDATE job_targets SET status = 'succeeded', completed_at = NOW() WHERE job_id = $1 AND device_id = $2",
      [jobId, deviceId2]
    );

    const { rows: targets } = await pool.query(
      "SELECT status FROM job_targets WHERE job_id = $1",
      [jobId]
    );
    const allDone = targets.every((t: any) => !["queued", "dispatched", "running"].includes(t.status));
    expect(allDone).toBe(true);
  });

  it("should fail job when any target fails", async () => {
    const deviceId1 = uuidv4();
    const deviceId2 = uuidv4();

    await pool.query(
      `INSERT INTO devices (id, organization_id, name, device_type, platform, os_name, os_version, architecture, agent_version, status)
       VALUES ($1, $2, 'fail-1', 'linux_server', 'linux', 'Ubuntu', '22.04', 'amd64', '0.1.0', 'online'),
              ($3, $2, 'fail-2', 'linux_server', 'linux', 'Ubuntu', '22.04', 'amd64', '0.1.0', 'online')`,
      [deviceId1, orgId, deviceId2]
    );

    const { rows: jobRows } = await pool.query(
      `INSERT INTO jobs (organization_id, created_by_user_id, job_type, status, parameters)
       VALUES ($1, $2, 'reboot_system', 'dispatched', '{}') RETURNING id`,
      [orgId, userId]
    );
    const jobId = jobRows[0].id;

    await pool.query(
      `INSERT INTO job_targets (job_id, device_id, status) VALUES ($1, $2, 'dispatched'), ($1, $3, 'dispatched')`,
      [jobId, deviceId1, deviceId2]
    );

    await pool.query(
      "UPDATE job_targets SET status = 'succeeded', completed_at = NOW() WHERE job_id = $1 AND device_id = $2",
      [jobId, deviceId1]
    );
    await pool.query(
      `UPDATE job_targets SET status = 'failed', completed_at = NOW(), result = '{"error": "timeout"}' WHERE job_id = $1 AND device_id = $2`,
      [jobId, deviceId2]
    );

    const { rows: targets } = await pool.query(
      "SELECT status FROM job_targets WHERE job_id = $1",
      [jobId]
    );
    const anyFailed = targets.some((t: any) => t.status === "failed");
    expect(anyFailed).toBe(true);
  });
});