import { mkdir, open, readFile, rename } from 'node:fs/promises';
import { dirname } from 'node:path';
import {
  ContinuationPolicySchema,
  DEFAULT_CONTINUATION_POLICY,
  WorkerContinuityStateSchema,
  type Assignment,
  type ContinuationPolicy,
  type ConversationBinding,
  type ExecutionAttempt,
  type ExecutorCommandRecord,
  type Message,
  type OrchestratorEvent,
  type Run,
  type Worker,
  type WorkerContinuityState,
  type WorkerDisposition,
} from '@platform-modules/chatgpt-orchestrator-protocol';

export interface StoredState {
  runs: Record<string, Run>;
  workers: Record<string, Worker>;
  assignments: Record<string, Assignment>;
  messages: Record<string, Message>;
  events: Record<string, OrchestratorEvent[]>;
  idempotency: Record<string, { resourceId: string; fingerprint: string }>;
  nextCursor: Record<string, number>;
  executorCommands: Record<string, ExecutorCommandRecord>;
  conversations: Record<string, ConversationBinding>;
  runContinuationPolicies: Record<string, ContinuationPolicy>;
  workerContinuity: Record<string, WorkerContinuityState>;
  executionAttempts: Record<string, ExecutionAttempt>;
  executionSequences: Record<string, number>;
  executionIdempotency: Record<string, { executionId: string; fingerprint: string }>;
}

export function emptyState(): StoredState {
  return {
    runs: {}, workers: {}, assignments: {}, messages: {}, events: {}, idempotency: {}, nextCursor: {},
    executorCommands: {}, conversations: {}, runContinuationPolicies: {}, workerContinuity: {},
    executionAttempts: {}, executionSequences: {}, executionIdempotency: {},
  };
}

function defaultDisposition(worker: Worker): WorkerDisposition {
  if (['completed', 'failed', 'cancelled'].includes(worker.state)) return 'terminal';
  if (['created', 'launch_requested', 'launching', 'launch_failed', 'waiting'].includes(worker.state)) return 'awaiting_execution';
  return 'running';
}

export function defaultWorkerContinuity(worker: Worker): WorkerContinuityState {
  return {
    workerId: worker.workerId,
    disposition: defaultDisposition(worker),
    policyOverride: null,
    humanWait: null,
    dependencyWait: null,
    pauseReason: null,
    pauseAfterCurrentTurn: false,
    resumeAfterCurrentTurn: false,
    pendingExecutionId: null,
    idleGraceDueAt: null,
    activeKeepaliveDueAt: null,
    wakeNotBefore: null,
    pendingWakeReasons: [],
    lastContinuationSubmittedAt: null,
    lastRoutineSteerAt: null,
    deadmanRecoveryCount: 0,
    lastDeadmanRecoveryAt: null,
    lastProgressCursor: null,
    lastProgressAt: null,
    consecutiveExecutionAttemptsWithoutProgress: 0,
    stallWarningAt: null,
    stallPausedAt: null,
    updatedAt: worker.updatedAt,
  };
}

function normalizeWorkerContinuity(worker: Worker, raw: unknown): WorkerContinuityState {
  const defaults = defaultWorkerContinuity(worker);
  if (!raw || typeof raw !== 'object') return defaults;
  const candidate = raw as Partial<WorkerContinuityState>;
  const policyOverride = candidate.policyOverride == null ? null : ContinuationPolicySchema.parse(candidate.policyOverride);
  return WorkerContinuityStateSchema.parse({
    ...defaults,
    ...candidate,
    workerId: worker.workerId,
    policyOverride,
  });
}

export function normalizeStoredState(input: Partial<StoredState>): StoredState {
  const normalized: StoredState = { ...emptyState(), ...structuredClone(input) };
  for (const [runId, policy] of Object.entries(normalized.runContinuationPolicies)) {
    normalized.runContinuationPolicies[runId] = ContinuationPolicySchema.parse(policy);
  }
  for (const [messageId, rawMessage] of Object.entries(normalized.messages)) {
    const message = rawMessage as Message & Partial<Message>;
    normalized.messages[messageId] = {
      ...message,
      delivery: message.delivery ?? 'record_only',
      priority: message.priority ?? 'normal',
      idempotencyKey: message.idempotencyKey ?? null,
      deliveryState: message.deliveryState ?? (message.delivery && message.delivery !== 'record_only' ? 'pending' : 'recorded'),
      notBefore: message.notBefore ?? null,
      deliveryCommandId: message.deliveryCommandId ?? null,
      deliveryAttempts: message.deliveryAttempts ?? 0,
      deliveredAt: message.deliveredAt ?? null,
      lastDeliveryError: message.lastDeliveryError ?? null,
    };
  }
  for (const worker of Object.values(normalized.workers)) {
    normalized.workerContinuity[worker.workerId] = normalizeWorkerContinuity(worker, normalized.workerContinuity[worker.workerId]);
    const attempts = Object.values(normalized.executionAttempts).filter((attempt) => attempt.workerId === worker.workerId);
    for (const attempt of attempts) {
      normalized.executionAttempts[attempt.executionId] = { ...attempt, wakeReasons: attempt.wakeReasons ?? [] };
    }
    const maxSequence = attempts.reduce((maximum, attempt) => Math.max(maximum, attempt.sequence), 0);
    normalized.executionSequences[worker.workerId] = Math.max(normalized.executionSequences[worker.workerId] ?? 0, maxSequence);
  }
  return normalized;
}

export function effectiveContinuationPolicy(state: Readonly<StoredState>, workerId: string): ContinuationPolicy {
  const worker = state.workers[workerId];
  if (!worker) return DEFAULT_CONTINUATION_POLICY;
  return state.workerContinuity[workerId]?.policyOverride
    ?? state.runContinuationPolicies[worker.runId]
    ?? DEFAULT_CONTINUATION_POLICY;
}

export interface OrchestratorRepository {
  read<T>(reader: (state: Readonly<StoredState>) => T): Promise<T>;
  mutate<T>(mutation: (state: StoredState) => T): Promise<T>;
}

export class InMemoryRepository implements OrchestratorRepository {
  private state: StoredState;
  private tail: Promise<void> = Promise.resolve();

  constructor(initial: Partial<StoredState> = emptyState()) {
    this.state = normalizeStoredState(initial);
  }

  async read<T>(reader: (state: Readonly<StoredState>) => T): Promise<T> {
    await this.tail;
    return reader(structuredClone(this.state));
  }

  async mutate<T>(mutation: (state: StoredState) => T): Promise<T> {
    let release!: () => void;
    const previous = this.tail;
    this.tail = new Promise<void>((resolve) => { release = resolve; });
    await previous;
    try {
      const draft = structuredClone(this.state);
      const result = mutation(draft);
      this.state = draft;
      return result;
    } finally {
      release();
    }
  }
}

export function pruneTerminalRuns(state: StoredState, before: string): string[] {
  const cutoff = Date.parse(before);
  if (!Number.isFinite(cutoff)) throw new Error('Retention cutoff must be an ISO date-time.');
  const terminal = new Set(['completed', 'failed', 'cancelled']);
  const runIds = Object.values(state.runs)
    .filter((run) => terminal.has(run.state) && Date.parse(run.updatedAt) < cutoff)
    .map((run) => run.runId);
  if (runIds.length === 0) return [];

  const removedRuns = new Set(runIds);
  const removedWorkers = new Set(Object.values(state.workers)
    .filter((worker) => removedRuns.has(worker.runId))
    .map((worker) => worker.workerId));
  const removedAssignments = new Set(Object.values(state.workers)
    .filter((worker) => removedWorkers.has(worker.workerId))
    .map((worker) => worker.assignmentId));
  const removedExecutions = new Set(Object.values(state.executionAttempts)
    .filter((attempt) => removedWorkers.has(attempt.workerId))
    .map((attempt) => attempt.executionId));

  for (const runId of removedRuns) {
    delete state.runs[runId];
    delete state.events[runId];
    delete state.nextCursor[runId];
    delete state.runContinuationPolicies[runId];
  }
  for (const workerId of removedWorkers) {
    delete state.workers[workerId];
    delete state.conversations[workerId];
    delete state.workerContinuity[workerId];
    delete state.executionSequences[workerId];
  }
  for (const assignmentId of removedAssignments) delete state.assignments[assignmentId];
  for (const executionId of removedExecutions) delete state.executionAttempts[executionId];

  for (const [messageId, message] of Object.entries(state.messages)) {
    if (removedRuns.has(message.runId)) delete state.messages[messageId];
  }
  for (const [commandId, record] of Object.entries(state.executorCommands)) {
    if (removedWorkers.has(record.command.workerId)) delete state.executorCommands[commandId];
  }
  for (const [key, record] of Object.entries(state.idempotency)) {
    if (removedRuns.has(record.resourceId) || removedWorkers.has(record.resourceId) || removedAssignments.has(record.resourceId)) {
      delete state.idempotency[key];
    }
  }
  for (const [key, record] of Object.entries(state.executionIdempotency)) {
    if (removedExecutions.has(record.executionId)) delete state.executionIdempotency[key];
  }
  return runIds;
}

export class JsonFileRepository implements OrchestratorRepository {
  private readonly path: string;
  private loaded = false;
  private state: StoredState = emptyState();
  private tail: Promise<void> = Promise.resolve();

  constructor(path: string) {
    this.path = path;
  }

  private async ensureLoaded(): Promise<void> {
    if (this.loaded) return;
    try {
      const raw = await readFile(this.path, 'utf8');
      this.state = normalizeStoredState(JSON.parse(raw) as Partial<StoredState>);
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error;
      this.state = emptyState();
    }
    this.loaded = true;
  }

  private async atomicWrite(state: StoredState): Promise<void> {
    const directory = dirname(this.path);
    await mkdir(directory, { recursive: true, mode: 0o700 });
    const temporary = `${this.path}.tmp-${process.pid}-${crypto.randomUUID()}`;
    const handle = await open(temporary, 'wx', 0o600);
    try {
      await handle.writeFile(`${JSON.stringify(state)}\n`, 'utf8');
      await handle.sync();
    } finally {
      await handle.close();
    }
    await rename(temporary, this.path);
    const directoryHandle = await open(directory, 'r');
    try { await directoryHandle.sync(); } finally { await directoryHandle.close(); }
  }

  private async lock<T>(body: () => Promise<T>): Promise<T> {
    let release!: () => void;
    const previous = this.tail;
    this.tail = new Promise<void>((resolve) => { release = resolve; });
    await previous;
    try {
      return await body();
    } finally {
      release();
    }
  }

  async read<T>(reader: (state: Readonly<StoredState>) => T): Promise<T> {
    return this.lock(async () => {
      await this.ensureLoaded();
      return reader(structuredClone(this.state));
    });
  }

  async mutate<T>(mutation: (state: StoredState) => T): Promise<T> {
    return this.lock(async () => {
      await this.ensureLoaded();
      const draft = structuredClone(this.state);
      const result = mutation(draft);
      await this.atomicWrite(draft);
      this.state = draft;
      return result;
    });
  }
}