import { agentRuns, type AgentTask, type Database } from '@pig/db';
import { eq } from 'drizzle-orm';
import type { AgentProvider } from './provider';
import { AgentTaskQueue } from './queue';
import { createPigTools } from './tools';

export interface PiggyWorkerOptions {
  pollIntervalMs: number;
  leaseSeconds: number;
}

export class PiggyWorker {
  constructor(
    private readonly db: Database,
    private readonly queue: AgentTaskQueue,
    private readonly provider: AgentProvider,
    private readonly options: PiggyWorkerOptions,
  ) {}

  async run(signal: AbortSignal): Promise<void> {
    while (!signal.aborted) {
      const handled = await this.runOnce(signal);
      if (!handled) await delay(this.options.pollIntervalMs, signal);
    }
  }

  async runOnce(signal?: AbortSignal): Promise<boolean> {
    const task = await this.queue.claimNext();
    if (!task) return false;
    await this.process(task, signal);
    return true;
  }

  private async process(task: AgentTask, parentSignal?: AbortSignal): Promise<void> {
    const [run] = await this.db
      .insert(agentRuns)
      .values({
        agentTaskId: task.id,
        principalUserId: task.requestedByUserId,
        model: this.provider.model,
        input: {
          kind: task.kind,
          subject: task.subject,
          reason: task.reason,
          payload: task.payload,
        },
      })
      .returning({ id: agentRuns.id });
    if (!run) throw new Error('Could not create an agent run.');

    const leaseAbort = new AbortController();
    const signal = parentSignal
      ? AbortSignal.any([parentSignal, leaseAbort.signal])
      : leaseAbort.signal;
    let renewalRunning = false;
    const renewal = setInterval(() => {
      if (renewalRunning) return;
      renewalRunning = true;
      void this.queue
        .renew(task.id)
        .then((owned) => {
          if (!owned) leaseAbort.abort(new Error('Piggy lost its task lease.'));
        })
        .catch((error) => leaseAbort.abort(error))
        .finally(() => {
          renewalRunning = false;
        });
    }, Math.max(1_000, Math.floor((this.options.leaseSeconds * 1_000) / 2)));
    renewal.unref();

    try {
      const result = await this.provider.run({
        task,
        tools: createPigTools(this.db, { task, agentRunId: run.id }),
        signal,
      });

      const owned = await this.queue.succeed(task.id);
      if (!owned) throw new Error('Piggy completed after losing its task lease.');

      await this.db
        .update(agentRuns)
        .set({
          status: 'succeeded',
          summary: result.summary,
          result: result.result,
          inputTokens: result.inputTokens,
          outputTokens: result.outputTokens,
          finishedAt: new Date(),
        })
        .where(eq(agentRuns.id, run.id));
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      await this.db
        .update(agentRuns)
        .set({ status: 'failed', error: message, finishedAt: new Date() })
        .where(eq(agentRuns.id, run.id));
      await this.queue.fail(task, message);
    } finally {
      clearInterval(renewal);
    }
  }
}

function delay(ms: number, signal: AbortSignal): Promise<void> {
  if (signal.aborted) return Promise.resolve();
  return new Promise((resolve) => {
    const timer = setTimeout(resolve, ms);
    signal.addEventListener(
      'abort',
      () => {
        clearTimeout(timer);
        resolve();
      },
      { once: true },
    );
  });
}
