import {
  pgTable,
  uuid,
  varchar,
  text,
  integer,
  jsonb,
  timestamp,
  pgEnum,
  index,
  check,
  type AnyPgColumn,
} from 'drizzle-orm/pg-core';
import { sql } from 'drizzle-orm';
import { projects } from './projects';
import { flows } from './flows';
import { apiKeys } from './apiKeys';

// Enum for generation status
export const generationStatusEnum = pgEnum('generation_status', [
  'pending',
  'processing',
  'success',
  'failed',
]);

// Type for referenced images JSONB
export type ReferencedImage = {
  imageId: string;
  alias: string;
};

export const generations = pgTable(
  'generations',
  {
    id: uuid('id').primaryKey().defaultRandom(),

    // Relations
    projectId: uuid('project_id')
      .notNull()
      .references(() => projects.id, { onDelete: 'cascade' }),
    flowId: uuid('flow_id').references(() => flows.id, { onDelete: 'set null' }),
    pendingFlowId: text('pending_flow_id'), // Temporary UUID for lazy flow pattern
    apiKeyId: uuid('api_key_id').references(() => apiKeys.id, { onDelete: 'set null' }),

    // Status
    status: generationStatusEnum('status').notNull().default('pending'),

    // Prompts (Section 2.1: Reversed semantics)
    // prompt: The prompt that was ACTUALLY USED for generation (enhanced OR original)
    // originalPrompt: User's ORIGINAL input, only stored if autoEnhance was used
    prompt: text('prompt').notNull(), // Prompt used for generation
    originalPrompt: text('original_prompt'), // User's original (nullable, only if enhanced)

    // Generation parameters
    aspectRatio: varchar('aspect_ratio', { length: 10 }),
    width: integer('width'),
    height: integer('height'),

    // AI Model
    modelName: varchar('model_name', { length: 100 }).notNull().default('gemini-flash-image-001'),
    modelVersion: varchar('model_version', { length: 50 }),

    // Result
    outputImageId: uuid('output_image_id').references(
      (): AnyPgColumn => {
        const { images } = require('./images');
        return images.id;
      },
      { onDelete: 'set null' },
    ),

    // Referenced images used in generation
    // Format: [{ "imageId": "uuid", "alias": "@product" }, ...]
    referencedImages: jsonb('referenced_images').$type<ReferencedImage[]>(),

    // Error handling
    errorMessage: text('error_message'),
    errorCode: varchar('error_code', { length: 50 }),
    retryCount: integer('retry_count').notNull().default(0),

    // Metrics
    processingTimeMs: integer('processing_time_ms'),
    cost: integer('cost'), // In cents (USD)

    // Request context
    requestId: uuid('request_id'),
    userAgent: text('user_agent'),
    ipAddress: text('ip_address'),

    // Metadata
    meta: jsonb('meta').$type<Record<string, unknown>>().notNull().default({}),

    // Audit
    createdAt: timestamp('created_at').notNull().defaultNow(),
    updatedAt: timestamp('updated_at')
      .notNull()
      .defaultNow()
      .$onUpdate(() => new Date()),
  },
  (table) => ({
    // CHECK constraints
    statusSuccessCheck: check(
      'status_success_check',
      sql`(${table.status} = 'success' AND ${table.outputImageId} IS NOT NULL) OR (${table.status} != 'success')`,
    ),
    statusFailedCheck: check(
      'status_failed_check',
      sql`(${table.status} = 'failed' AND ${table.errorMessage} IS NOT NULL) OR (${table.status} != 'failed')`,
    ),
    retryCountCheck: check('retry_count_check', sql`${table.retryCount} >= 0`),
    processingTimeCheck: check(
      'processing_time_check',
      sql`${table.processingTimeMs} IS NULL OR ${table.processingTimeMs} >= 0`,
    ),
    costCheck: check('cost_check', sql`${table.cost} IS NULL OR ${table.cost} >= 0`),

    // Indexes
    // Index for querying generations by project and status
    projectStatusIdx: index('idx_generations_project_status').on(
      table.projectId,
      table.status,
      table.createdAt.desc(),
    ),

    // Index for flow-scoped generations (partial index)
    flowIdx: index('idx_generations_flow')
      .on(table.flowId, table.createdAt.desc())
      .where(sql`${table.flowId} IS NOT NULL`),

    // Index for pending flow-scoped generations (partial index)
    pendingFlowIdx: index('idx_generations_pending_flow')
      .on(table.pendingFlowId, table.createdAt.desc())
      .where(sql`${table.pendingFlowId} IS NOT NULL`),

    // Index for output image lookup
    outputIdx: index('idx_generations_output').on(table.outputImageId),

    // Index for request correlation
    requestIdx: index('idx_generations_request').on(table.requestId),

    // Index for API key audit trail
    apiKeyIdx: index('idx_generations_api_key').on(table.apiKeyId),
  }),
);

export type Generation = typeof generations.$inferSelect;
export type NewGeneration = typeof generations.$inferInsert;
