import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { applyGraphPatch, validateWorkflowState } from "../domain/apply-patch.ts";
import { createInitialState } from "../domain/initial-state.ts";
import {
  isWorkflowStateV1,
  isWorkflowStateV2,
  migrateCheckpointToLatest,
  snapshotHashV1,
} from "../domain/migration.ts";
import { canonicalJson, patchRequestHash, snapshotHash } from "../domain/snapshot.ts";
import type { CommittedBatch, GraphCheckpointDetails, WorkflowState } from "../domain/types.ts";

function isCheckpoint(value: unknown): value is GraphCheckpointDetails {
  if (!value || typeof value !== "object") return false;
  const details = value as Partial<GraphCheckpointDetails>;
  return details.kind === "intent-petri/graph-checkpoint" && typeof details.revision === "number";
}

export interface HistoricalBatch {
  patch: {
    schemaVersion: 1 | 2;
    patchId: string;
    baseRevision: number;
    idempotencyKey: string;
    correlationId?: string;
    trustedHumanEvidenceIds?: string[];
    ops: Array<Record<string, unknown> & { op: string }>;
  };
  revision: number;
  events: Array<{
    kind: string;
    opIndex: number;
    eventIndex: number;
    payload: Record<string, unknown>;
  }>;
  snapshotHash: string;
  committedAt: string;
}

function isHistoricalBatch(value: unknown): value is Omit<HistoricalBatch, "events"> & { events?: unknown[] } {
  if (!value || typeof value !== "object") return false;
  const batch = value as Partial<HistoricalBatch>;
  return (
    !!batch.patch &&
    typeof batch.patch === "object" &&
    (batch.patch.schemaVersion === 1 || batch.patch.schemaVersion === 2) &&
    typeof batch.patch.patchId === "string" &&
    typeof batch.patch.baseRevision === "number" &&
    typeof batch.patch.idempotencyKey === "string" &&
    Array.isArray(batch.patch.ops) &&
    typeof batch.revision === "number" &&
    typeof batch.snapshotHash === "string" &&
    typeof batch.committedAt === "string"
  );
}

function normalizeHistoricalBatch(value: unknown): HistoricalBatch | undefined {
  if (!isHistoricalBatch(value)) return undefined;
  const rawEvents = Array.isArray(value.events) ? value.events : [];
  const counters = new Map<number, number>();
  const events = rawEvents.map((raw) => {
    if (!raw || typeof raw !== "object") throw new Error("invalid historical event");
    const item = raw as Record<string, unknown>;
    if (typeof item.kind !== "string" || typeof item.opIndex !== "number") {
      throw new Error("invalid historical event envelope");
    }
    const next = counters.get(item.opIndex) ?? 0;
    const eventIndex = typeof item.eventIndex === "number" ? item.eventIndex : next;
    counters.set(item.opIndex, eventIndex + 1);
    return {
      kind: item.kind,
      opIndex: item.opIndex,
      eventIndex,
      payload: item.payload && typeof item.payload === "object" ? (item.payload as Record<string, unknown>) : {},
    };
  });
  return {
    patch: structuredClone(value.patch),
    revision: value.revision,
    events,
    snapshotHash: value.snapshotHash,
    committedAt: value.committedAt,
  };
}

function validateEventOrdering(batch: HistoricalBatch): void {
  let previousOp = -1;
  const nextEventIndex = new Map<number, number>();
  for (const event of batch.events) {
    if (!Number.isInteger(event.opIndex) || event.opIndex < 0 || event.opIndex >= batch.patch.ops.length) {
      throw new Error("event opIndex is outside the patch");
    }
    if (event.opIndex < previousOp) throw new Error("events are not ordered by opIndex");
    previousOp = event.opIndex;
    const expected = nextEventIndex.get(event.opIndex) ?? 0;
    if (event.eventIndex !== expected) throw new Error("events are not contiguous by eventIndex");
    nextEventIndex.set(event.opIndex, expected + 1);
  }
}

function legacyRequestPatch(batch: HistoricalBatch): HistoricalBatch["patch"] {
  const patch = structuredClone(batch.patch);
  for (const op of patch.ops) {
    if (op.op !== "attach_evidence") continue;
    const evidence = op.evidence;
    if (evidence && typeof evidence === "object") delete (evidence as Record<string, unknown>).redacted;
  }
  return patch;
}

function requestHashMatches(
  recordHash: string,
  batch: HistoricalBatch,
  sourceSchemaVersion: 1 | 2,
): boolean {
  if (recordHash === patchRequestHash(batch.patch as never)) return true;
  return sourceSchemaVersion === 1 && recordHash === patchRequestHash(legacyRequestPatch(batch) as never);
}

function validateEnvelope(
  details: GraphCheckpointDetails,
  sourceState: WorkflowState | Parameters<typeof snapshotHashV1>[0],
  batch: HistoricalBatch | undefined,
  previousState: WorkflowState,
  sourceSchemaVersion: 1 | 2,
): void {
  if (details.status !== "applied") throw new Error("state-bearing checkpoint must be applied");
  if (details.revision !== sourceState.revision) throw new Error("checkpoint revision differs from state revision");
  if (sourceState.revision !== previousState.revision + 1) throw new Error("checkpoint revision is not contiguous");
  if (!batch) throw new Error(`schema ${sourceSchemaVersion} checkpoint is missing its committed batch`);
  if (details.patchId !== batch.patch.patchId) throw new Error("checkpoint patchId differs from batch");
  if (details.revision !== batch.revision) throw new Error("checkpoint revision differs from batch");
  if (details.snapshotHash !== batch.snapshotHash) throw new Error("checkpoint hash differs from batch");
  if (batch.patch.baseRevision !== previousState.revision) throw new Error("batch baseRevision breaks checkpoint continuity");
  const record = sourceState.appliedPatches[batch.patch.idempotencyKey];
  if (!record) throw new Error("state is missing the batch AppliedPatchRecord");
  if (
    record.patchId !== batch.patch.patchId ||
    record.revision !== batch.revision ||
    record.snapshotHash !== batch.snapshotHash ||
    !requestHashMatches(record.requestHash, batch, sourceSchemaVersion)
  ) {
    throw new Error("AppliedPatchRecord does not bind the checkpoint batch");
  }
  validateEventOrdering(batch);
}

export interface ReconstructedCheckpoint {
  entryId: string;
  parentEntryId?: string;
  toolCallId: string;
  state: WorkflowState;
  batch?: CommittedBatch | HistoricalBatch;
  sourceSchemaVersion: 1 | 2;
  migratedFromSchemaVersion?: 1;
}

export interface ReconstructionIssue {
  entryId: string;
  revision?: number;
  kind: "invalid" | "divergent";
  reason: string;
}

export interface ReconstructionResult {
  state: WorkflowState;
  history: ReconstructedCheckpoint[];
  checkpoints: number;
  invalidCheckpoints: number;
  ignoredDivergentCheckpoints: number;
  migratedCheckpoints: number;
  issues: ReconstructionIssue[];
  hashVerified: boolean;
  branchHeadEntryId?: string;
}

export function reconstructFromBranch(sessionManager: ExtensionContext["sessionManager"]): ReconstructionResult {
  let state = createInitialState();
  const history: ReconstructedCheckpoint[] = [];
  let checkpoints = 0;
  let invalidCheckpoints = 0;
  let ignoredDivergentCheckpoints = 0;
  let migratedCheckpoints = 0;
  let chainBroken = false;
  const issues: ReconstructionIssue[] = [];

  for (const entry of sessionManager.getBranch()) {
    if (entry.type !== "message") continue;
    const message = entry.message;
    if (message.role !== "toolResult" || message.toolName !== "update_action_path") continue;
    const details = message.details;
    if (!isCheckpoint(details) || !details.state) continue;
    checkpoints += 1;
    if (details.revision < state.revision) {
      ignoredDivergentCheckpoints += 1;
      issues.push({
        entryId: entry.id,
        revision: details.revision,
        kind: "divergent",
        reason: `Ignored rollback descendant r${details.revision}; verified authority is already r${state.revision}`,
      });
      continue;
    }
    if (chainBroken) {
      invalidCheckpoints += 1;
      issues.push({
        entryId: entry.id,
        revision: details.revision,
        kind: "invalid",
        reason: "Skipped after an earlier invalid checkpoint broke strict continuity",
      });
      continue;
    }

    try {
      let sourceSchemaVersion: 1 | 2;
      if (isWorkflowStateV2(details.state)) {
        sourceSchemaVersion = 2;
        if (snapshotHash(details.state) !== details.snapshotHash) throw new Error("schema 2 snapshot hash mismatch");
      } else if (isWorkflowStateV1(details.state)) {
        sourceSchemaVersion = 1;
        if (snapshotHashV1(details.state) !== details.snapshotHash) throw new Error("schema 1 snapshot hash mismatch");
      } else {
        throw new Error("unsupported workflow state schema");
      }

      const batch = normalizeHistoricalBatch(details.batch);
      validateEnvelope(details, details.state, batch, state, sourceSchemaVersion);
      const migrated = migrateCheckpointToLatest(details.state);
      if (!migrated) throw new Error("checkpoint migration failed");
      validateWorkflowState(migrated.state);

      if (sourceSchemaVersion === 2) {
        if (!batch || batch.patch.schemaVersion !== 2) throw new Error("schema 2 state requires a schema 2 batch");
        const replayed = applyGraphPatch(
          state,
          batch.patch as unknown as CommittedBatch["patch"],
          batch.committedAt,
        );
        if (replayed.status !== "applied") throw new Error(`batch replay failed: ${replayed.status}`);
        if (replayed.batch.snapshotHash !== details.snapshotHash) throw new Error("batch replay snapshot differs");
        if (canonicalJson(replayed.batch.events) !== canonicalJson(batch.events)) {
          throw new Error("batch events differ from deterministic replay");
        }
      }

      if (migrated.migratedFromSchemaVersion) migratedCheckpoints += 1;
      state = migrated.state;
      history.push({
        entryId: entry.id,
        ...(entry.parentId ? { parentEntryId: entry.parentId } : {}),
        toolCallId: message.toolCallId,
        state: structuredClone(state),
        ...(batch ? { batch: structuredClone(batch) } : {}),
        sourceSchemaVersion,
        ...(migrated.migratedFromSchemaVersion
          ? { migratedFromSchemaVersion: migrated.migratedFromSchemaVersion }
          : {}),
      });
    } catch (error) {
      invalidCheckpoints += 1;
      chainBroken = true;
      issues.push({
        entryId: entry.id,
        revision: details.revision,
        kind: "invalid",
        reason: error instanceof Error ? error.message : String(error),
      });
    }
  }

  const leafId = typeof sessionManager.getLeafId === "function" ? sessionManager.getLeafId() : undefined;
  return {
    state,
    history,
    checkpoints,
    invalidCheckpoints,
    ignoredDivergentCheckpoints,
    migratedCheckpoints,
    issues,
    hashVerified: invalidCheckpoints === 0,
    ...(leafId ? { branchHeadEntryId: leafId } : {}),
  };
}