import type { GenerationAsset } from '@lobechat/types';
import { index, integer, jsonb, pgTable, text, uuid, varchar } from 'drizzle-orm/pg-core';
import { createInsertSchema } from 'drizzle-zod';

import { idGenerator } from '../utils/idGenerator';
import { timestamps } from './_helpers';
import type { AsyncTaskSelectItem } from './asyncTask';
import { asyncTasks } from './asyncTask';
import { files } from './file';
import { users } from './user';
import { workspaces } from './workspace';

/**
 * Generation topics table - Used to organize and manage AI-generated content topics
 */
export const generationTopics = pgTable(
  'generation_topics',
  {
    id: text('id')
      .$defaultFn(() => idGenerator('generationTopics'))
      .notNull()
      .primaryKey(),

    userId: text('user_id')
      .references(() => users.id, { onDelete: 'cascade' })
      .notNull(),
    workspaceId: text('workspace_id').references(() => workspaces.id, { onDelete: 'cascade' }),

    /** Brief description of topic content, generated by LLM */
    title: text('title'),

    /** Topic cover image URL */
    coverUrl: text('cover_url'),

    /** Topic type: 'image' or 'video' */
    type: varchar('type', { length: 32 }).notNull().default('image'),

    ...timestamps,
  },
  (t) => [
    index('generation_topics_user_id_idx').on(t.userId),
    index('generation_topics_workspace_id_idx').on(t.workspaceId),
  ],
);

export const insertGenerationTopicSchema = createInsertSchema(generationTopics);

export type NewGenerationTopic = typeof generationTopics.$inferInsert;
export type GenerationTopicItem = typeof generationTopics.$inferSelect;

/**
 * Generation batches table - Stores configuration information for a single generation request
 */
export const generationBatches = pgTable(
  'generation_batches',
  {
    id: text('id')
      .$defaultFn(() => idGenerator('generationBatches'))
      .notNull()
      .primaryKey(),

    userId: text('user_id')
      .references(() => users.id, { onDelete: 'cascade' })
      .notNull(),
    workspaceId: text('workspace_id').references(() => workspaces.id, { onDelete: 'cascade' }),

    generationTopicId: text('generation_topic_id')
      .notNull()
      .references(() => generationTopics.id, { onDelete: 'cascade' }),

    /** Provider name */
    provider: text('provider').notNull(),

    /** Model name */
    model: text('model').notNull(),

    /** Generation prompt */
    prompt: text('prompt').notNull(),

    /** Image width */
    width: integer('width'),

    /** Image height */
    height: integer('height'),

    /** Image aspect ratio */
    ratio: varchar('ratio', { length: 64 }),

    /** Stores generation batch configuration for common settings that don't need indexing */
    config: jsonb('config'),

    ...timestamps,
  },
  (t) => [
    index('generation_batches_user_id_idx').on(t.userId),
    index('generation_batches_topic_id_idx').on(t.generationTopicId),
    index('generation_batches_workspace_id_idx').on(t.workspaceId),
  ],
);

export const insertGenerationBatchSchema = createInsertSchema(generationBatches);

export type NewGenerationBatch = typeof generationBatches.$inferInsert;
export type GenerationBatchItem = typeof generationBatches.$inferSelect;
export type GenerationBatchWithGenerations = GenerationBatchItem & {
  generations: GenerationWithAsyncTask[];
};

/**
 * Stores individual AI generation information
 */
export const generations = pgTable(
  'generations',
  {
    id: text('id')
      .$defaultFn(() => idGenerator('generations'))
      .notNull()
      .primaryKey(),

    userId: text('user_id')
      .references(() => users.id, { onDelete: 'cascade' })
      .notNull(),
    workspaceId: text('workspace_id').references(() => workspaces.id, { onDelete: 'cascade' }),

    generationBatchId: varchar('generation_batch_id', { length: 64 })
      .notNull()
      .references(() => generationBatches.id, { onDelete: 'cascade' }),

    /** Associated async task ID */
    asyncTaskId: uuid('async_task_id').references(() => asyncTasks.id, {
      onDelete: 'set null',
    }),

    /** Associated generated file ID, deletes generation record when file is deleted */
    fileId: text('file_id').references(() => files.id, { onDelete: 'cascade' }),

    /** Generation seed value */
    seed: integer('seed'),

    /** Generated asset information, including S3 storage key, actual image dimensions, thumbnail key, etc. */
    asset: jsonb('asset').$type<GenerationAsset>(),

    ...timestamps,
  },
  (t) => [
    index('generations_user_id_idx').on(t.userId),
    index('generations_batch_id_idx').on(t.generationBatchId),
    index('generations_file_id_idx').on(t.fileId),
    index('generations_workspace_id_idx').on(t.workspaceId),
  ],
);

export const insertGenerationSchema = createInsertSchema(generations);

export type NewGeneration = typeof generations.$inferInsert;
export type GenerationItem = typeof generations.$inferSelect;
export type GenerationWithAsyncTask = GenerationItem & { asyncTask?: AsyncTaskSelectItem };
