import { describe, it, expect, beforeEach } from 'vitest';
import { Repository } from '../../db/repository.js';
import { SqliteA2aTaskStore } from './task-store.js';
import type { Task } from '@a2a-js/sdk';

describe('SqliteA2aTaskStore', () => {
  let repo: Repository;
  let store: SqliteA2aTaskStore;

  beforeEach(() => {
    repo = new Repository(':memory:');
    store = new SqliteA2aTaskStore(repo);
  });

  const makeTask = (overrides: Partial<Task> = {}): Task => ({
    id: 'task-001',
    contextId: 'ctx-001',
    kind: 'task',
    status: { state: 'submitted' },
    metadata: { maestroJobId: 'job-1' },
    ...overrides,
  });

  it('save then load returns an equal Task', async () => {
    const task = makeTask();
    await store.save(task);
    const loaded = await store.load(task.id);
    expect(loaded).toBeDefined();
    expect(loaded!.id).toBe(task.id);
    expect(loaded!.contextId).toBe(task.contextId);
    expect(loaded!.kind).toBe('task');
    expect(loaded!.status.state).toBe('submitted');
    expect(loaded!.metadata?.maestroJobId).toBe('job-1');
  });

  it('load of unknown id returns undefined', async () => {
    const loaded = await store.load('no-such-task');
    expect(loaded).toBeUndefined();
  });

  it('job link persisted: getA2aTaskByJobId finds saved task', async () => {
    const task = makeTask();
    await store.save(task);
    const found = repo.getA2aTaskByJobId('job-1');
    expect(found).toBeDefined();
    expect(found!.id).toBe(task.id);
  });

  it('save overwrites existing task (upsert)', async () => {
    const task = makeTask();
    await store.save(task);
    const updated: Task = { ...task, status: { state: 'completed' } };
    await store.save(updated);
    const loaded = await store.load(task.id);
    expect(loaded!.status.state).toBe('completed');
  });

  it('localTaskId stored when metadata.maestroLocalTaskId provided', async () => {
    const task = makeTask({
      id: 'task-002',
      metadata: { maestroJobId: 'job-2', maestroLocalTaskId: 42 },
    });
    await store.save(task);
    const found = repo.getA2aTaskByJobId('job-2');
    expect(found).toBeDefined();
    expect(found!.id).toBe('task-002');
  });

  it('delegation metadata persisted: a2aDelegationId, a2aGrantId, a2aActingUserId', async () => {
    const task = makeTask({
      id: 'task-deleg-store',
      metadata: {
        maestroJobId: 'job-deleg',
        a2aDelegationId: 'deleg-abc',
        a2aGrantId: 'grant-def',
        a2aActingUserId: 'user-ghi',
      },
    });
    await store.save(task);

    // Verify via raw DB query that the columns were written
    const db = (repo as any).db;
    const row = db.prepare(
      'SELECT delegation_id, grant_id, acting_user_id FROM a2a_tasks WHERE id = ?'
    ).get('task-deleg-store') as { delegation_id: string; grant_id: string; acting_user_id: string };

    expect(row).toBeDefined();
    expect(row.delegation_id).toBe('deleg-abc');
    expect(row.grant_id).toBe('grant-def');
    expect(row.acting_user_id).toBe('user-ghi');
  });

  it('delegation columns null when metadata lacks a2a keys', async () => {
    const task = makeTask({ id: 'task-no-deleg-store' });
    await store.save(task);

    const db = (repo as any).db;
    const row = db.prepare(
      'SELECT delegation_id, grant_id, acting_user_id FROM a2a_tasks WHERE id = ?'
    ).get('task-no-deleg-store') as { delegation_id: string | null; grant_id: string | null; acting_user_id: string | null };

    expect(row.delegation_id).toBeNull();
    expect(row.grant_id).toBeNull();
    expect(row.acting_user_id).toBeNull();
  });
});
