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

import { eq, and, isNull, desc, count, inArray, sql } from 'drizzle-orm';
import { db } from '@/db';
import { images, flows, generations } from '@banatie/database';
import type { Image, NewImage, ImageFilters } from '@/types/models';
import { buildWhereClause, buildEqCondition, withoutDeleted } from '@/utils/helpers';
import { ERROR_MESSAGES } from '@/utils/constants';
import { AliasService } from './AliasService';
import { StorageFactory } from '../StorageFactory';

export class ImageService {
  private aliasService: AliasService;

  constructor() {
    this.aliasService = new AliasService();
  }

  async create(data: NewImage): Promise<Image> {
    const [image] = await db.insert(images).values(data).returning();
    if (!image) {
      throw new Error('Failed to create image record');
    }

    // Update flow timestamp if image is part of a flow
    if (image.flowId) {
      await db
        .update(flows)
        .set({ updatedAt: new Date() })
        .where(eq(flows.id, image.flowId));
    }

    return image;
  }

  async getById(id: string, includeDeleted = false): Promise<Image | null> {
    const image = await db.query.images.findFirst({
      where: and(
        eq(images.id, id),
        includeDeleted ? undefined : isNull(images.deletedAt)
      ),
    });

    return image || null;
  }

  async getByIdOrThrow(id: string, includeDeleted = false): Promise<Image> {
    const image = await this.getById(id, includeDeleted);
    if (!image) {
      throw new Error(ERROR_MESSAGES.IMAGE_NOT_FOUND);
    }
    return image;
  }

  async list(
    filters: ImageFilters,
    limit: number,
    offset: number
  ): Promise<{ images: Image[]; total: number }> {
    const conditions = [
      buildEqCondition(images, 'projectId', filters.projectId),
      buildEqCondition(images, 'flowId', filters.flowId),
      buildEqCondition(images, 'source', filters.source),
      buildEqCondition(images, 'alias', filters.alias),
      withoutDeleted(images, filters.deleted),
    ];

    const whereClause = buildWhereClause(conditions);

    const [imagesList, countResult] = await Promise.all([
      db.query.images.findMany({
        where: whereClause,
        orderBy: [desc(images.createdAt)],
        limit,
        offset,
      }),
      db
        .select({ count: count() })
        .from(images)
        .where(whereClause),
    ]);

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

    return {
      images: imagesList,
      total: Number(totalCount),
    };
  }

  async update(
    id: string,
    updates: {
      alias?: string | null;
      focalPoint?: { x: number; y: number };
      meta?: Record<string, unknown>;
    }
  ): Promise<Image> {
    const existing = await this.getByIdOrThrow(id);

    if (updates.alias && updates.alias !== existing.alias) {
      await this.aliasService.validateAliasForAssignment(
        updates.alias,
        existing.projectId,
        existing.flowId || undefined
      );
    }

    const [updated] = await db
      .update(images)
      .set({
        ...updates,
        updatedAt: new Date(),
      })
      .where(eq(images.id, id))
      .returning();

    if (!updated) {
      throw new Error(ERROR_MESSAGES.IMAGE_NOT_FOUND);
    }

    return updated;
  }

  async softDelete(id: string): Promise<Image> {
    const [deleted] = await db
      .update(images)
      .set({
        deletedAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(images.id, id))
      .returning();

    if (!deleted) {
      throw new Error(ERROR_MESSAGES.IMAGE_NOT_FOUND);
    }

    return deleted;
  }

  /**
   * Hard delete image with MinIO cleanup and cascades (Section 7.1)
   * 1. Delete physical file from MinIO storage
   * 2. Delete record from images table (hard delete)
   * 3. Cascade: set outputImageId = NULL in related generations
   * 4. Cascade: remove alias entries from flow.aliases
   * 5. Cascade: remove imageId from generation.referencedImages arrays
   */
  async hardDelete(id: string): Promise<void> {
    // Get image to retrieve storage info
    const image = await this.getById(id, true); // Include deleted
    if (!image) {
      throw new Error(ERROR_MESSAGES.IMAGE_NOT_FOUND);
    }

    try {
      // 1. Delete physical file from MinIO storage
      const storageService = await StorageFactory.getInstance();
      const storageParts = image.storageKey.split('/');

      if (storageParts.length >= 4) {
        const orgId = storageParts[0]!;
        const projectId = storageParts[1]!;
        const category = storageParts[2]! as 'uploads' | 'generated' | 'references';
        const filename = storageParts.slice(3).join('/');

        await storageService.deleteFile(orgId, projectId, category, filename);
      }

      // 2. Cascade: Set outputImageId = NULL in related generations
      await db
        .update(generations)
        .set({ outputImageId: null })
        .where(eq(generations.outputImageId, id));

      // 3. Cascade: Remove alias entries from flow.aliases where this imageId is referenced
      const allFlows = await db.query.flows.findMany();
      for (const flow of allFlows) {
        const aliases = (flow.aliases as Record<string, string>) || {};
        let modified = false;

        // Remove all entries where value equals this imageId
        const newAliases: Record<string, string> = {};
        for (const [key, value] of Object.entries(aliases)) {
          if (value !== id) {
            newAliases[key] = value;
          } else {
            modified = true;
          }
        }

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

      // 4. Cascade: Remove imageId from generation.referencedImages JSON arrays
      const affectedGenerations = await db.query.generations.findMany({
        where: sql`${generations.referencedImages}::jsonb @> ${JSON.stringify([{ imageId: id }])}`,
      });

      for (const gen of affectedGenerations) {
        const refs = (gen.referencedImages as Array<{ imageId: string; alias: string }>) || [];
        const filtered = refs.filter(ref => ref.imageId !== id);

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

      // 5. Delete record from images table
      await db.delete(images).where(eq(images.id, id));

    } catch (error) {
      // Per Section 7.4: If MinIO delete fails, do NOT proceed with DB cleanup
      // This prevents orphaned files in MinIO
      console.error('MinIO delete failed, aborting image deletion:', error);
      throw new Error(ERROR_MESSAGES.STORAGE_DELETE_FAILED || 'Failed to delete file from storage');
    }
  }

  async assignProjectAlias(imageId: string, alias: string): Promise<Image> {
    const image = await this.getByIdOrThrow(imageId);

    if (image.flowId) {
      throw new Error('Cannot assign project alias to flow-scoped image');
    }

    await this.aliasService.validateAliasForAssignment(
      alias,
      image.projectId
    );

    const [updated] = await db
      .update(images)
      .set({
        alias,
        updatedAt: new Date(),
      })
      .where(eq(images.id, imageId))
      .returning();

    if (!updated) {
      throw new Error(ERROR_MESSAGES.IMAGE_NOT_FOUND);
    }

    return updated;
  }

  /**
   * Reassign a project-scoped alias to a new image
   * Clears the alias from any existing image and assigns it to the new image
   * Implements override behavior per Section 5.2 of api-refactoring-final.md
   *
   * @param alias - The alias to reassign (e.g., "@hero")
   * @param newImageId - ID of the image to receive the alias
   * @param projectId - Project ID for scope validation
   */
  async reassignProjectAlias(
    alias: string,
    newImageId: string,
    projectId: string
  ): Promise<void> {
    // Step 1: Clear alias from any existing image with this alias
    // Project aliases can exist on images with or without flowId
    await db
      .update(images)
      .set({
        alias: null,
        updatedAt: new Date()
      })
      .where(
        and(
          eq(images.projectId, projectId),
          eq(images.alias, alias),
          isNull(images.deletedAt)
        )
      );

    // Step 2: Assign alias to new image
    await db
      .update(images)
      .set({
        alias: alias,
        updatedAt: new Date()
      })
      .where(eq(images.id, newImageId));
  }

  async getByStorageKey(storageKey: string): Promise<Image | null> {
    const image = await db.query.images.findFirst({
      where: and(
        eq(images.storageKey, storageKey),
        isNull(images.deletedAt)
      ),
    });

    return image || null;
  }

  async getByFileHash(fileHash: string, projectId: string): Promise<Image | null> {
    const image = await db.query.images.findFirst({
      where: and(
        eq(images.fileHash, fileHash),
        eq(images.projectId, projectId),
        isNull(images.deletedAt)
      ),
    });

    return image || null;
  }

  async getMultipleByIds(ids: string[]): Promise<Image[]> {
    if (ids.length === 0) {
      return [];
    }

    return await db.query.images.findMany({
      where: and(
        inArray(images.id, ids),
        isNull(images.deletedAt)
      ),
    });
  }

  /**
   * Link all pending images to a flow
   * Called when flow is created to attach all images with matching pendingFlowId
   */
  async linkPendingImagesToFlow(
    flowId: string,
    projectId: string
  ): Promise<void> {
    // Find all images with pendingFlowId matching this flowId
    const pendingImages = await db.query.images.findMany({
      where: and(
        eq(images.pendingFlowId, flowId),
        eq(images.projectId, projectId)
      ),
    });

    if (pendingImages.length === 0) {
      return;
    }

    // Update images: set flowId and clear pendingFlowId
    await db
      .update(images)
      .set({
        flowId: flowId,
        pendingFlowId: null,
        updatedAt: new Date(),
      })
      .where(
        and(
          eq(images.pendingFlowId, flowId),
          eq(images.projectId, projectId)
        )
      );
  }
}