import { Resolver, Query, Mutation, Arg, Ctx, Int, Directive } from 'type-graphql'
import { getRepository } from '@things-factory/shell'
import { AIModelClientFactory } from '@things-factory/ai-inference'
import { DataSample, DataSet } from '@things-factory/dataset'
import { labelStudioApi } from '../utils/label-studio-api-client.js'
import {
  CreateLabelingTasksRequest,
  TaskCreationResult,
  SyncAnnotationsRequest,
  SyncAnnotationsResult,
  DatasetLabelingStatus,
  DatasetLabelingStatusRequest,
  GeneratePredictionsForDatasetRequest
} from '../types/dataset-labeling-types.js'
import { BatchPredictionResult } from '../types/prediction-types.js'
import { extractMediaUrls, getBaseUrl } from '../utils/media-url-extractor.js'

/**
 * Dataset Labeling Integration Service
 *
 * Integrates Things-Factory Dataset module with Label Studio for AI-assisted labeling
 *
 * Features:
 * - Create Label Studio tasks from DataSamples
 * - Auto-generate AI predictions for labeling tasks
 * - Sync human annotations back to DataSamples
 * - Track labeling progress and status
 */
@Resolver()
export class DatasetLabelingIntegration {
  /**
   * Create Label Studio labeling tasks from DataSamples in a DataSet
   * Optionally generates AI predictions automatically
   */
  @Mutation(returns => TaskCreationResult, {
    description:
      'Create Label Studio labeling tasks from DataSamples in a DataSet. Optionally auto-generates AI predictions for each task.'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async createLabelingTasksFromDataset(
    @Arg('input') input: CreateLabelingTasksRequest,
    @Ctx() context: ResolverContext
  ): Promise<TaskCreationResult> {
    const { domain } = context.state

    console.log(
      `[Dataset Labeling] Creating tasks for DataSet ${input.dataSetId} in LS project ${input.projectId}`
    )

    try {
      // 1. Get DataSet configuration
      const dataSet = await getRepository(DataSet).findOne({
        where: { domain: { id: domain.id }, id: input.dataSetId }
      })

      if (!dataSet) {
        throw new Error(`DataSet not found: ${input.dataSetId}`)
      }

      // 2. Query DataSamples
      const queryBuilder = getRepository(DataSample)
        .createQueryBuilder('sample')
        .where('sample.domainId = :domainId', { domainId: domain.id })
        .andWhere('sample.dataSetId = :dataSetId', { dataSetId: input.dataSetId })

      if (input.sinceDate) {
        queryBuilder.andWhere('sample.collectedAt >= :sinceDate', { sinceDate: input.sinceDate })
      }

      if (input.limit) {
        queryBuilder.limit(input.limit)
      }

      queryBuilder.orderBy('sample.collectedAt', 'DESC')

      const samples = await queryBuilder.getMany()

      console.log(`[Dataset Labeling] Found ${samples.length} DataSamples to process`)

      // 3. Create tasks for each sample
      const result: TaskCreationResult = {
        totalSamples: samples.length,
        tasksCreated: 0,
        tasksFailed: 0,
        tasksSkipped: 0,
        taskIds: [],
        predictionsCreated: 0
      }

      // Get base URL for converting relative attachment paths to absolute URLs
      const baseUrl = getBaseUrl(context)

      for (const sample of samples) {
        try {
          // Extract all media URLs from DataSample (image, video, audio)
          const mediaUrls = extractMediaUrls(sample, baseUrl)

          if (Object.keys(mediaUrls).length === 0) {
            console.log(`[Dataset Labeling] Skipping sample ${sample.id}: no media found`)
            result.tasksSkipped++
            continue
          }

          // Create Label Studio task with all media URLs
          const task = await labelStudioApi.createTask(input.projectId, {
            data: mediaUrls,
            meta: {
              dataSampleId: sample.id,
              dataSampleName: sample.name,
              dataSetId: input.dataSetId,
              dataSetName: dataSet.name,
              collectedAt: sample.collectedAt?.toISOString(),
              source: 'things-factory-dataset'
            }
          })

          result.taskIds.push(task.id)
          result.tasksCreated++

          console.log(`[Dataset Labeling] Created task ${task.id} for sample ${sample.id}`)

          // 4. Auto-generate AI prediction if requested (for image detection)
          if (input.autoGeneratePredictions && mediaUrls.image) {
            try {
              const modelClient = input.modelId
                ? AIModelClientFactory.getClient(input.modelId)
                : AIModelClientFactory.getDefaultClient()

              // Use first image URL if multiple images
              const imageUrl = Array.isArray(mediaUrls.image) ? mediaUrls.image[0] : mediaUrls.image

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

              if (objects.length > 0) {
                const labelStudioResult = this.convertToLabelStudioFormat(objects)
                const avgConfidence = objects.reduce((sum, obj) => sum + obj.confidence, 0) / objects.length

                await labelStudioApi.createPrediction({
                  task: task.id,
                  result: labelStudioResult,
                  score: avgConfidence,
                  model_version: input.modelId || 'default-model-v1.0'
                })

                result.predictionsCreated = (result.predictionsCreated || 0) + 1

                console.log(`[Dataset Labeling] Created prediction for task ${task.id} (${objects.length} objects)`)
              }
            } catch (predError) {
              console.error(`[Dataset Labeling] Failed to create prediction for task ${task.id}:`, predError)
              // Continue even if prediction fails
            }
          }
        } catch (taskError) {
          console.error(`[Dataset Labeling] Failed to create task for sample ${sample.id}:`, taskError)
          result.tasksFailed++
        }
      }

      console.log(
        `[Dataset Labeling] Completed: ${result.tasksCreated} created, ${result.tasksFailed} failed, ${result.tasksSkipped} skipped`
      )

      return result
    } catch (error) {
      console.error('[Dataset Labeling] Failed to create tasks:', error)
      throw new Error(`Failed to create labeling tasks: ${error.message}`)
    }
  }

  /**
   * Sync completed annotations from Label Studio back to DataSamples
   */
  @Mutation(returns => SyncAnnotationsResult, {
    description: 'Sync Label Studio annotations back to DataSamples, updating judgment fields'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async syncAnnotationsToDataset(
    @Arg('input') input: SyncAnnotationsRequest,
    @Ctx() context: ResolverContext
  ): Promise<SyncAnnotationsResult> {
    const { domain } = context.state

    console.log(`[Dataset Labeling] Syncing annotations from LS project ${input.projectId} to DataSet ${input.dataSetId}`)

    try {
      // 1. Get all tasks from Label Studio project
      const tasksResponse = await labelStudioApi.getTasks(input.projectId, {
        page_size: 1000 // Adjust as needed
      })

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

      console.log(`[Dataset Labeling] Found ${tasks.length} tasks in Label Studio`)

      // 2. Filter tasks related to this DataSet
      const dataSetTasks = tasks.filter((task: any) => task.meta?.dataSetId === input.dataSetId)

      console.log(`[Dataset Labeling] ${dataSetTasks.length} tasks belong to DataSet ${input.dataSetId}`)

      const result: SyncAnnotationsResult = {
        totalAnnotations: 0,
        samplesUpdated: 0,
        updatesFailed: 0,
        skipped: 0
      }

      // 3. Process each task
      for (const task of dataSetTasks) {
        const dataSampleId = task.meta?.dataSampleId

        if (!dataSampleId) {
          result.skipped++
          continue
        }

        // Get annotations for this task
        const annotations = task.annotations || []

        if (annotations.length === 0) {
          result.skipped++
          continue
        }

        // Filter by completion status if requested
        let relevantAnnotations = annotations
        if (input.completedOnly) {
          relevantAnnotations = annotations.filter((ann: any) => ann.was_cancelled === false)
        }

        // Filter by date if requested
        if (input.sinceDate) {
          relevantAnnotations = relevantAnnotations.filter(
            (ann: any) => new Date(ann.updated_at) >= input.sinceDate
          )
        }

        if (relevantAnnotations.length === 0) {
          result.skipped++
          continue
        }

        result.totalAnnotations += relevantAnnotations.length

        try {
          // 4. Get DataSample
          const sample = await getRepository(DataSample).findOne({
            where: { domain: { id: domain.id }, id: dataSampleId }
          })

          if (!sample) {
            console.log(`[Dataset Labeling] DataSample not found: ${dataSampleId}`)
            result.updatesFailed++
            continue
          }

          // 5. Convert annotations to judgment data
          const latestAnnotation = relevantAnnotations[relevantAnnotations.length - 1]
          const judgment = this.convertAnnotationToJudgment(latestAnnotation, task)

          // 6. Update DataSample
          sample.judgment = judgment

          // Optionally update OOC/OOS flags based on annotation
          // This is application-specific logic
          if (judgment.defectsFound > 0) {
            sample.ooc = true
          }

          await getRepository(DataSample).save(sample)

          result.samplesUpdated++

          console.log(`[Dataset Labeling] Updated DataSample ${dataSampleId} with annotation ${latestAnnotation.id}`)
        } catch (updateError) {
          console.error(`[Dataset Labeling] Failed to update DataSample ${dataSampleId}:`, updateError)
          result.updatesFailed++
        }
      }

      console.log(
        `[Dataset Labeling] Sync completed: ${result.samplesUpdated} updated, ${result.updatesFailed} failed, ${result.skipped} skipped`
      )

      return result
    } catch (error) {
      console.error('[Dataset Labeling] Failed to sync annotations:', error)
      throw new Error(`Failed to sync annotations: ${error.message}`)
    }
  }

  /**
   * Generate AI predictions for existing DataSet samples
   */
  @Mutation(returns => BatchPredictionResult, {
    description: 'Generate AI predictions for DataSet samples that already have Label Studio tasks'
  })
  @Directive('@privilege(category: "label-studio", privilege: "mutation")')
  async generatePredictionsForDataset(
    @Arg('input') input: GeneratePredictionsForDatasetRequest,
    @Ctx() context: ResolverContext
  ): Promise<BatchPredictionResult> {
    const { domain } = context.state

    console.log(
      `[Dataset Labeling] Generating predictions for DataSet ${input.dataSetId} in LS project ${input.projectId}`
    )

    try {
      // 1. Get all tasks from Label Studio project
      const tasksResponse = await labelStudioApi.getTasks(input.projectId, {
        page_size: 1000
      })

      const tasks = tasksResponse.tasks || tasksResponse.results || []
      const dataSetTasks = tasks.filter((task: any) => task.meta?.dataSetId === input.dataSetId)

      console.log(`[Dataset Labeling] Found ${dataSetTasks.length} tasks for DataSet`)

      const result: BatchPredictionResult = {
        total: dataSetTasks.length,
        succeeded: 0,
        failed: 0,
        results: [],
        modelVersion: input.modelId || 'default-model-v1.0'
      }

      // 2. Get AI model client
      const modelClient = input.modelId
        ? AIModelClientFactory.getClient(input.modelId)
        : AIModelClientFactory.getDefaultClient()

      // 3. Process tasks in batches
      const BATCH_SIZE = 5
      for (let i = 0; i < dataSetTasks.length; i += BATCH_SIZE) {
        const batch = dataSetTasks.slice(i, i + BATCH_SIZE)

        const batchPromises = batch.map(async (task: any) => {
          try {
            // Check if prediction already exists
            if (!input.forceRegenerate && task.predictions && task.predictions.length > 0) {
              result.results.push({
                taskId: task.id,
                success: true,
                objectCount: 0,
                error: 'Prediction already exists (skipped)'
              })
              return
            }

            const imageUrl = task.data?.image

            if (!imageUrl) {
              result.failed++
              result.results.push({
                taskId: task.id,
                success: false,
                objectCount: 0,
                error: 'No image URL in task data'
              })
              return
            }

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

            if (objects.length === 0) {
              result.succeeded++
              result.results.push({
                taskId: task.id,
                success: true,
                objectCount: 0,
                error: 'No objects detected'
              })
              return
            }

            // Create prediction
            const labelStudioResult = this.convertToLabelStudioFormat(objects)
            const avgConfidence = objects.reduce((sum, obj) => sum + obj.confidence, 0) / objects.length

            const prediction = await labelStudioApi.createPrediction({
              task: task.id,
              result: labelStudioResult,
              score: avgConfidence,
              model_version: result.modelVersion
            })

            result.succeeded++
            result.results.push({
              taskId: task.id,
              predictionId: prediction.id,
              success: true,
              objectCount: objects.length,
              avgConfidence
            })

            console.log(`[Dataset Labeling] Created prediction for task ${task.id} (${objects.length} objects)`)
          } catch (error) {
            result.failed++
            result.results.push({
              taskId: task.id,
              success: false,
              objectCount: 0,
              error: error.message
            })
          }
        })

        await Promise.all(batchPromises)
      }

      console.log(`[Dataset Labeling] Prediction generation completed: ${result.succeeded}/${result.total} succeeded`)

      return result
    } catch (error) {
      console.error('[Dataset Labeling] Failed to generate predictions:', error)
      throw new Error(`Failed to generate predictions: ${error.message}`)
    }
  }

  /**
   * Query labeling status for a DataSet
   */
  @Query(returns => DatasetLabelingStatus, {
    description: 'Get labeling status and progress for a DataSet'
  })
  @Directive('@privilege(category: "label-studio", privilege: "query")')
  async datasetLabelingStatus(
    @Arg('input') input: DatasetLabelingStatusRequest,
    @Ctx() context: ResolverContext
  ): Promise<DatasetLabelingStatus> {
    const { domain } = context.state

    try {
      // 1. Get DataSet
      const dataSet = await getRepository(DataSet).findOne({
        where: { domain: { id: domain.id }, id: input.dataSetId }
      })

      if (!dataSet) {
        throw new Error(`DataSet not found: ${input.dataSetId}`)
      }

      // 2. Count total DataSamples
      const totalSamples = await getRepository(DataSample).count({
        where: { domain: { id: domain.id }, dataSetId: input.dataSetId }
      })

      // 3. Get Label Studio tasks if projectId provided
      let tasksCreated = 0
      let withPredictions = 0
      let withAnnotations = 0
      let annotationsCompleted = 0

      if (input.projectId) {
        const tasksResponse = await labelStudioApi.getTasks(input.projectId, {
          page_size: 1000
        })

        const tasks = tasksResponse.tasks || tasksResponse.results || []
        const dataSetTasks = tasks.filter((task: any) => task.meta?.dataSetId === input.dataSetId)

        tasksCreated = dataSetTasks.length

        for (const task of dataSetTasks) {
          if (task.predictions && task.predictions.length > 0) {
            withPredictions++
          }

          if (task.annotations && task.annotations.length > 0) {
            withAnnotations++

            const completedAnnotations = task.annotations.filter((ann: any) => ann.was_cancelled === false)
            if (completedAnnotations.length > 0) {
              annotationsCompleted++
            }
          }
        }
      }

      const notProcessed = totalSamples - tasksCreated
      const completionRate = totalSamples > 0 ? annotationsCompleted / totalSamples : 0

      return {
        dataSetId: input.dataSetId,
        dataSetName: dataSet.name,
        totalSamples,
        tasksCreated,
        withPredictions,
        withAnnotations,
        annotationsCompleted,
        notProcessed,
        completionRate,
        projectId: input.projectId,
        lastSyncedAt: new Date()
      }
    } catch (error) {
      console.error('[Dataset Labeling] Failed to get status:', error)
      throw new Error(`Failed to get labeling status: ${error.message}`)
    }
  }

  /**
   * Helper: Extract image URL from DataSample
   */
  private extractImageUrl(sample: DataSample, imageField: string): string | null {
    // Try data field first
    if (sample.data && typeof sample.data === 'object') {
      if (sample.data[imageField]) {
        return sample.data[imageField]
      }
    }

    // Try rawData field
    if (sample.rawData) {
      try {
        const rawData = typeof sample.rawData === 'string' ? JSON.parse(sample.rawData) : sample.rawData
        if (rawData[imageField]) {
          return rawData[imageField]
        }
        // If rawData is just a string URL
        if (typeof rawData === 'string' && rawData.startsWith('http')) {
          return rawData
        }
      } catch (e) {
        // If rawData is a plain URL string
        if (typeof sample.rawData === 'string' && sample.rawData.startsWith('http')) {
          return sample.rawData
        }
      }
    }

    return null
  }

  /**
   * Helper: Convert AI detection results to Label Studio format
   */
  private convertToLabelStudioFormat(objects: any[]) {
    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]
      }
    }))
  }

  /**
   * Helper: Convert Label Studio annotation to DataSample judgment format
   */
  private convertAnnotationToJudgment(annotation: any, task: any): any {
    const result = annotation.result || []

    const objects = result
      .filter((item: any) => item.type === 'rectanglelabels')
      .map((item: any) => ({
        type: item.value.rectanglelabels?.[0] || 'unknown',
        bbox: {
          x: item.value.x,
          y: item.value.y,
          width: item.value.width,
          height: item.value.height
        },
        source: 'human-annotation'
      }))

    return {
      annotationId: annotation.id,
      taskId: task.id,
      annotatedBy: annotation.completed_by,
      annotatedAt: annotation.updated_at,
      objects,
      objectCount: objects.length,
      defectsFound: objects.length, // Application-specific logic
      verified: true,
      source: 'label-studio'
    }
  }
}