import type { Pool, PoolClient, QueryResult } from 'pg';
import { RepositoryConflictError, RepositoryNotFoundError, RepositoryValidationError } from '../../src/persistence/repository-errors';
import { translateDatabaseError } from '../postgres/postgres-assessor-repository';
import type { ComplianceProjectExport, ComplianceStore } from '../../src/compliance/compliance-store';
import {
  regulatoryMappingBatchRunValidationErrors,
  regulatoryMappingDecisionValidationErrors,
  regulatoryMappingRunValidationErrors,
  type RegulatoryMappingBatchRun,
  type RegulatoryMappingDecision,
  type RegulatoryMappingRun
} from '../../src/compliance/regulatory-mapping';
import { aiOverviewRunValidationErrors, assertAiOverviewTransition, type AiOverviewRun } from '../../src/compliance/ai-overview';

interface DataRow<T> { data: T }

function clone<T>(value: T): T {
  return JSON.parse(JSON.stringify(value)) as T;
}

function json(value: unknown): string {
  return JSON.stringify(value);
}

function dataRows<T>(result: QueryResult<DataRow<T>>): T[] {
  return result.rows.map(row => clone(row.data));
}

function dataRow<T>(result: QueryResult<DataRow<T>>): T | null {
  return result.rows[0] ? clone(result.rows[0].data) : null;
}

/**
 * The Compliance Engine's own tables (migration 026, and the AI overview runs of migration 025). Writes run in the engine's own transactions:
 * they never advance the ASSESSOR repository revision and never trigger a Result DB rebuild.
 */
export class PostgresComplianceStore implements ComplianceStore {
  constructor(private readonly pool: Pool) {}

  private async transaction<T>(work: (client: PoolClient) => Promise<T>): Promise<T> {
    const client = await this.pool.connect();
    try {
      await client.query('BEGIN');
      const result = await work(client);
      await client.query('COMMIT');
      return result;
    } catch (error) {
      try { await client.query('ROLLBACK'); } catch { /* connection failure keeps original error */ }
      throw translateDatabaseError(error);
    } finally {
      client.release();
    }
  }

  async getRegulatoryMappingRun(id: string): Promise<RegulatoryMappingRun | null> {
    return dataRow(await this.pool.query<DataRow<RegulatoryMappingRun>>(
      'SELECT data FROM assessor_regulatory_mappings WHERE id = $1', [id]
    ));
  }

  async claimRegulatoryMappingRun(run: RegulatoryMappingRun, forceRegeneration: boolean): Promise<{ run: RegulatoryMappingRun; claimed: boolean }> {
    const errors = regulatoryMappingRunValidationErrors(run);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      const inserted = await client.query(`
        INSERT INTO assessor_regulatory_mappings
          (id, rule_key, source_tool, source_rule_id, run_status, query_hash, corpus_id, corpus_version,
           embedding_model, retrieval_version, prompt_version, output_schema_version, model, created_at, updated_at, data)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16::jsonb)
        ON CONFLICT DO NOTHING
        RETURNING id
      `, [run.id, run.ruleKey, run.sourceTool, run.sourceRuleId ?? null, run.status, run.queryHash, run.corpusId,
        run.corpusVersion, run.embeddingModel, run.retrievalVersion, run.promptVersion, run.outputSchemaVersion,
        run.model, run.createdAt, run.updatedAt, json(run)]);
      if (inserted.rowCount) return { run: clone(run), claimed: true };
      const existing = dataRow(await client.query<DataRow<RegulatoryMappingRun>>(
        'SELECT data FROM assessor_regulatory_mappings WHERE id = $1 FOR UPDATE', [run.id]
      ));
      if (!existing) throw new RepositoryConflictError('Regulatory mapping identity conflicted with another pipeline tuple.');
      if (['queued', 'running'].includes(existing.status) || existing.status === 'succeeded' && !forceRegeneration) {
        return { run: existing, claimed: false };
      }
      await client.query(`
        UPDATE assessor_regulatory_mappings SET
          run_status = $2, query_hash = $3, created_at = $4, updated_at = $5, data = $6::jsonb
        WHERE id = $1
      `, [run.id, run.status, run.queryHash, run.createdAt, run.updatedAt, json(run)]);
      return { run: clone(run), claimed: true };
    });
  }

  async saveRegulatoryMappingRun(run: RegulatoryMappingRun): Promise<RegulatoryMappingRun> {
    const errors = regulatoryMappingRunValidationErrors(run);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      await client.query(`
        INSERT INTO assessor_regulatory_mappings
          (id, rule_key, source_tool, source_rule_id, run_status, query_hash, corpus_id, corpus_version,
           embedding_model, retrieval_version, prompt_version, output_schema_version, model, created_at, updated_at, data)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16::jsonb)
        ON CONFLICT (id) DO UPDATE SET
          run_status = excluded.run_status, updated_at = excluded.updated_at, data = excluded.data
      `, [run.id, run.ruleKey, run.sourceTool, run.sourceRuleId ?? null, run.status, run.queryHash, run.corpusId,
        run.corpusVersion, run.embeddingModel, run.retrievalVersion, run.promptVersion, run.outputSchemaVersion,
        run.model, run.createdAt, run.updatedAt, json(run)]);
      return clone(run);
    });
  }

  async listRegulatoryMappingRuns(): Promise<RegulatoryMappingRun[]> {
    return dataRows(await this.pool.query<DataRow<RegulatoryMappingRun>>(
      'SELECT data FROM assessor_regulatory_mappings ORDER BY created_at DESC, id DESC'
    ));
  }

  async createRegulatoryMappingBatchRun(run: RegulatoryMappingBatchRun): Promise<RegulatoryMappingBatchRun> {
    const errors = regulatoryMappingBatchRunValidationErrors(run);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      const inserted = await client.query(`
        INSERT INTO assessor_regulatory_mapping_batches
          (id, project_id, scope_key, filter_hash, idempotency_key, run_status, created_at, updated_at, data)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::jsonb)
        ON CONFLICT DO NOTHING
        RETURNING data
      `, [run.id, run.projectId, run.scopeKey, run.filterHash, run.idempotencyKey, run.status,
        run.createdAt, run.updatedAt, json(run)]);
      if (inserted.rowCount) return clone(run);
      const existing = dataRow(await client.query<DataRow<RegulatoryMappingBatchRun>>(`
        SELECT data FROM assessor_regulatory_mapping_batches
        WHERE project_id = $1 AND (idempotency_key = $2 OR (scope_key = $3 AND run_status IN ('queued', 'running')))
        ORDER BY (idempotency_key = $2) DESC, created_at DESC, id DESC
        LIMIT 1
      `, [run.projectId, run.idempotencyKey, run.scopeKey]));
      if (!existing) throw new RepositoryConflictError('Regulatory mapping batch conflicted with another active request.');
      return existing;
    });
  }

  async updateRegulatoryMappingBatchRun(run: RegulatoryMappingBatchRun): Promise<RegulatoryMappingBatchRun> {
    const errors = regulatoryMappingBatchRunValidationErrors(run);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      const existing = dataRow(await client.query<DataRow<RegulatoryMappingBatchRun>>(
        'SELECT data FROM assessor_regulatory_mapping_batches WHERE id = $1 FOR UPDATE', [run.id]
      ));
      if (!existing) throw new RepositoryNotFoundError('Regulatory mapping batch does not exist.');
      if (existing.projectId !== run.projectId || existing.scopeKey !== run.scopeKey || existing.idempotencyKey !== run.idempotencyKey) {
        throw new RepositoryConflictError('Regulatory mapping batch identity is immutable.');
      }
      if (['completed', 'failed'].includes(existing.status)) throw new RepositoryConflictError('A completed regulatory mapping batch is immutable.');
      const allowed = existing.status === 'queued' ? ['running', 'failed'] : ['running', 'completed', 'failed'];
      if (!allowed.includes(run.status)) throw new RepositoryConflictError(`Invalid regulatory mapping batch transition from ${existing.status} to ${run.status}.`);
      await client.query(`
        UPDATE assessor_regulatory_mapping_batches SET run_status = $2, updated_at = $3, data = $4::jsonb
        WHERE id = $1
      `, [run.id, run.status, run.updatedAt, json(run)]);
      return clone(run);
    });
  }

  async listRegulatoryMappingBatchRuns(projectId: string, scopeKey?: string): Promise<RegulatoryMappingBatchRun[]> {
    return dataRows(await this.pool.query<DataRow<RegulatoryMappingBatchRun>>(`
      SELECT data FROM assessor_regulatory_mapping_batches
      WHERE project_id = $1 AND ($2::text IS NULL OR scope_key = $2)
      ORDER BY created_at DESC, id DESC
    `, [projectId, scopeKey ?? null]));
  }

  async createRegulatoryMappingDecision(
    decision: RegulatoryMappingDecision,
    expectedLatestDecisionId?: string
  ): Promise<RegulatoryMappingDecision> {
    const errors = regulatoryMappingDecisionValidationErrors(decision);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      // Serializes decisions for one rule and control, so two writers cannot both see the same latest decision.
      await client.query('SELECT pg_advisory_xact_lock(hashtext($1))', [`regulatory-decision:${decision.ruleKey}:${decision.controlId}`]);
      const run = dataRow(await client.query<DataRow<RegulatoryMappingRun>>(
        'SELECT data FROM assessor_regulatory_mappings WHERE id = $1 FOR UPDATE', [decision.mappingId]
      ));
      if (!run || run.ruleKey !== decision.ruleKey || run.status !== 'succeeded') {
        throw new RepositoryNotFoundError('Completed regulatory mapping does not exist for this rule.');
      }
      if (!run.candidates.some(candidate => candidate.controlId === decision.controlId)) {
        throw new RepositoryValidationError('Decision control was not a grounded candidate.');
      }
      const latest = dataRow(await client.query<DataRow<RegulatoryMappingDecision>>(`
        SELECT data FROM assessor_regulatory_mapping_decisions
        WHERE rule_key = $1 AND control_id = $2
        ORDER BY created_at DESC, id DESC LIMIT 1
      `, [decision.ruleKey, decision.controlId]));
      if ((expectedLatestDecisionId ?? undefined) !== (latest?.id ?? undefined)) {
        throw new RepositoryConflictError('Regulatory mapping decision changed before this update.');
      }
      const stored = { ...decision, ...(latest ? { previousDecisionId: latest.id } : {}) };
      await client.query(`
        INSERT INTO assessor_regulatory_mapping_decisions
          (id, rule_key, source_tool, source_rule_id, mapping_id, control_id, decision, previous_decision_id, created_at, data)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::jsonb)
      `, [stored.id, stored.ruleKey, stored.sourceTool, stored.sourceRuleId ?? null, stored.mappingId, stored.controlId,
        stored.decision, stored.previousDecisionId ?? null, stored.createdAt, json(stored)]);
      return clone(stored);
    });
  }

  async listRegulatoryMappingDecisions(ruleKey: string): Promise<RegulatoryMappingDecision[]> {
    return dataRows(await this.pool.query<DataRow<RegulatoryMappingDecision>>(`
      SELECT data FROM assessor_regulatory_mapping_decisions
      WHERE rule_key = $1
      ORDER BY created_at DESC, id DESC
    `, [ruleKey]));
  }

  async createAiOverviewRun(run: AiOverviewRun): Promise<AiOverviewRun> {
    const errors = aiOverviewRunValidationErrors(run);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      const idempotent = dataRow(await client.query<DataRow<AiOverviewRun>>(`
        SELECT data FROM assessor_ai_overview_runs
        WHERE project_id = $1 AND scope_key = $2 AND idempotency_key = $3
      `, [run.projectId, run.scopeKey, run.idempotencyKey]));
      if (idempotent) return idempotent;
      await client.query(`
        INSERT INTO assessor_ai_overview_runs
          (id, project_id, scope_key, assessment_ids, scan_run_ids, idempotency_key, run_status,
           input_hash, prompt_version, output_schema_version, model, created_at, updated_at, data)
        VALUES ($1, $2, $3, $4::text[], $5::text[], $6, $7, $8, $9, $10, $11, $12, $13, $14::jsonb)
      `, [run.id, run.projectId, run.scopeKey, run.assessmentIds, run.scanRunIds, run.idempotencyKey,
        run.status, run.inputHash, run.promptVersion, run.outputSchemaVersion, run.model,
        run.createdAt, run.updatedAt, json(run)]);
      return clone(run);
    });
  }

  async updateAiOverviewRun(run: AiOverviewRun): Promise<AiOverviewRun> {
    const errors = aiOverviewRunValidationErrors(run);
    if (errors.length) throw new RepositoryValidationError(errors.join(' '));
    return this.transaction(async client => {
      const existing = dataRow(await client.query<DataRow<AiOverviewRun>>(
        'SELECT data FROM assessor_ai_overview_runs WHERE id = $1 FOR UPDATE', [run.id]
      ));
      if (!existing) throw new RepositoryNotFoundError(`AI overview run ${run.id} does not exist.`);
      assertAiOverviewTransition(existing, run);
      await client.query(`
        UPDATE assessor_ai_overview_runs SET
          run_status = $2, updated_at = $3, data = $4::jsonb
        WHERE id = $1
      `, [run.id, run.status, run.updatedAt, json(run)]);
      return clone(run);
    });
  }

  async listAiOverviewRuns(projectId: string, scopeKey?: string): Promise<AiOverviewRun[]> {
    return dataRows(await this.pool.query<DataRow<AiOverviewRun>>(`
      SELECT data FROM assessor_ai_overview_runs
      WHERE project_id = $1 AND ($2::text IS NULL OR scope_key = $2)
      ORDER BY created_at DESC, id DESC
    `, [projectId, scopeKey ?? null]));
  }

  async deleteAiOverviewRuns(projectId: string, runIds: string[]): Promise<void> {
    if (!runIds.length) return;
    await this.transaction(client => client.query(
      'DELETE FROM assessor_ai_overview_runs WHERE project_id = $1 AND id = ANY($2::text[])', [projectId, runIds]
    ));
  }

  async exportProject(projectId: string): Promise<ComplianceProjectExport> {
    return {
      regulatoryMappingBatchRuns: await this.listRegulatoryMappingBatchRuns(projectId),
      regulatoryMappingRuns: (await this.listRegulatoryMappingRuns()).filter(item => item.status === 'succeeded'),
      regulatoryMappingDecisions: dataRows(await this.pool.query<DataRow<RegulatoryMappingDecision>>(
        'SELECT data FROM assessor_regulatory_mapping_decisions ORDER BY created_at, id'
      )),
      aiOverviewRuns: (await this.listAiOverviewRuns(projectId)).filter(item => item.status === 'succeeded')
    };
  }
}