import { prisma } from "@codeflash-ai/common"
import { Prisma } from "@prisma/client"
import {
  StagingStorageStrategy,
  StagingStorageContext,
  StagingStorageResult,
  parseSpeedupValue,
  buildMergedMetadata,
} from "./staging-storage-strategy.js"

export interface PlainTextStagingDependencies {
  prisma: {
    optimization_events: {
      update: typeof prisma.optimization_events.update
      findFirst: typeof prisma.optimization_events.findFirst
    }
  }
}

let dependencies: PlainTextStagingDependencies = {
  prisma: {
    optimization_events: {
      update: prisma.optimization_events.update,
      findFirst: prisma.optimization_events.findFirst,
    },
  },
}

export function setPlainTextStagingDependencies(newDependencies: PlainTextStagingDependencies) {
  dependencies = newDependencies
}

export function resetPlainTextStagingDependencies() {
  dependencies = {
    prisma: {
      optimization_events: {
        update: prisma.optimization_events.update,
        findFirst: prisma.optimization_events.findFirst,
      },
    },
  }
}

export class PlainTextStagingStrategy extends StagingStorageStrategy {
  getStorageType(): "plain_text" {
    return "plain_text"
  }

  async save(context: StagingStorageContext): Promise<StagingStorageResult> {
    const {
      userId,
      traceId,
      baseBranch,
      prCommentFields,
      metadata,
      diffContents,
      originalLineProfiler,
      optimizedLineProfiler,
    } = context

    const {
      function_name,
      file_path,
      speedup_x: speedupXRaw,
      speedup_pct: speedupPctRaw,
    } = prCommentFields

    const speedupX = parseSpeedupValue(speedupXRaw, "x")
    if (speedupX === null) {
      return { success: false, storageType: "plain_text", error: "Invalid speedup_x value" }
    }

    const speedupPct = speedupPctRaw ? parseSpeedupValue(speedupPctRaw, "%") : undefined
    if (speedupPctRaw && speedupPct === null) {
      return { success: false, storageType: "plain_text", error: "Invalid speedup_pct value" }
    }

    try {
      const existingEvent = await dependencies.prisma.optimization_events.findFirst({
        where: { trace_id: traceId, user_id: userId },
        select: { metadata: true },
      })

      if (!existingEvent) {
        return {
          success: false,
          storageType: "plain_text",
          error: "No optimization event found for the given traceId",
        }
      }

      const mergedMetadata = buildMergedMetadata(
        existingEvent.metadata,
        metadata,
        {
          storageType: "plain_text",
          diffContents,
        },
        { originalLineProfiler, optimizedLineProfiler },
      )

      await dependencies.prisma.optimization_events.update({
        where: { trace_id: traceId },
        data: {
          function_name: function_name as string,
          file_path: file_path as string,
          speedup_x: speedupX,
          speedup_pct: speedupPct,
          metadata: mergedMetadata as Prisma.InputJsonValue,
          is_optimization_found: true,
          is_staging: true,
          staging_storage_type: "plain_text",
          baseBranch: baseBranch,
          event_type: "no-pr",
        },
      })

      return { success: true, storageType: "plain_text" }
    } catch (error) {
      return {
        success: false,
        storageType: "plain_text",
        error: "Failed to update optimization event",
      }
    }
  }
}
