import { ExperimentMetadata } from "./types"
import { prisma } from "@/lib/prisma"

/**
 * Get the modified code for a trace, falling back to original optimized code if no modifications exist
 * @param traceId - The trace ID to get code for
 * @param fileKey - Optional specific file key, if not provided returns all files
 * @returns Modified code or original code if no modifications exist, plus review quality data
 */
export async function getModifiedCodeForTrace(
  traceId: string,
  fileKey?: string,
): Promise<{
  code: { [fileKey: string]: string } | string | null
  review_quality?: string | null
  review_explanation?: string | null
} | null> {
  try {
    const optimizationFeature = await prisma.optimization_features.findUnique({
      where: { trace_id: traceId },
      select: {
        experiment_metadata: true,
        metadata: true,
        review_quality: true,
        review_explanation: true,
      },
    })

    if (!optimizationFeature) {
      return null
    }

    const experimentMetadata = optimizationFeature.experiment_metadata as ExperimentMetadata | null
    const additionalMetadata = (optimizationFeature.metadata as any) || {}

    // Get modified code from metadata if available
    const modifiedCode = additionalMetadata.modifiedCode as { [key: string]: string } | undefined

    let codeResult: { [fileKey: string]: string } | string | null = null

    if (fileKey) {
      // Return specific file
      if (modifiedCode && modifiedCode[fileKey]) {
        codeResult = modifiedCode[fileKey]
      } else if (experimentMetadata?.diffContents?.[fileKey]) {
        // Fall back to original optimized code
        codeResult =
          experimentMetadata.diffContents[fileKey].newContent ||
          experimentMetadata.diffContents[fileKey].oldContent ||
          ""
      }
    } else {
      // Return all files
      const result: { [fileKey: string]: string } = {}

      // Start with original optimized code
      if (experimentMetadata?.diffContents) {
        Object.entries(experimentMetadata.diffContents).forEach(([key, diff]) => {
          result[key] = diff.newContent || diff.oldContent || ""
        })
      }

      // Override with modified code where available
      if (modifiedCode) {
        Object.entries(modifiedCode).forEach(([key, code]) => {
          result[key] = code
        })
      }

      codeResult = Object.keys(result).length > 0 ? result : null
    }

    return {
      code: codeResult,
      review_quality: optimizationFeature.review_quality,
      review_explanation: optimizationFeature.review_explanation,
    }
  } catch (error) {
    console.error("Error getting modified code for trace:", error)
    return null
  }
}

/**
 * Check if a trace has any modified code
 * @param traceId - The trace ID to check
 * @returns True if trace has modified code, false otherwise
 */
export async function hasModifiedCode(traceId: string): Promise<boolean> {
  try {
    const optimizationFeature = await prisma.optimization_features.findUnique({
      where: { trace_id: traceId },
      select: {
        metadata: true,
      },
    })

    if (!optimizationFeature) {
      return false
    }

    const additionalMetadata = (optimizationFeature.metadata as any) || {}
    const modifiedCode = additionalMetadata.modifiedCode as { [key: string]: string } | undefined

    return !!(modifiedCode && Object.keys(modifiedCode).length > 0)
  } catch (error) {
    console.error("Error checking for modified code:", error)
    return false
  }
}

/**
 * Get the most recent code for suggestions - modified code takes precedence over original
 * This is the main function other services should use when they need code for suggestions
 * @param traceId - The trace ID
 * @returns Object with file paths as keys and code content as values
 */
export async function getCodeForSuggestions(
  traceId: string,
): Promise<{ [fileKey: string]: string } | null> {
  const result = await getModifiedCodeForTrace(traceId)
  if (!result) return null
  return typeof result.code === "object" && result.code !== null ? result.code : null
}
