banatie-service/apps/api-service/src/services/core/GenerationService.ts

import { randomUUID } from 'crypto';
import { eq, desc, count } from 'drizzle-orm';
import { db } from '@/db';
import { generations, flows } from '@banatie/database';
import type {
  Generation,
  NewGeneration,
  GenerationWithRelations,
  GenerationFilters,
} from '@/types/models';
import { ImageService } from './ImageService';
import { AliasService } from './AliasService';
import { ImageGenService } from '../ImageGenService';
import { StorageFactory } from '../StorageFactory';
import { buildWhereClause, buildEqCondition } from '@/utils/helpers';
import { ERROR_MESSAGES, GENERATION_LIMITS } from '@/utils/constants';
import { extractAliasesFromPrompt } from '@/utils/validators';
import type { ReferenceImage } from '@/types/api';

export interface CreateGenerationParams {
  projectId: string;
  apiKeyId: string;
  prompt: string;
  referenceImages?: string[] | undefined; // Aliases to resolve
  aspectRatio?: string | undefined;
  flowId?: string | undefined;
  alias?: string | undefined;
  flowAlias?: string | undefined;
  autoEnhance?: boolean | undefined;
  enhancedPrompt?: string | undefined;
  meta?: Record<string, unknown> | undefined;
  requestId?: string | undefined;
}

export class GenerationService {
  private imageService: ImageService;
  private aliasService: AliasService;
  private imageGenService: ImageGenService;

  constructor() {
    this.imageService = new ImageService();
    this.aliasService = new AliasService();

    const geminiApiKey = process.env['GEMINI_API_KEY'];
    if (!geminiApiKey) {
      throw new Error('GEMINI_API_KEY environment variable is required');
    }
    this.imageGenService = new ImageGenService(geminiApiKey);
  }

  async create(params: CreateGenerationParams): Promise<GenerationWithRelations> {
    const startTime = Date.now();

    // Auto-detect aliases from prompt and merge with manual references
    const autoDetectedAliases = extractAliasesFromPrompt(params.prompt);
    const manualReferences = params.referenceImages || [];

    // Merge: manual references first, then auto-detected (remove duplicates)
    const allReferences = Array.from(new Set([...manualReferences, ...autoDetectedAliases]));

    // FlowId logic (Section 10.1):
    // - If undefined (not provided) → generate new UUID
    // - If null (explicitly null) → keep null
    // - If string (specific value) → use that value
    let finalFlowId: string | null;
    if (params.flowId === undefined) {
      finalFlowId = randomUUID();
    } else if (params.flowId === null) {
      finalFlowId = null;
    } else {
      finalFlowId = params.flowId;
    }

    // Prompt semantics (Section 2.1):
    // - If autoEnhance = false OR no enhancedPrompt: prompt = user input, originalPrompt = null
    // - If autoEnhance = true AND enhancedPrompt: prompt = enhanced, originalPrompt = user input
    const usedPrompt = params.enhancedPrompt || params.prompt;
    const preservedOriginal = params.enhancedPrompt ? params.prompt : null;

    const generationRecord: NewGeneration = {
      projectId: params.projectId,
      flowId: finalFlowId,
      apiKeyId: params.apiKeyId,
      status: 'pending',
      prompt: usedPrompt, // Prompt actually used for generation
      originalPrompt: preservedOriginal, // User's original (only if enhanced)
      aspectRatio: params.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO,
      referencedImages: null,
      requestId: params.requestId || null,
      meta: params.meta || {},
    };

    const [generation] = await db
      .insert(generations)
      .values(generationRecord)
      .returning();

    if (!generation) {
      throw new Error('Failed to create generation record');
    }

    try {
      await this.updateStatus(generation.id, 'processing');

      let referenceImageBuffers: ReferenceImage[] = [];
      let referencedImagesMetadata: Array<{ imageId: string; alias: string }> = [];

      if (allReferences.length > 0) {
        const resolved = await this.resolveReferenceImages(
          allReferences,
          params.projectId,
          params.flowId
        );
        referenceImageBuffers = resolved.buffers;
        referencedImagesMetadata = resolved.metadata;

        await db
          .update(generations)
          .set({ referencedImages: referencedImagesMetadata })
          .where(eq(generations.id, generation.id));
      }

      const genResult = await this.imageGenService.generateImage({
        prompt: usedPrompt, // Use the prompt that was stored (enhanced or original)
        filename: `gen_${generation.id}`,
        referenceImages: referenceImageBuffers,
        aspectRatio: params.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO,
        orgId: 'default',
        projectId: params.projectId,
        meta: params.meta || {},
      });

      if (!genResult.success) {
        const processingTime = Date.now() - startTime;
        await this.updateStatus(generation.id, 'failed', {
          errorMessage: genResult.error || 'Generation failed',
          processingTimeMs: processingTime,
        });
        throw new Error(genResult.error || 'Generation failed');
      }

      const storageKey = genResult.filepath!;
      // TODO: Add file hash computation when we have a helper to download by storageKey
      const fileHash = null;

      const imageRecord = await this.imageService.create({
        projectId: params.projectId,
        flowId: finalFlowId,
        generationId: generation.id,
        apiKeyId: params.apiKeyId,
        storageKey,
        storageUrl: genResult.url!,
        mimeType: 'image/jpeg',
        fileSize: genResult.size || 0,
        fileHash,
        source: 'generated',
        alias: params.alias || null,
        meta: params.meta || {},
      });

      // Eager flow creation if flowAlias is provided (Section 4.2)
      if (params.flowAlias && finalFlowId) {
        // Check if flow exists, create if not
        const existingFlow = await db.query.flows.findFirst({
          where: eq(flows.id, finalFlowId),
        });

        if (!existingFlow) {
          await db.insert(flows).values({
            id: finalFlowId,
            projectId: params.projectId,
            aliases: {},
            meta: {},
          });
        }

        await this.assignFlowAlias(finalFlowId, params.flowAlias, imageRecord.id);
      }

      if (finalFlowId) {
        await db
          .update(flows)
          .set({ updatedAt: new Date() })
          .where(eq(flows.id, finalFlowId));
      }

      const processingTime = Date.now() - startTime;
      await this.updateStatus(generation.id, 'success', {
        outputImageId: imageRecord.id,
        processingTimeMs: processingTime,
      });

      return await this.getByIdWithRelations(generation.id);
    } catch (error) {
      const processingTime = Date.now() - startTime;
      await this.updateStatus(generation.id, 'failed', {
        errorMessage: error instanceof Error ? error.message : 'Unknown error',
        processingTimeMs: processingTime,
      });
      throw error;
    }
  }

  private async resolveReferenceImages(
    aliases: string[],
    projectId: string,
    flowId?: string
  ): Promise<{
    buffers: ReferenceImage[];
    metadata: Array<{ imageId: string; alias: string }>;
  }> {
    const resolutions = await this.aliasService.resolveMultiple(aliases, projectId, flowId);

    const buffers: ReferenceImage[] = [];
    const metadata: Array<{ imageId: string; alias: string }> = [];

    const storageService = await StorageFactory.getInstance();

    for (const [alias, resolution] of resolutions) {
      if (!resolution.image) {
        throw new Error(`${ERROR_MESSAGES.ALIAS_NOT_FOUND}: ${alias}`);
      }

      const parts = resolution.image.storageKey.split('/');
      if (parts.length < 4) {
        throw new Error(`Invalid storage key format: ${resolution.image.storageKey}`);
      }

      const orgId = parts[0]!;
      const projId = parts[1]!;
      const category = parts[2]! as 'uploads' | 'generated' | 'references';
      const filename = parts.slice(3).join('/');

      const buffer = await storageService.downloadFile(
        orgId,
        projId,
        category,
        filename
      );

      buffers.push({
        buffer,
        mimetype: resolution.image.mimeType,
        originalname: filename,
      });

      metadata.push({
        imageId: resolution.imageId,
        alias,
      });
    }

    return { buffers, metadata };
  }

  private async assignFlowAlias(
    flowId: string,
    flowAlias: string,
    imageId: string
  ): Promise<void> {
    const flow = await db.query.flows.findFirst({
      where: eq(flows.id, flowId),
    });

    if (!flow) {
      throw new Error(ERROR_MESSAGES.FLOW_NOT_FOUND);
    }

    const currentAliases = (flow.aliases as Record<string, string>) || {};
    const updatedAliases = { ...currentAliases };

    // Assign the flow alias to the image
    updatedAliases[flowAlias] = imageId;

    await db
      .update(flows)
      .set({ aliases: updatedAliases, updatedAt: new Date() })
      .where(eq(flows.id, flowId));
  }

  private async updateStatus(
    id: string,
    status: 'pending' | 'processing' | 'success' | 'failed',
    additionalUpdates?: {
      errorMessage?: string;
      outputImageId?: string;
      processingTimeMs?: number;
    }
  ): Promise<void> {
    await db
      .update(generations)
      .set({
        status,
        ...additionalUpdates,
        updatedAt: new Date(),
      })
      .where(eq(generations.id, id));
  }

  async getById(id: string): Promise<Generation | null> {
    const generation = await db.query.generations.findFirst({
      where: eq(generations.id, id),
    });

    return generation || null;
  }

  async getByIdWithRelations(id: string): Promise<GenerationWithRelations> {
    const generation = await db.query.generations.findFirst({
      where: eq(generations.id, id),
      with: {
        outputImage: true,
        flow: true,
      },
    });

    if (!generation) {
      throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND);
    }

    if (generation.referencedImages && Array.isArray(generation.referencedImages)) {
      const refImageIds = (generation.referencedImages as Array<{ imageId: string; alias: string }>)
        .map((ref) => ref.imageId);
      const refImages = await this.imageService.getMultipleByIds(refImageIds);
      return {
        ...generation,
        referenceImages: refImages,
      } as GenerationWithRelations;
    }

    return generation as GenerationWithRelations;
  }

  async list(
    filters: GenerationFilters,
    limit: number,
    offset: number
  ): Promise<{ generations: GenerationWithRelations[]; total: number }> {
    const conditions = [
      buildEqCondition(generations, 'projectId', filters.projectId),
      buildEqCondition(generations, 'flowId', filters.flowId),
      buildEqCondition(generations, 'status', filters.status),
    ];

    const whereClause = buildWhereClause(conditions);

    const [generationsList, countResult] = await Promise.all([
      db.query.generations.findMany({
        where: whereClause,
        orderBy: [desc(generations.createdAt)],
        limit,
        offset,
        with: {
          outputImage: true,
          flow: true,
        },
      }),
      db
        .select({ count: count() })
        .from(generations)
        .where(whereClause),
    ]);

    const totalCount = countResult[0]?.count || 0;

    return {
      generations: generationsList as GenerationWithRelations[],
      total: Number(totalCount),
    };
  }

  /**
   * Regenerate an existing generation (Section 3)
   * - Allows regeneration for any status (no status checks)
   * - Uses exact same parameters as original
   * - Updates existing image (same ID, path, URL)
   * - No retry count logic
   */
  async regenerate(id: string): Promise<GenerationWithRelations> {
    const generation = await this.getById(id);
    if (!generation) {
      throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND);
    }

    if (!generation.outputImageId) {
      throw new Error('Cannot regenerate generation without output image');
    }

    const startTime = Date.now();

    try {
      // Update status to processing
      await this.updateStatus(id, 'processing');

      // Use EXACT same parameters as original (no overrides)
      const genResult = await this.imageGenService.generateImage({
        prompt: generation.prompt,
        filename: `gen_${id}`,
        referenceImages: [], // TODO: Re-resolve referenced images if needed
        aspectRatio: generation.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO,
        orgId: 'default',
        projectId: generation.projectId,
        meta: generation.meta as Record<string, unknown> || {},
      });

      if (!genResult.success) {
        const processingTime = Date.now() - startTime;
        await this.updateStatus(id, 'failed', {
          errorMessage: genResult.error || 'Regeneration failed',
          processingTimeMs: processingTime,
        });
        throw new Error(genResult.error || 'Regeneration failed');
      }

      // Note: Physical file in MinIO is overwritten by ImageGenService
      // Image record preserves: imageId, storageKey, storageUrl, alias, createdAt
      // Image record updates: fileSize (if changed), updatedAt

      const processingTime = Date.now() - startTime;
      await this.updateStatus(id, 'success', {
        processingTimeMs: processingTime,
      });

      return await this.getByIdWithRelations(id);
    } catch (error) {
      const processingTime = Date.now() - startTime;
      await this.updateStatus(id, 'failed', {
        errorMessage: error instanceof Error ? error.message : 'Unknown error',
        processingTimeMs: processingTime,
      });
      throw error;
    }
  }

  // Keep retry() for backward compatibility, delegate to regenerate()
  async retry(id: string, overrides?: { prompt?: string; aspectRatio?: string }): Promise<GenerationWithRelations> {
    // Ignore overrides, regenerate with original parameters
    return await this.regenerate(id);
  }

  async update(
    id: string,
    updates: {
      prompt?: string;
      aspectRatio?: string;
      flowId?: string | null;
      meta?: Record<string, unknown>;
    }
  ): Promise<GenerationWithRelations> {
    const generation = await this.getById(id);
    if (!generation) {
      throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND);
    }

    // Check if generative parameters changed (prompt or aspectRatio)
    const shouldRegenerate =
      (updates.prompt !== undefined && updates.prompt !== generation.prompt) ||
      (updates.aspectRatio !== undefined && updates.aspectRatio !== generation.aspectRatio);

    // Handle flowId change (Section 9.2)
    if (updates.flowId !== undefined && updates.flowId !== null) {
      // If flowId provided and not null, create flow if it doesn't exist (eager creation)
      const existingFlow = await db.query.flows.findFirst({
        where: eq(flows.id, updates.flowId),
      });

      if (!existingFlow) {
        await db.insert(flows).values({
          id: updates.flowId,
          projectId: generation.projectId,
          aliases: {},
          meta: {},
        });
      }
    }

    // Update database fields
    const updateData: Partial<NewGeneration> = {};
    if (updates.prompt !== undefined) {
      updateData.prompt = updates.prompt; // Update the prompt used for generation
    }
    if (updates.aspectRatio !== undefined) {
      updateData.aspectRatio = updates.aspectRatio;
    }
    if (updates.flowId !== undefined) {
      updateData.flowId = updates.flowId;
    }
    if (updates.meta !== undefined) {
      updateData.meta = updates.meta;
    }

    if (Object.keys(updateData).length > 0) {
      await db
        .update(generations)
        .set({ ...updateData, updatedAt: new Date() })
        .where(eq(generations.id, id));
    }

    // If generative parameters changed, trigger regeneration
    if (shouldRegenerate && generation.outputImageId) {
      // Update status to processing
      await this.updateStatus(id, 'processing');

      try {
        // Use updated prompt/aspectRatio or fall back to existing
        const promptToUse = updates.prompt || generation.prompt;
        const aspectRatioToUse = updates.aspectRatio || generation.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO;

        // Regenerate image
        const genResult = await this.imageGenService.generateImage({
          prompt: promptToUse,
          filename: `gen_${id}`,
          referenceImages: [],
          aspectRatio: aspectRatioToUse,
          orgId: 'default',
          projectId: generation.projectId,
          meta: updates.meta || generation.meta || {},
        });

        if (!genResult.success) {
          await this.updateStatus(id, 'failed', {
            errorMessage: genResult.error || 'Regeneration failed',
          });
          throw new Error(genResult.error || 'Regeneration failed');
        }

        // Note: Physical file in MinIO is overwritten by ImageGenService
        // TODO: Update fileSize and other metadata when ImageService.update() supports it

        await this.updateStatus(id, 'success');
      } catch (error) {
        await this.updateStatus(id, 'failed', {
          errorMessage: error instanceof Error ? error.message : 'Unknown error',
        });
        throw error;
      }
    }

    return await this.getByIdWithRelations(id);
  }

  async delete(id: string): Promise<void> {
    const generation = await this.getById(id);
    if (!generation) {
      throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND);
    }

    if (generation.outputImageId) {
      await this.imageService.softDelete(generation.outputImageId);
    }

    await db.delete(generations).where(eq(generations.id, id));
  }
}