import prisma from "@shieldai/db";
import { EmbeddingService } from "../embedding/EmbeddingService";
import { FAISSIndex } from "../indexer/FAISSIndex";
import { AudioPreprocessor } from "../preprocessor/AudioPreprocessor";
import { VoiceEnrollmentInput, VoiceEnrollmentOutput } from "@shieldai/types";

export class VoiceEnrollmentService {
  private embeddingService: EmbeddingService;
  private faissIndex: FAISSIndex;
  private preprocessor: AudioPreprocessor;

  constructor() {
    this.embeddingService = new EmbeddingService();
    this.faissIndex = new FAISSIndex();
    this.preprocessor = new AudioPreprocessor();
  }

  async enroll(input: VoiceEnrollmentInput, userId: string): Promise<VoiceEnrollmentOutput> {
    const preprocessed = await this.preprocessor.preprocess(input.audioBuffer, input.sampleRate);

    const embedding = await this.embeddingService.extract(preprocessed.audio);

    const enrollment = await prisma.voiceEnrollment.create({
      data: {
        userId,
        name: input.label,
        voiceHash: this.computeVoiceHash(embedding.vector),
        audioMetadata: {
          sampleRate: preprocessed.sampleRate,
          durationSec: preprocessed.durationSec,
          embeddingDim: embedding.dimension,
        },
      },
    });

    await this.faissIndex.add(enrollment.id, embedding.vector);

    return {
      id: enrollment.id,
      label: enrollment.name,
      embeddingDim: preprocessed.sampleRate,
      sampleRate: preprocessed.sampleRate,
      durationSec: preprocessed.durationSec,
      createdAt: enrollment.createdAt,
    };
  }

  async listEnrollments(userId: string): Promise<VoiceEnrollmentOutput[]> {
    const enrollments = await prisma.voiceEnrollment.findMany({
      where: { userId },
      orderBy: { createdAt: "desc" },
      select: {
        id: true,
        label: true,
        embeddingDim: true,
        sampleRate: true,
        durationSec: true,
        createdAt: true,
      },
    });

    return enrollments.map((e) => ({
      id: e.id,
      label: e.label,
      embeddingDim: e.embeddingDim,
      sampleRate: e.sampleRate,
      durationSec: e.durationSec,
      createdAt: e.createdAt,
    }));
  }

  async removeEnrollment(userId: string, enrollmentId: string): Promise<boolean> {
    const enrollment = await prisma.voiceEnrollment.findFirst({
      where: { id: enrollmentId, userId },
    });

    if (!enrollment) {
      throw new Error(`Enrollment ${enrollmentId} not found for user ${userId}`);
    }

    await prisma.voiceEnrollment.delete({ where: { id: enrollmentId } });
    await this.faissIndex.remove(enrollmentId);

    return true;
  }

  async getEnrollment(enrollmentId: string): Promise<VoiceEnrollmentOutput | null> {
    const enrollment = await prisma.voiceEnrollment.findUnique({
      where: { id: enrollmentId },
      select: {
        id: true,
        label: true,
        embeddingDim: true,
        sampleRate: true,
        durationSec: true,
        createdAt: true,
      },
    });

    if (!enrollment) return null;

    return {
      id: enrollment.id,
      label: enrollment.label,
      embeddingDim: enrollment.embeddingDim,
      sampleRate: enrollment.sampleRate,
      durationSec: enrollment.durationSec,
      createdAt: enrollment.createdAt,
    };
  }

  async matchVoice(
    audioBuffer: Buffer,
    userId: string,
    threshold: number = 0.75
  ): Promise<MatchResult | null> {
    const preprocessed = await this.preprocessor.preprocess(audioBuffer);
    const embedding = await this.embeddingService.extract(preprocessed.audio);

    const matches = await this.faissIndex.search(embedding.vector, 5);

    const enrollmentIds = matches.map((m) => m.id);
    const enrollments = await prisma.voiceEnrollment.findMany({
      where: {
        id: { in: enrollmentIds },
        userId,
      },
    });

    let bestMatch: MatchResult | null = null;

    for (const match of matches) {
      const enrollment = enrollments.find((e) => e.id === match.id);
      if (enrollment && match.similarity >= threshold) {
        if (!bestMatch || match.similarity > bestMatch.similarity) {
          bestMatch = {
            enrollmentId: enrollment.id,
            label: enrollment.label,
            similarity: match.similarity,
          };
        }
      }
    }

    return bestMatch;
  }
}

export interface MatchResult {
  enrollmentId: string;
  label: string;
  similarity: number;
}
