import type { PostgresClient, PostgresPool } from '@qinglong/runtime-core';

import { PostgresModelInvocationRepository } from '../../../model-invocation/postgres-model-invocation-repository/repository';
import { completeWithAtomicOutputOperation } from '../../../model-invocation/postgres-model-invocation-repository/completionOperations';
import {
  normalizeModelInvocationCompletionRecord,
  type ModelInvocationCompletionCommand,
  type ModelInvocationCompletionRecord,
} from '../../../model-invocation/modelInvocation';
import { normalizeCopilotFailureDiagnosisExecutionPlan } from '../admission/plan';
import type { CopilotFailureDiagnosisExecutionPlan } from '../admission/contracts';
import {
  assertCopilotFailureDiagnosisOutputCompletionBinding,
  type CommitCopilotFailureDiagnosisOutputResult,
  type CopilotFailureDiagnosisOutputCompletionRepository,
} from './completion';
import {
  CopilotFailureDiagnosisOutputArtifactConflictError,
  CopilotFailureDiagnosisOutputArtifactUnavailableError,
  copilotFailureDiagnosisOutputArtifactIdentity,
  normalizeCopilotFailureDiagnosisOutputArtifact,
  type CopilotFailureDiagnosisOutputArtifact,
} from './outputArtifact';
import {
  CopilotFailureDiagnosisFinalizationConflictError,
  CopilotFailureDiagnosisFinalizationUnavailableError,
  CopilotFailureDiagnosisModelExecutionInProgressError,
  CopilotFailureDiagnosisModelResolutionRequiredError,
  createCopilotFailureDiagnosisFinalizationReceipt,
  normalizeCopilotFailureDiagnosisFinalizationReceipt,
  type CopilotFailureDiagnosisFinalizationReceipt,
  type CopilotFailureDiagnosisFinalizationRepository,
  type CopilotFailureDiagnosisFinalOutcome,
} from './finalization';

const TABLE = '"ql3_ai"."copilot_failure_diagnosis_model_outputs"';
const FINALIZATION_TABLE =
  '"ql3_ai"."copilot_failure_diagnosis_finalizations"';

interface OutputRow extends Record<string, unknown> {
  readonly artifactJson: unknown;
}

function integer(value: unknown): number {
  if (typeof value === 'number' && Number.isSafeInteger(value) && value >= 0) {
    return value;
  }
  if (typeof value === 'string' && /^(0|[1-9]\d*)$/.test(value)) {
    const parsed = Number(value);
    if (Number.isSafeInteger(parsed)) return parsed;
  }
  throw new CopilotFailureDiagnosisFinalizationConflictError();
}

function object(value: unknown): Record<string, unknown> {
  if (value && typeof value === 'object' && !Array.isArray(value)) {
    return value as Record<string, unknown>;
  }
  if (typeof value === 'string') {
    try {
      return JSON.parse(value) as Record<string, unknown>;
    } catch {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
  }
  throw new CopilotFailureDiagnosisFinalizationConflictError();
}

function string(value: unknown): string {
  if (typeof value !== 'string') {
    throw new CopilotFailureDiagnosisFinalizationConflictError();
  }
  return value;
}

function unavailable(cause?: unknown): never {
  throw new CopilotFailureDiagnosisOutputArtifactUnavailableError({
    cause: cause instanceof Error ? cause : undefined,
  });
}

function parse(row: OutputRow): Readonly<CopilotFailureDiagnosisOutputArtifact> {
  try {
    return normalizeCopilotFailureDiagnosisOutputArtifact(
      row.artifactJson as CopilotFailureDiagnosisOutputArtifact,
    );
  } catch (cause) {
    return unavailable(cause);
  }
}

async function read(
  queryable: Pick<PostgresPool, 'query'> | Pick<PostgresClient, 'query'>,
  artifactId: string,
): Promise<Readonly<CopilotFailureDiagnosisOutputArtifact> | null> {
  if (!/^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/.test(artifactId)) {
    throw new CopilotFailureDiagnosisOutputArtifactConflictError();
  }
  try {
    const result = await queryable.query<OutputRow>(
      `SELECT artifact_json AS "artifactJson"
         FROM ${TABLE}
        WHERE artifact_id = $1`,
      [artifactId],
    );
    if (result.rows.length > 1) {
      throw new CopilotFailureDiagnosisOutputArtifactConflictError();
    }
    return result.rows[0] ? parse(result.rows[0]) : null;
  } catch (cause) {
    if (cause instanceof CopilotFailureDiagnosisOutputArtifactConflictError) {
      throw cause;
    }
    return unavailable(cause);
  }
}

async function put(
  client: PostgresClient,
  artifactValue: CopilotFailureDiagnosisOutputArtifact,
): Promise<Readonly<CopilotFailureDiagnosisOutputArtifact>> {
  const artifact = normalizeCopilotFailureDiagnosisOutputArtifact(artifactValue);
  const existing = await read(client, artifact.artifactId);
  if (existing) {
    if (JSON.stringify(existing) !== JSON.stringify(artifact)) {
      throw new CopilotFailureDiagnosisOutputArtifactConflictError();
    }
    return existing;
  }
  try {
    await client.query(
      `INSERT INTO ${TABLE} (
         artifact_id, request_id, plan_digest, tool_completion_digest,
         project_id, run_id, step_run_id, invocation_id, provider, model,
         egress_evidence_digest, content_digest, output_bytes, key_id,
         algorithm, sealed_at_ms, artifact_digest, artifact_json
       ) VALUES (
         $1, $2, $3, $4, $5, $6, $7, $8, $9, $10,
         $11, $12, $13, $14, $15, $16, $17, $18::jsonb
       )`,
      [
        artifact.artifactId,
        artifact.requestId,
        artifact.planDigest,
        artifact.toolCompletionDigest,
        artifact.projectId,
        artifact.runId,
        artifact.stepRunId,
        artifact.invocationId,
        artifact.provider,
        artifact.model,
        artifact.egressEvidenceDigest,
        artifact.contentDigest,
        artifact.outputBytes,
        artifact.keyId,
        artifact.algorithm,
        artifact.sealedAtMs,
        artifact.artifactDigest,
        JSON.stringify(artifact),
      ],
    );
    return artifact;
  } catch (cause) {
    return unavailable(cause);
  }
}

export class PostgresCopilotFailureDiagnosisModelRepository
  extends PostgresModelInvocationRepository
  implements
    CopilotFailureDiagnosisOutputCompletionRepository,
    CopilotFailureDiagnosisFinalizationRepository
{
  readonly #pool: PostgresPool;

  constructor(pool: PostgresPool) {
    super(pool);
    this.#pool = pool;
  }

  findCopilotFailureDiagnosisOutput(
    artifactId: string,
  ): Promise<Readonly<CopilotFailureDiagnosisOutputArtifact> | null> {
    return read(this.#pool, artifactId);
  }

  async completeWithCopilotFailureDiagnosisOutput(
    commandValue: Readonly<ModelInvocationCompletionCommand>,
    artifactValue: Readonly<CopilotFailureDiagnosisOutputArtifact>,
  ): Promise<Readonly<CommitCopilotFailureDiagnosisOutputResult>> {
    const binding = assertCopilotFailureDiagnosisOutputCompletionBinding(
      commandValue,
      artifactValue,
    );
    const result = await completeWithAtomicOutputOperation(
      this.#pool,
      commandValue,
      {
        artifact: binding.artifact,
        reference: binding.reference,
        read: (client) => read(client, binding.artifact.artifactId),
        put: (client) => put(client, binding.artifact),
        matches: (stored) =>
          JSON.stringify(stored) === JSON.stringify(binding.artifact),
      },
    );
    return Object.freeze({
      status: result.status,
      reference: result.reference,
    });
  }

  async findFinalization(
    requestId: string,
  ): Promise<Readonly<CopilotFailureDiagnosisFinalizationReceipt> | null> {
    if (!/^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/.test(requestId)) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    try {
      const result = await this.#pool.query<Record<string, unknown>>(
        `SELECT request_id AS "requestId", plan_digest AS "planDigest",
                run_id AS "runId", model_step_run_id AS "modelStepRunId",
                invocation_id AS "invocationId",
                completion_digest AS "completionDigest", outcome,
                output_artifact_id AS "outputArtifactId",
                final_run_version AS "finalRunVersion",
                final_run_event_sequence AS "finalRunEventSequence",
                run_event_id AS "runEventId",
                finalized_at_ms AS "finalizedAtMs",
                receipt_digest AS "receiptDigest",
                receipt_json AS "receiptJson"
           FROM ${FINALIZATION_TABLE}
          WHERE request_id = $1`,
        [requestId],
      );
      if (result.rows.length > 1) {
        throw new CopilotFailureDiagnosisFinalizationConflictError();
      }
      if (!result.rows[0]) return null;
      const row = result.rows[0];
      const receipt = normalizeCopilotFailureDiagnosisFinalizationReceipt(
        object(row.receiptJson) as unknown as CopilotFailureDiagnosisFinalizationReceipt,
      );
      if (
        receipt.requestId !== string(row.requestId) ||
        receipt.planDigest !== string(row.planDigest) ||
        receipt.runId !== string(row.runId) ||
        receipt.modelStepRunId !== string(row.modelStepRunId) ||
        receipt.invocationId !== string(row.invocationId) ||
        receipt.completionDigest !== string(row.completionDigest) ||
        receipt.outcome !== string(row.outcome) ||
        receipt.outputArtifactId !== row.outputArtifactId ||
        receipt.finalRunVersion !== integer(row.finalRunVersion) ||
        receipt.finalRunEventSequence !== integer(row.finalRunEventSequence) ||
        receipt.runEventId !== string(row.runEventId) ||
        receipt.finalizedAtMs !== integer(row.finalizedAtMs) ||
        receipt.receiptDigest !== string(row.receiptDigest)
      ) {
        throw new CopilotFailureDiagnosisFinalizationConflictError();
      }
      const durable = await this.#pool.query<Record<string, unknown>>(
        `SELECT run.status AS "runStatus", run.version AS "runVersion",
                run.event_sequence AS "runEventSequence",
                run.finished_at_ms AS "finishedAtMs",
                run.output_ref AS "outputRef",
                step.status AS "stepStatus",
                event.type AS "eventType", event.dedupe_key AS "dedupeKey",
                event.step_run_id AS "eventStepRunId",
                event.payload, event.created_at_ms AS "eventCreatedAtMs",
                completion.completion_digest AS "completionDigest",
                completion.outcome AS "completionOutcome"
           FROM "ql3"."runs" AS run
           JOIN "ql3"."step_runs" AS step
             ON step.run_id = run.id AND step.id = $1
           JOIN "ql3"."run_events" AS event
             ON event.run_id = run.id AND event.id = $2
           JOIN "ql3_ai"."model_invocation_completions" AS completion
             ON completion.invocation_id = $3
          WHERE run.id = $4`,
        [
          receipt.modelStepRunId,
          receipt.runEventId,
          receipt.invocationId,
          receipt.runId,
        ],
      );
      if (durable.rows.length !== 1) {
        throw new CopilotFailureDiagnosisFinalizationConflictError();
      }
      const proof = durable.rows[0]!;
      const payload = object(proof.payload);
      if (
        string(proof.runStatus) !== receipt.outcome ||
        integer(proof.runVersion) !== receipt.finalRunVersion ||
        integer(proof.runEventSequence) !== receipt.finalRunEventSequence ||
        integer(proof.finishedAtMs) !== receipt.finalizedAtMs ||
        proof.outputRef !== receipt.outputArtifactId ||
        string(proof.stepStatus) !== receipt.outcome ||
        string(proof.eventType) !== `copilot.diagnosis.${receipt.outcome}` ||
        string(proof.dedupeKey) !== receipt.runEventId ||
        string(proof.eventStepRunId) !== receipt.modelStepRunId ||
        integer(proof.eventCreatedAtMs) !== receipt.finalizedAtMs ||
        string(proof.completionDigest) !== receipt.completionDigest ||
        string(proof.completionOutcome) !== receipt.outcome ||
        payload.requestId !== receipt.requestId ||
        payload.planDigest !== receipt.planDigest ||
        payload.invocationId !== receipt.invocationId ||
        payload.completionDigest !== receipt.completionDigest ||
        payload.outcome !== receipt.outcome ||
        payload.outputArtifactId !== receipt.outputArtifactId
      ) {
        throw new CopilotFailureDiagnosisFinalizationConflictError();
      }
      return receipt;
    } catch (cause) {
      if (cause instanceof CopilotFailureDiagnosisFinalizationConflictError) {
        throw cause;
      }
      throw new CopilotFailureDiagnosisFinalizationUnavailableError({ cause });
    }
  }

  async finalize(requestId: string): Promise<Readonly<{
    status: 'created' | 'existing';
    receipt: Readonly<CopilotFailureDiagnosisFinalizationReceipt>;
  }>> {
    const existing = await this.findFinalization(requestId);
    if (existing) {
      return Object.freeze({ status: 'existing' as const, receipt: existing });
    }
    for (let attempt = 0; attempt < 3; attempt += 1) {
      let client: PostgresClient;
      try {
        client = await this.#pool.connect();
      } catch (cause) {
        throw new CopilotFailureDiagnosisFinalizationUnavailableError({ cause });
      }
      let began = false;
      try {
        await client.query('BEGIN ISOLATION LEVEL SERIALIZABLE');
        began = true;
        await client.query(`SELECT set_config('statement_timeout', $1, true)`, [
          '5s',
        ]);
        await client.query(`SELECT set_config('lock_timeout', $1, true)`, [
          '2s',
        ]);
        const result = await this.#finalizeInTransaction(client, requestId);
        await client.query('COMMIT');
        began = false;
        return result;
      } catch (cause) {
        if (began) {
          try {
            await client.query('ROLLBACK');
          } catch {
            // Preserve the original failure.
          }
        }
        const state =
          cause && typeof cause === 'object' && 'code' in cause
            ? String(cause.code)
            : '';
        if ((state === '40001' || state === '40P01') && attempt < 2) {
          continue;
        }
        const recovered = await this.findFinalization(requestId);
        if (recovered) {
          return Object.freeze({
            status: 'existing' as const,
            receipt: recovered,
          });
        }
        if (
          cause instanceof CopilotFailureDiagnosisFinalizationConflictError ||
          cause instanceof CopilotFailureDiagnosisModelExecutionInProgressError ||
          cause instanceof CopilotFailureDiagnosisModelResolutionRequiredError ||
          cause instanceof CopilotFailureDiagnosisFinalizationUnavailableError
        ) {
          throw cause;
        }
        throw new CopilotFailureDiagnosisFinalizationUnavailableError({ cause });
      } finally {
        client.release();
      }
    }
    throw new CopilotFailureDiagnosisFinalizationUnavailableError();
  }

  async #finalizeInTransaction(
    client: PostgresClient,
    requestId: string,
  ): Promise<Readonly<{
    status: 'created' | 'existing';
    receipt: Readonly<CopilotFailureDiagnosisFinalizationReceipt>;
  }>> {
    const admission = await client.query<Record<string, unknown>>(
      `SELECT plan_json AS "planJson"
         FROM "ql3_ai"."copilot_failure_diagnosis_admissions"
        WHERE request_id = $1`,
      [requestId],
    );
    if (admission.rows.length !== 1) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    const plan = normalizeCopilotFailureDiagnosisExecutionPlan(
      object(admission.rows[0]!.planJson) as unknown as CopilotFailureDiagnosisExecutionPlan,
    );
    const durable = await client.query<Record<string, unknown>>(
      `SELECT run.status AS "runStatus", run.version AS "runVersion",
              run.event_sequence AS "runEventSequence",
              step.status AS "stepStatus",
              step.step_run_digest AS "stepRunDigest",
              completion.record_json AS "completionJson"
         FROM "ql3"."runs" AS run
         JOIN "ql3"."step_runs" AS step
           ON step.run_id = run.id AND step.id = $1
         LEFT JOIN "ql3_ai"."model_invocation_completions" AS completion
           ON completion.invocation_id = $2
        WHERE run.id = $3
        FOR UPDATE OF run, step`,
      [plan.modelStepRunId, plan.modelInvocationId, plan.runId],
    );
    if (durable.rows.length !== 1) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    const row = durable.rows[0]!;
    if (row.completionJson === null || row.completionJson === undefined) {
      throw new CopilotFailureDiagnosisModelExecutionInProgressError();
    }
    const completion = normalizeModelInvocationCompletionRecord(
      object(row.completionJson) as unknown as ModelInvocationCompletionRecord,
    );
    if (
      completion.invocationId !== plan.modelInvocationId ||
      completion.projectId !== plan.projectId ||
      completion.runId !== plan.runId ||
      completion.stepRunId !== plan.modelStepRunId ||
      completion.traceId !== plan.traceId ||
      string(row.runStatus) !== 'running' ||
      string(row.stepRunDigest) !== completion.completedStepRunDigest
    ) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    if (completion.outcome === 'outcome_unknown') {
      throw new CopilotFailureDiagnosisModelResolutionRequiredError();
    }
    const outcome: CopilotFailureDiagnosisFinalOutcome = completion.outcome;
    if (string(row.stepStatus) !== outcome) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    const output =
      outcome === 'succeeded'
        ? await read(
            client,
            copilotFailureDiagnosisOutputArtifactIdentity(
              plan.modelInvocationId,
            ),
          )
        : null;
    if (
      (outcome === 'succeeded' &&
        (!output ||
          output.requestId !== plan.requestId ||
          output.planDigest !== plan.planDigest ||
          output.invocationId !== completion.invocationId ||
          output.outputBytes !== completion.outputBytes)) ||
      (outcome !== 'succeeded' && output !== null)
    ) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    const runVersion = integer(row.runVersion);
    const eventSequence = integer(row.runEventSequence);
    if (runVersion !== eventSequence) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    const receipt = createCopilotFailureDiagnosisFinalizationReceipt({
      requestId: plan.requestId,
      planDigest: plan.planDigest,
      runId: plan.runId,
      modelStepRunId: plan.modelStepRunId,
      invocationId: plan.modelInvocationId,
      completionDigest: completion.completionDigest,
      outcome,
      outputArtifactId: output?.artifactId ?? null,
      finalRunVersion: runVersion + 1,
      finalRunEventSequence: eventSequence + 1,
      finalizedAtMs: completion.completedAtMs,
    });
    const failure = outcome === 'succeeded'
      ? { code: null, summary: null }
      : outcome === 'timed_out'
      ? {
          code: 'COPILOT_FAILURE_DIAGNOSIS_TIMED_OUT',
          summary: 'Copilot failure diagnosis timed out',
        }
      : {
          code: 'COPILOT_FAILURE_DIAGNOSIS_FAILED',
          summary: 'Copilot failure diagnosis failed',
        };
    const updated = await client.query(
      `UPDATE "ql3"."runs"
          SET status = $1, version = $2, event_sequence = $3,
              output_ref = $4, finished_at_ms = $5,
              error_code = $6, error_summary = $7
        WHERE id = $8 AND status = 'running'
          AND version = $9 AND event_sequence = $10`,
      [
        outcome,
        receipt.finalRunVersion,
        receipt.finalRunEventSequence,
        receipt.outputArtifactId,
        receipt.finalizedAtMs,
        failure.code,
        failure.summary,
        receipt.runId,
        runVersion,
        eventSequence,
      ],
    );
    if ((updated.rowCount ?? updated.rows.length) !== 1) {
      throw new CopilotFailureDiagnosisFinalizationConflictError();
    }
    const payload = JSON.stringify({
      requestId: receipt.requestId,
      planDigest: receipt.planDigest,
      invocationId: receipt.invocationId,
      completionDigest: receipt.completionDigest,
      outcome: receipt.outcome,
      outputArtifactId: receipt.outputArtifactId,
    });
    await client.query(
      `INSERT INTO "ql3"."run_events" (
         id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
         attempt_id, step_run_id, payload, created_at_ms
       ) VALUES ($1, $2, $3, $4, $1, 'system', NULL, NULL, $5, $6::jsonb, $7)`,
      [
        receipt.runEventId,
        receipt.runId,
        receipt.finalRunEventSequence,
        `copilot.diagnosis.${receipt.outcome}`,
        plan.modelStepRunId,
        payload,
        receipt.finalizedAtMs,
      ],
    );
    await client.query(
      `INSERT INTO ${FINALIZATION_TABLE} (
         request_id, plan_digest, run_id, model_step_run_id, invocation_id,
         completion_digest,
         outcome, output_artifact_id, final_run_version,
         final_run_event_sequence, run_event_id, finalized_at_ms,
         receipt_digest, receipt_json
       ) VALUES (
         $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
         $14::jsonb
       )`,
      [
        receipt.requestId,
        receipt.planDigest,
        receipt.runId,
        receipt.modelStepRunId,
        receipt.invocationId,
        receipt.completionDigest,
        receipt.outcome,
        receipt.outputArtifactId,
        receipt.finalRunVersion,
        receipt.finalRunEventSequence,
        receipt.runEventId,
        receipt.finalizedAtMs,
        receipt.receiptDigest,
        JSON.stringify(receipt),
      ],
    );
    return Object.freeze({ status: 'created' as const, receipt });
  }
}