import { Resolver, Query, Mutation, Arg, Ctx, Int, Directive } from 'type-graphql'
import { labelStudioApi } from '../../utils/label-studio-api-client.js'
import {
  LabelStudioPrediction,
  PredictionInput,
  PredictionImportResult,
  BulkPredictionInput
} from '../../types/label-studio-types.js'

@Resolver()
export class PredictionManagement {
  /**
   * Get predictions for a task
   */
  @Query(returns => [LabelStudioPrediction], {
    description: 'Get all predictions for a Label Studio task'
  })
  @Directive('@privilege(category: "label-studio", privilege: "query")')
  async labelStudioTaskPredictions(
    @Arg('taskId', type => Int) taskId: number,
    @Ctx() context: ResolverContext
  ): Promise<LabelStudioPrediction[]> {
    try {
      const predictions = await labelStudioApi.getPredictions(taskId)

      return predictions.map((prediction: any) => ({
        id: prediction.id,
        taskId: prediction.task || taskId,
        result: JSON.stringify(prediction.result),
        score: prediction.score || undefined,
        modelVersion: prediction.model_version || undefined,
        createdAt: prediction.created_at ? new Date(prediction.created_at) : undefined
      }))
    } catch (error) {
      console.error(`Failed to fetch predictions for task ${taskId}:`, error)
      throw new Error(`Failed to fetch predictions: ${error.message}`)
    }
  }

  /**
   * Get predictions for a project
   */
  @Query(returns => [LabelStudioPrediction], {
    description: 'Get all predictions for a Label Studio project'
  })
  @Directive('@privilege(category: "label-studio", privilege: "query")')
  async labelStudioProjectPredictions(
    @Arg('projectId', type => Int) projectId: number,
    @Ctx() context: ResolverContext
  ): Promise<LabelStudioPrediction[]> {
    try {
      const predictions = await labelStudioApi.getProjectPredictions(projectId)

      return predictions.map((prediction: any) => ({
        id: prediction.id,
        taskId: prediction.task,
        result: JSON.stringify(prediction.result),
        score: prediction.score || undefined,
        modelVersion: prediction.model_version || undefined,
        createdAt: prediction.created_at ? new Date(prediction.created_at) : undefined
      }))
    } catch (error) {
      console.error(`Failed to fetch predictions for project ${projectId}:`, error)
      throw new Error(`Failed to fetch predictions: ${error.message}`)
    }
  }

  /**
   * Get single prediction by ID
   */
  @Query(returns => LabelStudioPrediction, {
    description: 'Get a single prediction by ID',
    nullable: true
  })
  @Directive('@privilege(category: "label-studio", privilege: "query")')
  async labelStudioPrediction(
    @Arg('predictionId', type => Int) predictionId: number,
    @Ctx() context: ResolverContext
  ): Promise<LabelStudioPrediction | null> {
    try {
      const prediction = await labelStudioApi.getPrediction(predictionId)

      if (!prediction) {
        return null
      }

      return {
        id: prediction.id,
        taskId: prediction.task,
        result: JSON.stringify(prediction.result),
        score: prediction.score || undefined,
        modelVersion: prediction.model_version || undefined,
        createdAt: prediction.created_at ? new Date(prediction.created_at) : undefined
      }
    } catch (error) {
      console.error(`Failed to fetch prediction ${predictionId}:`, error)
      return null
    }
  }

  /**
   * Create a prediction for a task
   */
  @Mutation(returns => LabelStudioPrediction, {
    description: 'Create a prediction for a Label Studio task'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async createLabelStudioPrediction(
    @Arg('input') input: PredictionInput,
    @Ctx() context: ResolverContext
  ): Promise<LabelStudioPrediction> {
    try {
      // Parse result JSON
      const result = JSON.parse(input.result)

      // Create prediction
      const prediction = await labelStudioApi.createPrediction({
        task: input.taskId,
        result,
        score: input.score,
        model_version: input.modelVersion
      })

      return {
        id: prediction.id,
        taskId: prediction.task,
        result: JSON.stringify(prediction.result),
        score: prediction.score || undefined,
        modelVersion: prediction.model_version || undefined,
        createdAt: prediction.created_at ? new Date(prediction.created_at) : undefined
      }
    } catch (error) {
      console.error(`Failed to create prediction for task ${input.taskId}:`, error)
      throw new Error(`Failed to create prediction: ${error.message}`)
    }
  }

  /**
   * Create predictions in bulk
   */
  @Mutation(returns => PredictionImportResult, {
    description: 'Create multiple predictions at once'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async createBulkPredictions(
    @Arg('predictions', type => [BulkPredictionInput]) predictions: BulkPredictionInput[],
    @Ctx() context: ResolverContext
  ): Promise<PredictionImportResult> {
    const errors: string[] = []
    let created = 0
    let failed = 0

    try {
      // Parse and validate predictions
      const parsedPredictions = predictions.map((pred, index) => {
        try {
          return {
            task: pred.taskId,
            result: JSON.parse(pred.result),
            score: pred.score,
            model_version: pred.modelVersion
          }
        } catch (e) {
          failed++
          errors.push(`Invalid JSON in prediction ${index}: ${pred.result}`)
          return null
        }
      })

      const validPredictions = parsedPredictions.filter(p => p !== null)

      if (validPredictions.length === 0) {
        return {
          created: 0,
          failed: predictions.length,
          errors
        }
      }

      // Create predictions
      const response = await labelStudioApi.createPredictions(validPredictions)
      created = response.created || validPredictions.length

      return {
        created,
        failed,
        errors: errors.length > 0 ? errors : undefined
      }
    } catch (error) {
      console.error('Failed to create bulk predictions:', error)
      throw new Error(`Failed to create predictions: ${error.message}`)
    }
  }

  /**
   * Import predictions for a project
   * This is the recommended way to bulk import AI model predictions
   */
  @Mutation(returns => PredictionImportResult, {
    description: 'Import predictions for a Label Studio project in bulk'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async importPredictionsToProject(
    @Arg('projectId', type => Int) projectId: number,
    @Arg('predictions', type => [BulkPredictionInput]) predictions: BulkPredictionInput[],
    @Ctx() context: ResolverContext
  ): Promise<PredictionImportResult> {
    const errors: string[] = []
    let created = 0
    let failed = 0

    try {
      // Parse and validate predictions
      const parsedPredictions = predictions.map((pred, index) => {
        try {
          return {
            task: pred.taskId,
            result: JSON.parse(pred.result),
            score: pred.score,
            model_version: pred.modelVersion
          }
        } catch (e) {
          failed++
          errors.push(`Invalid JSON in prediction ${index}: ${pred.result}`)
          return null
        }
      })

      const validPredictions = parsedPredictions.filter(p => p !== null)

      if (validPredictions.length === 0) {
        return {
          created: 0,
          failed: predictions.length,
          errors
        }
      }

      // Import predictions using project endpoint
      const response = await labelStudioApi.importPredictions(projectId, validPredictions)

      // Parse response
      if (response.created !== undefined) {
        created = response.created
      } else if (response.task_count !== undefined) {
        created = response.task_count
      } else if (Array.isArray(response)) {
        created = response.length
      }

      return {
        created,
        failed,
        errors: errors.length > 0 ? errors : undefined
      }
    } catch (error) {
      console.error(`Failed to import predictions to project ${projectId}:`, error)
      throw new Error(`Failed to import predictions: ${error.message}`)
    }
  }

  /**
   * Delete prediction
   */
  @Mutation(returns => Boolean, {
    description: 'Delete a prediction from Label Studio'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async deleteLabelStudioPrediction(
    @Arg('predictionId', type => Int) predictionId: number,
    @Ctx() context: ResolverContext
  ): Promise<boolean> {
    try {
      await labelStudioApi.deletePrediction(predictionId)
      return true
    } catch (error) {
      console.error(`Failed to delete prediction ${predictionId}:`, error)
      throw new Error(`Failed to delete prediction: ${error.message}`)
    }
  }
}