import pg from 'pg';
import { createInMemoryExecutorJobStore } from './executor-gateway.mjs';

const { Pool } = pg;
const DEFAULT_EXECUTOR_SCHEMA = 'memind_orchestrator';
const EXECUTOR_JOB_TABLE = 'executor_jobs';
const EXECUTOR_JOB_EVENT_TABLE = 'executor_job_events';
const EXECUTOR_WORKER_TABLE = 'executor_workers';

function normalizeMode(value) {
  const mode = String(value ?? 'postgres').trim().toLowerCase();
  if (mode === 'memory' || mode === 'postgres') return mode;
  throw new Error(`Unsupported executor job store mode: ${value}`);
}

function normalizeSchema(value) {
  const schema = String(value ?? DEFAULT_EXECUTOR_SCHEMA).trim();
  if (!/^[a-z_][a-z0-9_]{0,62}$/i.test(schema)) {
    throw new Error('Invalid Executor Gateway PostgreSQL schema');
  }
  return schema;
}

function parseState(value) {
  if (value == null) return null;
  if (typeof value === 'object') return structuredClone(value);
  try {
    return JSON.parse(String(value));
  } catch {
    return null;
  }
}

function projectRow(row) {
  return parseState(row?.state_json);
}

function projectEventRow(row) {
  return parseState(row?.event_json);
}

function eventWithSequence(event, jobId, sequence) {
  return {
    ...structuredClone(event),
    jobId,
    sequence,
  };
}

async function appendEvent(client, eventTable, jobId, event) {
  if (!event) return null;
  const sequenceResult = await client.query(
    `SELECT COALESCE(MAX(sequence), 0) + 1 AS next_sequence
       FROM ${eventTable}
      WHERE job_id = $1`,
    [jobId],
  );
  const sequence = Number(sequenceResult.rows?.[0]?.next_sequence ?? 1);
  const persistedEvent = eventWithSequence(event, jobId, sequence);
  await client.query(
    `INSERT INTO ${eventTable}
      (event_id, job_id, sequence, event_json, created_at)
     VALUES ($1, $2, $3, $4::jsonb, $5)`,
    [
      persistedEvent.eventId,
      jobId,
      sequence,
      JSON.stringify(persistedEvent),
      persistedEvent.timestamp,
    ],
  );
  return persistedEvent;
}

async function withTransaction(pool, operation) {
  if (typeof pool?.connect !== 'function') {
    throw new Error('PostgreSQL Executor Job Store requires transactional pool.connect');
  }
  const client = await pool.connect();
  try {
    await client.query('BEGIN');
    const result = await operation(client);
    await client.query('COMMIT');
    return result;
  } catch (error) {
    try {
      await client.query('ROLLBACK');
    } catch {
      // Preserve the transactional operation failure.
    }
    throw error;
  } finally {
    client.release?.();
  }
}

export function createPostgresExecutorJobStore({
  pool,
  schema = DEFAULT_EXECUTOR_SCHEMA,
} = {}) {
  if (!pool?.query) throw new Error('PostgreSQL Executor Job Store requires a pool');
  const normalizedSchema = normalizeSchema(schema);
  const table = `"${normalizedSchema}"."${EXECUTOR_JOB_TABLE}"`;
  const eventTable = `"${normalizedSchema}"."${EXECUTOR_JOB_EVENT_TABLE}"`;
  const workerTable = `"${normalizedSchema}"."${EXECUTOR_WORKER_TABLE}"`;

  return {
    kind: 'postgres',
    durable: true,

    async ensureSchema() {
      await pool.query(`CREATE SCHEMA IF NOT EXISTS "${normalizedSchema}"`);
      await pool.query(`
        CREATE TABLE IF NOT EXISTS ${table} (
          job_id VARCHAR(128) PRIMARY KEY,
          idempotency_key VARCHAR(200) NOT NULL UNIQUE,
          request_fingerprint CHAR(64) NOT NULL,
          state_json JSONB NOT NULL,
          created_at BIGINT NOT NULL,
          updated_at BIGINT NOT NULL
        )
      `);
      await pool.query(`
        CREATE INDEX IF NOT EXISTS executor_jobs_updated_at_idx
          ON ${table} (updated_at DESC)
      `);
      await pool.query(`
        CREATE TABLE IF NOT EXISTS ${eventTable} (
          event_id VARCHAR(128) PRIMARY KEY,
          job_id VARCHAR(128) NOT NULL REFERENCES ${table} (job_id) ON DELETE CASCADE,
          sequence INTEGER NOT NULL,
          event_json JSONB NOT NULL,
          created_at BIGINT NOT NULL,
          UNIQUE (job_id, sequence)
        )
      `);
      await pool.query(`
        CREATE INDEX IF NOT EXISTS executor_job_events_job_sequence_idx
          ON ${eventTable} (job_id, sequence)
      `);
      await pool.query(`
        INSERT INTO ${eventTable} (event_id, job_id, sequence, event_json, created_at)
        SELECT
          'evt_' || md5(job_id || ':snapshot-import'),
          job_id,
          1,
          jsonb_build_object(
            'version', 'executor-job-event-v1',
            'eventId', 'evt_' || md5(job_id || ':snapshot-import'),
            'jobId', job_id,
            'sequence', 1,
            'type', 'executor_job_snapshot_imported',
            'timestamp', created_at,
            'data', jsonb_build_object(
              'status', state_json->>'status',
              'reason', state_json->>'reason'
            )
          ),
          created_at
        FROM ${table} jobs
        WHERE NOT EXISTS (
          SELECT 1 FROM ${eventTable} events WHERE events.job_id = jobs.job_id
        )
        ON CONFLICT DO NOTHING
      `);
      await pool.query(`
        CREATE TABLE IF NOT EXISTS ${workerTable} (
          worker_id VARCHAR(128) PRIMARY KEY,
          state_json JSONB NOT NULL,
          last_seen_at BIGINT NOT NULL
        )
      `);
      await pool.query(`
        CREATE INDEX IF NOT EXISTS executor_workers_last_seen_idx
          ON ${workerTable} (last_seen_at DESC)
      `);
    },

    async getById(jobId) {
      const result = await pool.query(
        `SELECT state_json FROM ${table} WHERE job_id = $1 LIMIT 1`,
        [String(jobId ?? '').trim()],
      );
      return projectRow(result.rows?.[0]);
    },

    async getByIdempotencyKey(key) {
      const result = await pool.query(
        `SELECT state_json FROM ${table} WHERE idempotency_key = $1 LIMIT 1`,
        [String(key ?? '').trim()],
      );
      return projectRow(result.rows?.[0]);
    },

    async deleteById(jobId) {
      const result = await pool.query(
        `DELETE FROM ${table}
          WHERE job_id = $1
          RETURNING job_id`,
        [String(jobId ?? '').trim()],
      );
      return Boolean(result.rows?.[0]);
    },

    async listTerminalBefore({ before, limit = 100 } = {}) {
      const cutoff = Number(before);
      if (!Number.isFinite(cutoff)) return [];
      const pageSize = Math.min(500, Math.max(1, Number(limit) || 100));
      const result = await pool.query(
        `SELECT
           state_json->>'jobId' AS job_id,
           state_json->>'workflowRunId' AS workflow_run_id,
           (state_json->>'completedAt')::BIGINT AS completed_at
         FROM ${table}
         WHERE state_json->>'status' = ANY($1::text[])
           AND state_json->>'workflowRunId' IS NOT NULL
           AND COALESCE((state_json->>'completedAt')::BIGINT, 0) > 0
           AND (state_json->>'completedAt')::BIGINT <= $2
         ORDER BY (state_json->>'completedAt')::BIGINT ASC, job_id ASC
         LIMIT $3`,
        [['succeeded', 'failed', 'cancelled', 'timed_out', 'blocked'], cutoff, pageSize],
      );
      return (result.rows ?? []).map((row) => ({
        jobId: row.job_id,
        workflowRunId: row.workflow_run_id,
        completedAt: Number(row.completed_at),
      }));
    },

    async createIfAbsent(record, { initialEvent = null } = {}) {
      const created = await withTransaction(pool, async (client) => {
        const inserted = await client.query(
          `INSERT INTO ${table}
            (job_id, idempotency_key, request_fingerprint, state_json, created_at, updated_at)
           VALUES ($1, $2, $3, $4::jsonb, $5, $6)
           ON CONFLICT DO NOTHING
           RETURNING state_json`,
          [
            record.jobId,
            record.idempotencyKey,
            record.requestFingerprint,
            JSON.stringify(record),
            record.createdAt,
            record.updatedAt,
          ],
        );
        const insertedRecord = projectRow(inserted.rows?.[0]);
        if (!insertedRecord || !initialEvent) return insertedRecord;
        const event = eventWithSequence(initialEvent, record.jobId, 1);
        await client.query(
          `INSERT INTO ${eventTable}
            (event_id, job_id, sequence, event_json, created_at)
           VALUES ($1, $2, $3, $4::jsonb, $5)`,
          [
            event.eventId,
            record.jobId,
            event.sequence,
            JSON.stringify(event),
            event.timestamp,
          ],
        );
        return insertedRecord;
      });
      if (created) return { created: true, record: created };

      const existingResult = await pool.query(
        `SELECT state_json FROM ${table} WHERE idempotency_key = $1 LIMIT 1`,
        [record.idempotencyKey],
      );
      const existing = projectRow(existingResult.rows?.[0]);
      if (existing) return { created: false, record: existing };

      const error = new Error(`Executor job id already exists: ${record.jobId}`);
      error.code = 'EXECUTOR_JOB_ID_CONFLICT';
      error.status = 409;
      throw error;
    },

    async update(record, { event = null } = {}) {
      return withTransaction(pool, async (client) => {
        const locked = await client.query(
          `SELECT job_id FROM ${table} WHERE job_id = $1 FOR UPDATE`,
          [record.jobId],
        );
        if (!locked.rows?.[0]) return null;
        const result = await client.query(
          `UPDATE ${table}
              SET state_json = $2::jsonb,
                  updated_at = $3
            WHERE job_id = $1
            RETURNING state_json`,
          [record.jobId, JSON.stringify(record), record.updatedAt],
        );
        await appendEvent(client, eventTable, record.jobId, event);
        return projectRow(result.rows?.[0]);
      });
    },

    async claimNext({
      workerId,
      executors,
      now,
      leaseToken,
      leaseDurationMs,
      eventFactory,
    }) {
      return withTransaction(pool, async (client) => {
        const selected = await client.query(
          `SELECT state_json
             FROM ${table}
            WHERE (
                state_json->>'status' = 'queued'
                OR (
                  state_json->>'status' = 'retryable'
                  AND COALESCE((state_json->>'nextAttemptAt')::BIGINT, 0) <= $1
                )
              )
              AND state_json->>'executor' = ANY($2::text[])
            ORDER BY created_at ASC, job_id ASC
            FOR UPDATE SKIP LOCKED
            LIMIT 1`,
          [now, executors],
        );
        const current = projectRow(selected.rows?.[0]);
        if (!current) return null;
        const updated = {
          ...current,
          status: 'leased',
          reason: null,
          attempts: Number(current.attempts ?? 0) + 1,
          nextAttemptAt: null,
          lease: {
            token: leaseToken,
            workerId,
            acquiredAt: now,
            heartbeatAt: now,
            expiresAt: now + leaseDurationMs,
          },
          updatedAt: now,
          completedAt: null,
        };
        const result = await client.query(
          `UPDATE ${table}
              SET state_json = $2::jsonb,
                  updated_at = $3
            WHERE job_id = $1
            RETURNING state_json`,
          [updated.jobId, JSON.stringify(updated), updated.updatedAt],
        );
        await appendEvent(
          client,
          eventTable,
          updated.jobId,
          typeof eventFactory === 'function' ? eventFactory(updated, current) : null,
        );
        return projectRow(result.rows?.[0]);
      });
    },

    async transition(jobId, {
      expectedStatuses = [],
      leaseToken = null,
      updater,
      eventFactory = null,
    } = {}) {
      return withTransaction(pool, async (client) => {
        const locked = await client.query(
          `SELECT state_json FROM ${table} WHERE job_id = $1 FOR UPDATE`,
          [String(jobId ?? '').trim()],
        );
        const current = projectRow(locked.rows?.[0]);
        if (!current) return { applied: false, record: null, reason: 'not_found' };
        if (expectedStatuses.length && !expectedStatuses.includes(current.status)) {
          return { applied: false, record: current, reason: 'status_mismatch' };
        }
        if (leaseToken != null && current.lease?.token !== leaseToken) {
          return { applied: false, record: current, reason: 'lease_mismatch' };
        }
        const updated = updater(structuredClone(current));
        if (!updated || updated.jobId !== current.jobId) {
          throw new Error('Executor job transition must preserve job identity');
        }
        const result = await client.query(
          `UPDATE ${table}
              SET state_json = $2::jsonb,
                  updated_at = $3
            WHERE job_id = $1
            RETURNING state_json`,
          [current.jobId, JSON.stringify(updated), updated.updatedAt],
        );
        await appendEvent(
          client,
          eventTable,
          current.jobId,
          typeof eventFactory === 'function' ? eventFactory(updated, current) : null,
        );
        return {
          applied: true,
          record: projectRow(result.rows?.[0]),
          reason: null,
        };
      });
    },

    async listExpiredLeases({ now, limit = 100 } = {}) {
      const pageSize = Math.min(500, Math.max(1, Number(limit) || 100));
      const result = await pool.query(
        `SELECT state_json
           FROM ${table}
          WHERE state_json->>'status' = ANY($1::text[])
            AND COALESCE((state_json->'lease'->>'expiresAt')::BIGINT, 0) <= $2
          ORDER BY COALESCE((state_json->'lease'->>'expiresAt')::BIGINT, 0) ASC
          LIMIT $3`,
        [['leased', 'running'], now, pageSize],
      );
      return (result.rows ?? []).map(projectRow).filter(Boolean);
    },

    async getQueueStats({ now = Date.now() } = {}) {
      const result = await pool.query(
        `SELECT
            COUNT(*)::BIGINT AS total,
            COUNT(*) FILTER (
              WHERE state_json->>'status' = 'queued'
                 OR (
                   state_json->>'status' = 'retryable'
                   AND COALESCE((state_json->>'nextAttemptAt')::BIGINT, 0) <= $1
                 )
            )::BIGINT AS claimable,
            COUNT(*) FILTER (
              WHERE state_json->>'status' = ANY($2::text[])
                AND COALESCE((state_json->'lease'->>'expiresAt')::BIGINT, 0) <= $1
            )::BIGINT AS expired_leases,
            state_json->>'status' AS status,
            COUNT(*)::BIGINT AS status_count
           FROM ${table}
          GROUP BY state_json->>'status'`,
        [now, ['leased', 'running']],
      );
      const counts = {};
      let total = 0;
      let claimable = 0;
      let expiredLeases = 0;
      for (const row of result.rows ?? []) {
        if (row.status) counts[row.status] = Number(row.status_count ?? 0);
        total += Number(row.status_count ?? 0);
        claimable += Number(row.claimable ?? 0);
        expiredLeases += Number(row.expired_leases ?? 0);
      }
      return { total, counts, claimable, expiredLeases };
    },

    async getAdmissionStats({
      tenantId = null,
      userId = null,
      now = Date.now(),
      windowMs = 60_000,
    } = {}) {
      const result = await pool.query(
        `SELECT
            COUNT(*) FILTER (
              WHERE state_json->>'status' = ANY($1::text[])
            )::BIGINT AS global_active,
            COUNT(*) FILTER (
              WHERE state_json->>'status' = ANY($1::text[])
                AND $2::text IS NOT NULL
                AND state_json->'request'->'subject'->>'tenantId' = $2
            )::BIGINT AS tenant_active,
            COUNT(*) FILTER (
              WHERE state_json->>'status' = ANY($1::text[])
                AND $3::text IS NOT NULL
                AND state_json->'request'->'subject'->>'userId' = $3
            )::BIGINT AS user_active,
            COUNT(*) FILTER (
              WHERE $2::text IS NOT NULL
                AND state_json->'request'->'subject'->>'tenantId' = $2
                AND created_at >= $4
            )::BIGINT AS tenant_recent,
            COUNT(*) FILTER (
              WHERE $3::text IS NOT NULL
                AND state_json->'request'->'subject'->>'userId' = $3
                AND created_at >= $4
            )::BIGINT AS user_recent
           FROM ${table}`,
        [
          ['queued', 'retryable', 'leased', 'running'],
          tenantId,
          userId,
          now - windowMs,
        ],
      );
      const row = result.rows?.[0] ?? {};
      return {
        globalActive: Number(row.global_active ?? 0),
        tenantActive: Number(row.tenant_active ?? 0),
        userActive: Number(row.user_active ?? 0),
        tenantRecent: Number(row.tenant_recent ?? 0),
        userRecent: Number(row.user_recent ?? 0),
      };
    },

    async recordWorkerHeartbeat(record) {
      await pool.query(
        `INSERT INTO ${workerTable} (worker_id, state_json, last_seen_at)
         VALUES ($1, $2::jsonb, $3)
         ON CONFLICT (worker_id) DO UPDATE
           SET state_json = EXCLUDED.state_json,
               last_seen_at = EXCLUDED.last_seen_at`,
        [record.workerId, JSON.stringify(record), record.lastSeenAt],
      );
      return structuredClone(record);
    },

    async listWorkers({ now = Date.now(), staleAfterMs = 60_000 } = {}) {
      const result = await pool.query(
        `SELECT state_json
           FROM ${workerTable}
          ORDER BY last_seen_at DESC
          LIMIT 500`,
      );
      return (result.rows ?? [])
        .map(projectRow)
        .filter(Boolean)
        .map((worker) => ({
          ...worker,
          stale: Number(worker.lastSeenAt ?? 0) < now - staleAfterMs,
        }));
    },

    async listEvents(jobId, { after = 0, limit = 500 } = {}) {
      const cursor = Math.max(0, Number(after) || 0);
      const pageSize = Math.min(500, Math.max(1, Number(limit) || 500));
      const result = await pool.query(
        `SELECT event_json
           FROM ${eventTable}
          WHERE job_id = $1 AND sequence > $2
          ORDER BY sequence ASC
          LIMIT $3`,
        [String(jobId ?? '').trim(), cursor, pageSize],
      );
      return (result.rows ?? []).map(projectEventRow).filter(Boolean);
    },
  };
}

export async function createExecutorJobPersistence({
  mode = process.env.MEMIND_ORCHESTRATOR_CHECKPOINT_MODE,
  connectionString = process.env.MEMIND_ORCHESTRATOR_DATABASE_URL,
  schema = process.env.MEMIND_ORCHESTRATOR_DATABASE_SCHEMA,
  poolFactory = (config) => new Pool(config),
} = {}) {
  const normalizedMode = normalizeMode(mode);
  if (normalizedMode === 'memory') {
    const store = createInMemoryExecutorJobStore();
    return {
      kind: 'memory',
      durable: false,
      store,
      async probe() {
        return true;
      },
      async close() {},
    };
  }

  const normalizedConnectionString = String(connectionString ?? '').trim();
  if (!normalizedConnectionString) {
    const error = new Error(
      'MEMIND_ORCHESTRATOR_DATABASE_URL is required for PostgreSQL Executor Job Store',
    );
    error.code = 'ORCHESTRATOR_DATABASE_URL_REQUIRED';
    throw error;
  }

  const normalizedSchema = normalizeSchema(schema);
  const pool = poolFactory({ connectionString: normalizedConnectionString });
  const store = createPostgresExecutorJobStore({ pool, schema: normalizedSchema });
  try {
    await store.ensureSchema();
  } catch (error) {
    try {
      await pool.end?.();
    } catch {
      // Preserve the schema initialization failure.
    }
    throw error;
  }
  return {
    kind: 'postgres',
    durable: true,
    store,
    async probe() {
      await pool.query('SELECT 1 AS ok');
      return true;
    },
    async close() {
      await pool.end?.();
    },
  };
}

export const executorJobStoreInternals = {
  DEFAULT_EXECUTOR_SCHEMA,
  EXECUTOR_JOB_TABLE,
  EXECUTOR_JOB_EVENT_TABLE,
  EXECUTOR_WORKER_TABLE,
  appendEvent,
  eventWithSequence,
  normalizeMode,
  normalizeSchema,
  parseState,
  projectEventRow,
  projectRow,
  withTransaction,
};