import { Resolver, Mutation, Arg, Ctx, Int, Float, Directive } from 'type-graphql'
import { AIModelClientFactory, DetectedObject } from '@things-factory/ai-inference'
import { labelStudioApi } from '../utils/label-studio-api-client.js'
import {
  GeneratePredictionRequest,
  BatchGeneratePredictionRequest,
  PredictionGenerationResult,
  BatchPredictionResult
} from '../types/prediction-types.js'

/**
 * Label Studio AI Prediction Service
 *
 * Integrates AI inference with Label Studio prediction system
 * Uses ai-inference module for pure AI operations
 */
@Resolver()
export class LabelStudioAIPredictionService {
  /**
   * Generate Label Studio prediction for a task
   * Runs AI model and creates prediction in Label Studio
   */
  @Mutation(returns => PredictionGenerationResult, {
    description: 'Generate Label Studio prediction for a task using AI model'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async generatePrediction(
    @Arg('input') input: GeneratePredictionRequest,
    @Ctx() context: ResolverContext
  ): Promise<PredictionGenerationResult> {
    try {
      // 1. Run AI model using ai-inference module
      const modelClient = input.modelId
        ? AIModelClientFactory.getClient(input.modelId)
        : AIModelClientFactory.getDefaultClient()

      const objects = await modelClient.detectObjects(input.imageUrl, {
        confidenceThreshold: input.confidenceThreshold
      })

      if (objects.length === 0) {
        return {
          taskId: input.taskId,
          success: true,
          objectCount: 0,
          error: 'No objects detected'
        }
      }

      // 2. Convert AI results to Label Studio format
      const labelStudioResult = this.convertToLabelStudioFormat(objects)

      // 3. Calculate average confidence
      const avgConfidence = objects.reduce((sum, obj) => sum + obj.confidence, 0) / objects.length

      // 4. Create prediction in Label Studio
      const prediction = await labelStudioApi.createPrediction({
        task: input.taskId,
        result: labelStudioResult,
        score: avgConfidence,
        model_version: input.modelId || 'default-model-v1.0'
      })

      console.log(
        `[LS AI Prediction] Created prediction ${prediction.id} for task ${input.taskId} with ${objects.length} objects`
      )

      return {
        taskId: input.taskId,
        predictionId: prediction.id,
        success: true,
        objectCount: objects.length,
        avgConfidence
      }
    } catch (error) {
      console.error(`[LS AI Prediction] Failed for task ${input.taskId}:`, error)
      return {
        taskId: input.taskId,
        success: false,
        objectCount: 0,
        error: error.message
      }
    }
  }

  /**
   * Generate predictions for multiple tasks in batch
   */
  @Mutation(returns => BatchPredictionResult, {
    description: 'Generate Label Studio predictions for multiple tasks in batch'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async generateBatchPredictions(
    @Arg('input') input: BatchGeneratePredictionRequest,
    @Ctx() context: ResolverContext
  ): Promise<BatchPredictionResult> {
    const results: PredictionGenerationResult[] = []
    let succeeded = 0
    let failed = 0

    const modelId = input.modelId || 'default-model-v1.0'

    try {
      // Process tasks in parallel (limit concurrency to avoid overload)
      const BATCH_SIZE = 5
      for (let i = 0; i < input.taskIds.length; i += BATCH_SIZE) {
        const batch = input.taskIds.slice(i, i + BATCH_SIZE)

        const batchPromises = batch.map(async taskId => {
          try {
            // Get task data from Label Studio
            const task = await labelStudioApi.getTask(taskId)

            if (!task || !task.data || !task.data.image) {
              failed++
              return {
                taskId,
                success: false,
                objectCount: 0,
                error: 'Task has no image data'
              }
            }

            // Generate prediction
            const result = await this.generatePrediction(
              {
                taskId,
                imageUrl: task.data.image,
                modelId: input.modelId,
                confidenceThreshold: input.confidenceThreshold
              },
              context
            )

            if (result.success) {
              succeeded++
            } else {
              failed++
            }

            return result
          } catch (error) {
            failed++
            return {
              taskId,
              success: false,
              objectCount: 0,
              error: error.message
            }
          }
        })

        const batchResults = await Promise.all(batchPromises)
        results.push(...batchResults)
      }

      console.log(
        `[LS AI Prediction] Batch completed: ${succeeded} succeeded, ${failed} failed out of ${input.taskIds.length} tasks`
      )

      return {
        total: input.taskIds.length,
        succeeded,
        failed,
        results,
        modelVersion: modelId
      }
    } catch (error) {
      console.error('[LS AI Prediction] Batch processing failed:', error)
      throw new Error(`Batch prediction generation failed: ${error.message}`)
    }
  }

  /**
   * Auto-generate predictions for all unlabeled tasks in a project
   */
  @Mutation(returns => BatchPredictionResult, {
    description: 'Auto-generate predictions for all unlabeled tasks in a Label Studio project'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async autoGeneratePredictions(
    @Arg('projectId', type => Int) projectId: number,
    @Arg('modelId', { nullable: true }) modelId: string,
    @Arg('confidenceThreshold', type => Float, { nullable: true }) confidenceThreshold: number,
    @Ctx() context: ResolverContext
  ): Promise<BatchPredictionResult> {
    try {
      // Get all tasks from project
      const tasksResponse = await labelStudioApi.getTasks(projectId, {
        page_size: 1000 // Adjust as needed
      })

      const tasks = tasksResponse.tasks || tasksResponse.results || []

      // Filter unlabeled tasks (no annotations)
      const unlabeledTasks = tasks.filter((task: any) => {
        const annotationCount = task.annotations?.length || task.total_annotations || 0
        return annotationCount === 0
      })

      console.log(
        `[LS AI Prediction] Auto-generating predictions for ${unlabeledTasks.length} unlabeled tasks in project ${projectId}`
      )

      // Generate predictions in batch
      return await this.generateBatchPredictions(
        {
          projectId,
          taskIds: unlabeledTasks.map((t: any) => t.id),
          modelId,
          confidenceThreshold
        },
        context
      )
    } catch (error) {
      console.error(`[LS AI Prediction] Auto-generation failed for project ${projectId}:`, error)
      throw new Error(`Auto prediction generation failed: ${error.message}`)
    }
  }

  /**
   * Convert AI detection results to Label Studio format
   * This handles the Label Studio-specific data structure
   */
  private convertToLabelStudioFormat(objects: DetectedObject[]) {
    return objects.map(obj => ({
      from_name: 'label',
      to_name: 'image',
      type: 'rectanglelabels',
      value: {
        x: obj.bbox.x,
        y: obj.bbox.y,
        width: obj.bbox.width,
        height: obj.bbox.height,
        rectanglelabels: [obj.className]
      }
    }))
  }
}