import type {
  AgentRun,
  AgentRunAttempt,
  AgentRunStep,
  AttemptId,
  ChildTaskRecord,
  HostStatus,
  PromptAssemblySnapshot,
  ProviderTurnTraceSnapshot,
  RoutingPressureSnapshot,
  RunId,
  ThreadDiagnostics,
  ThreadId,
  ThreadMetadata,
  ThreadRunDiagnostics,
  TurnId,
  VerdictRecord,
  WorkflowName,
  WorkflowPhaseName,
  WorkerLeaseStatus,
} from "@best-agent/protocol";
import type { ApprovalStore } from "../core/approval-store.js";
import { buildWorkflowChildTaskGraph, collectWorkflowTurnStates } from "../core/workflow-graph.js";
import type { RunCheckpoint, ThreadStore } from "../core/thread-store.js";

type DiagnosticsCollectorDeps = {
  approvalStore: ApprovalStore;
  buildRoutingPressureSnapshot(input: {
    attemptId?: AttemptId;
    childTaskFanout: number;
    createdAt: string;
    diagnostics: ThreadDiagnostics;
    runId?: RunId;
    threadId: ThreadId;
    turnId?: TurnId;
  }): RoutingPressureSnapshot;
  computeRunSuspectReason(
    run: AgentRun,
    latestAttempt?: AgentRunAttempt,
    checkpoint?: RunCheckpoint | null,
    workerLeases?: Array<{ kind: string; status: WorkerLeaseStatus; updatedAt: string }>,
  ): string | undefined;
  readHostStatus(): Promise<HostStatus>;
  readProfileAllowedTools(profileName: string): string[] | undefined;
  readPromptAssembly(threadId: ThreadId): PromptAssemblySnapshot | undefined;
  readProviderTurnTrace(threadId: ThreadId): ProviderTurnTraceSnapshot | undefined;
  readTranscriptStats(threadId: ThreadId): ThreadDiagnostics["transcript"];
  readWorkspaceDiagnostics(threadId: ThreadId): ThreadDiagnostics["workspace"];
  runStepFromCheckpoint(threadId: ThreadId, checkpoint: RunCheckpoint): AgentRunStep;
  store: ThreadStore;
};

export class DiagnosticsCollector {
  readonly #deps: DiagnosticsCollectorDeps;

  constructor(deps: DiagnosticsCollectorDeps) {
    this.#deps = deps;
  }

  async decorateThreadDiagnostics(
    threadId: ThreadId,
    diagnostics: ThreadDiagnostics,
  ): Promise<ThreadDiagnostics> {
    const approvals = await this.#collectThreadApprovalDiagnostics(threadId);
    const childTasks = await this.#collectThreadChildTaskDiagnostics(threadId);
    const host = await this.#deps.readHostStatus();
    const routing = await this.#collectThreadRoutingDiagnostics(threadId);
    const runs = await this.#collectThreadRunDiagnostics(threadId);
    const projections = await this.#refreshThreadProjectionSnapshots(threadId, runs);
    const workflow = await this.#collectThreadWorkflowDiagnostics(threadId);
    const workers = await this.#collectThreadWorkerDiagnostics(threadId);
    const latestRun = runs.at(-1);
    const latestRoutingPressure =
      routing.records.at(-1)?.pressureSummary ??
      this.#deps.buildRoutingPressureSnapshot({
        attemptId: latestRun?.latestAttempt?.id,
        childTaskFanout: childTasks.activeCount,
        createdAt: new Date().toISOString(),
        diagnostics,
        runId: latestRun?.run.id,
        threadId,
        turnId: latestRun?.run.turnId,
      });
    return {
      ...diagnostics,
      approvals,
      childTasks,
      host,
      projections,
      promptAssembly: this.#deps.readPromptAssembly(threadId),
      providerTurnTrace: this.#deps.readProviderTurnTrace(threadId),
      routing,
      routingPressure: latestRoutingPressure,
      runs,
      transcript: this.#deps.readTranscriptStats(threadId),
      workers,
      workflow,
      workspace: this.#deps.readWorkspaceDiagnostics(threadId),
    };
  }

  async #refreshThreadProjectionSnapshots(
    threadId: ThreadId,
    runs: ThreadRunDiagnostics[],
  ): Promise<ThreadDiagnostics["projections"]> {
    const latestSequence = await this.#deps.store.getLatestEventSequence(threadId);
    const existingSnapshots = await this.#deps.store.listProjectionSnapshots(threadId);
    const existingByKind = new Map(existingSnapshots.map((snapshot) => [snapshot.kind, snapshot]));
    const degradedByKind = new Map<string, ThreadDiagnostics["projections"]["snapshots"][number]>();
    const desiredSnapshots = [
      {
        kind: "thread-events" as const,
        sequence: latestSequence,
        snapshot: {
          latestSequence,
        },
      },
      {
        kind: "thread-runs" as const,
        sequence: latestSequence,
        snapshot: {
          latestRunId: runs[0]?.run.id,
          runCount: runs.length,
          suspectRunCount: runs.filter((run) => run.suspect).length,
        },
      },
    ];

    for (const desired of desiredSnapshots) {
      try {
        await this.#deps.store.saveProjectionSnapshot({
          kind: desired.kind,
          sequence: desired.sequence,
          snapshot: desired.snapshot,
          threadId,
        });
      } catch {
        degradedByKind.set(
          desired.kind,
          await this.#persistDegradedProjectionSnapshot(
            threadId,
            desired,
            existingByKind.get(desired.kind),
          ),
        );
      }
    }

    const snapshots = await this.#deps.store.listProjectionSnapshots(threadId);
    const latestByKind: ThreadDiagnostics["projections"]["latestByKind"] = {};
    for (const snapshot of snapshots) {
      if (!latestByKind[snapshot.kind]) {
        latestByKind[snapshot.kind] = snapshot;
      }
    }

    for (const [kind, degraded] of degradedByKind) {
      latestByKind[kind as keyof typeof latestByKind] = degraded;
      const existingIndex = snapshots.findIndex((snapshot) => snapshot.kind === degraded.kind);
      if (existingIndex >= 0) {
        snapshots[existingIndex] = degraded;
      } else {
        snapshots.push(degraded);
      }
    }

    return {
      latestByKind,
      snapshots,
    };
  }

  async #persistDegradedProjectionSnapshot(
    threadId: ThreadId,
    desired: {
      kind: "thread-events" | "thread-runs";
      sequence: number;
      snapshot: ThreadMetadata;
    },
    previous?: ThreadDiagnostics["projections"]["snapshots"][number],
  ): Promise<ThreadDiagnostics["projections"]["snapshots"][number]> {
    const degradedSequence = previous?.sequence ?? desired.sequence;
    const degradedSnapshot = previous?.snapshot ?? desired.snapshot;
    try {
      return await this.#deps.store.saveProjectionSnapshot({
        kind: desired.kind,
        sequence: degradedSequence,
        snapshot: degradedSnapshot,
        status: "degraded",
        threadId,
      });
    } catch {
      const timestamp = new Date().toISOString();
      return {
        createdAt: previous?.createdAt ?? timestamp,
        kind: desired.kind,
        schemaVersion: previous?.schemaVersion ?? 1,
        sequence: degradedSequence,
        snapshot: degradedSnapshot,
        status: "degraded",
        threadId,
        updatedAt: timestamp,
      };
    }
  }

  async #collectThreadChildTaskDiagnostics(
    threadId: ThreadId,
  ): Promise<ThreadDiagnostics["childTasks"]> {
    const records = await this.#deps.store.listChildTaskRecords(threadId, 50);
    const effectiveVerdicts = await this.#deps.store.getEffectiveVerdicts(
      threadId,
      Array.from(new Set(records.map((record) => record.childTaskId))),
    );
    const decoratedRecords = records.map((record) => {
      const effectiveVerdict = effectiveVerdicts.get(record.childTaskId);
      return effectiveVerdict ? { ...record, effectiveVerdict } : record;
    });
    const latestById: ThreadDiagnostics["childTasks"]["latestById"] = {};
    for (const record of decoratedRecords) {
      if (!latestById[record.childTaskId]) {
        latestById[record.childTaskId] = record;
      }
    }
    const latestRecords = Object.values(latestById);
    return {
      activeCount: latestRecords.filter(
        (record) => record.status === "running" || record.status === "completed",
      ).length,
      latestById,
      records: [...decoratedRecords].reverse(),
    };
  }

  async #collectThreadApprovalDiagnostics(
    threadId: ThreadId,
  ): Promise<ThreadDiagnostics["approvals"]> {
    const records = await this.#deps.approvalStore.listReviews(threadId, 50);
    const latestById: ThreadDiagnostics["approvals"]["latestById"] = {};
    for (const record of records) {
      if (!latestById[record.id]) {
        latestById[record.id] = record;
      }
    }
    const latestRecords = Object.values(latestById);
    return {
      activeCount: latestRecords.filter((record) => record.status === "pending").length,
      latestById,
      records: [...records].reverse(),
    };
  }

  async #collectThreadWorkflowDiagnostics(
    threadId: ThreadId,
  ): Promise<ThreadDiagnostics["workflow"]> {
    const records = await this.#deps.store.listChildTaskRecords(threadId, 100);
    const effectiveVerdicts = await this.#deps.store.getEffectiveVerdicts(
      threadId,
      Array.from(new Set(records.map((record) => record.childTaskId))),
    );
    const workflowTurns = await collectWorkflowTurnStates({
      effectiveVerdicts,
      records,
    });

    const latestByTurnId: ThreadDiagnostics["workflow"]["latestByTurnId"] = {};
    let activeCount = 0;
    for (const workflowTurn of workflowTurns) {
      const graph = buildWorkflowDiagnosticsGraph({
        effectiveVerdicts: workflowTurn.effectiveVerdicts,
        phases: workflowTurn.phases,
        readProfileAllowedTools: this.#deps.readProfileAllowedTools,
        workflowName: workflowTurn.workflowName,
      });
      if (graph.nodes.some((node) => node.status === "running" || node.status === "completed")) {
        activeCount += 1;
      }
      latestByTurnId[workflowTurn.turnId] = {
        graph,
        name: workflowTurn.workflowName,
        turnId: workflowTurn.turnId,
      };
    }

    return {
      activeCount,
      latestByTurnId,
      records: [...records].reverse(),
    };
  }

  async #collectThreadRoutingDiagnostics(
    threadId: ThreadId,
  ): Promise<ThreadDiagnostics["routing"]> {
    const records = await this.#deps.store.listRoutingRecords(threadId, 20);
    const latestByRole: ThreadDiagnostics["routing"]["latestByRole"] = {};
    for (const record of records) {
      if (!latestByRole[record.routingRole]) {
        latestByRole[record.routingRole] = record;
      }
    }
    return {
      latestByRole,
      records: [...records].reverse(),
    };
  }

  async #collectThreadWorkerDiagnostics(threadId: ThreadId): Promise<ThreadDiagnostics["workers"]> {
    const leases = await this.#deps.store.listWorkerLeases(threadId, 50);
    const latestById: ThreadDiagnostics["workers"]["latestById"] = {};
    for (const lease of leases) {
      if (!latestById[lease.id]) {
        latestById[lease.id] = lease;
      }
    }
    return {
      activeCount: leases.filter((lease) => lease.status === "running").length,
      latestById,
      leases: [...leases].reverse(),
    };
  }

  async #collectThreadRunDiagnostics(threadId: ThreadId): Promise<ThreadRunDiagnostics[]> {
    const runs = await this.#deps.store.listRuns(threadId);
    const workerLeases = await this.#deps.store.listWorkerLeases(threadId, 200);
    const diagnostics = await Promise.all(
      runs.map(async (run) => {
        const latestAttempt = (await this.#deps.store.listRunAttempts(run.id)).at(-1);
        const checkpoint = await this.#deps.store.getPendingRunCheckpointByTurn(run.turnId);
        const effectiveVerdict = run.nodeId
          ? ((await this.#deps.store.getEffectiveVerdict(threadId, run.nodeId)) ?? undefined)
          : undefined;
        const suspectReason = this.#deps.computeRunSuspectReason(
          run,
          latestAttempt,
          checkpoint,
          workerLeases.filter((lease) => lease.runId === run.id),
        );
        return {
          ...(effectiveVerdict ? { effectiveVerdict } : {}),
          latestAttempt,
          pendingStep: checkpoint
            ? this.#deps.runStepFromCheckpoint(threadId, checkpoint)
            : undefined,
          run,
          suspect: Boolean(suspectReason),
          ...(suspectReason ? { suspectReason } : {}),
        } satisfies ThreadRunDiagnostics;
      }),
    );
    return diagnostics.sort(
      (left, right) => Date.parse(right.run.updatedAt) - Date.parse(left.run.updatedAt),
    );
  }
}

function buildWorkflowDiagnosticsGraph(input: {
  effectiveVerdicts: ReadonlyMap<string, VerdictRecord>;
  phases: ReadonlyMap<
    WorkflowPhaseName,
    ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
  >;
  readProfileAllowedTools: (profileName: string) => string[] | undefined;
  workflowName: WorkflowName;
}): {
  edges: Array<{ from: WorkflowPhaseName; to: WorkflowPhaseName }>;
  nodes: Array<{
    allowedTools: string[];
    childThreadId?: ThreadId;
    correlationId?: string;
    effectiveVerdict?: VerdictRecord;
    nodeId?: string;
    phaseName: WorkflowPhaseName;
    profileName: string;
    status: ChildTaskRecord["status"] | "pending";
  }>;
} {
  return buildWorkflowChildTaskGraph({
    effectiveVerdicts: input.effectiveVerdicts,
    phases: input.phases,
    readProfileAllowedTools: input.readProfileAllowedTools,
    workflowName: input.workflowName,
  });
}
