import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../infrastructure/database/prisma.service';

@Injectable()
export class DocumentImportRepository {
  constructor(private readonly prisma: PrismaService) {}

  async create(data: {
    userId?: string;
    knowledgeBaseId?: string;
    sourceId?: string;
    fileId?: string;
    sourceType?: string;
    sourceName?: string;
    sourceUrl?: string;
    status?: string;
  }) {
    return this.prisma.documentImport.create({
      data: {
        userId: data.userId ?? '',
        knowledgeBaseId: data.knowledgeBaseId,
        sourceId: data.sourceId,
        fileId: data.fileId,
        sourceType: data.sourceType ?? 'file',
        sourceName: data.sourceName ?? 'unknown',
        sourceUrl: data.sourceUrl,
        status: data.status ?? 'QUEUED',
      },
    });
  }

  async findById(id: string) {
    return this.prisma.documentImport.findUnique({ where: { id } });
  }

  async updateStatus(id: string, status: string, data?: {
    step?: string;
    progress?: number;
    errorCode?: string;
    errorMessage?: string;
    workerId?: string;
  }) {
    return this.prisma.documentImport.update({
      where: { id },
      data: { status, ...data },
    });
  }

  async claimNext(workerId: string) {
    return this.prisma.documentImport.findFirst({
      where: { status: 'QUEUED' },
      orderBy: { createdAt: 'asc' },
    });
  }

  async claim(id: string, workerId: string) {
    return this.prisma.documentImport.updateMany({
      where: { id, status: 'QUEUED' },
      data: { status: 'CLAIMED', workerId, heartbeatAt: new Date(), startedAt: new Date() },
    });
  }

  async heartbeat(id: string) {
    return this.prisma.documentImport.update({
      where: { id },
      data: { heartbeatAt: new Date() },
    });
  }

  async findStaleJobs() {
    const fiveMinutesAgo = new Date(Date.now() - 5 * 60 * 1000);
    return this.prisma.documentImport.findMany({
      where: {
        status: { in: ['CLAIMED', 'DOWNLOADING', 'PARSING', 'OCR_PROCESSING', 'VISION_PROCESSING', 'CLEANING', 'CHUNKING', 'EMBEDDING', 'INDEXING', 'GENERATING_CANDIDATES'] },
        heartbeatAt: { lt: fiveMinutesAgo },
      },
    });
  }

  async resetStaleJob(id: string, retryCount: number, maxRetries: number) {
    const newStatus = retryCount >= maxRetries ? 'FAILED_FINAL' : 'QUEUED';
    return this.prisma.documentImport.update({
      where: { id },
      data: {
        status: newStatus,
        workerId: null,
        retryCount: { increment: 1 },
      },
    });
  }

  async findBySourceId(sourceId: string) {
    return this.prisma.documentImport.findFirst({
      where: { sourceId },
      orderBy: { createdAt: 'desc' },
    });
  }
}
