import apiClient from '../client'
import type { TrainingModelsResponse, DocumentListResponse } from '../types'

export const trainingApi = {
  getDocumentsForTraining: async (params?: {
    has_annotations?: boolean
    min_annotation_count?: number
    exclude_used_in_training?: boolean
    limit?: number
    offset?: number
  }): Promise<DocumentListResponse> => {
    const { data } = await apiClient.get('/api/v1/admin/training/documents', {
      params,
    })
    return data
  },

  getModels: async (params?: {
    status?: string
    limit?: number
    offset?: number
  }): Promise<TrainingModelsResponse> => {
    const { data} = await apiClient.get('/api/v1/admin/training/models', {
      params,
    })
    return data
  },

  getTaskDetail: async (taskId: string) => {
    const { data } = await apiClient.get(`/api/v1/admin/training/tasks/${taskId}`)
    return data
  },

  startTraining: async (config: {
    name: string
    description?: string
    document_ids: string[]
    epochs?: number
    batch_size?: number
    model_base?: string
  }) => {
    // Convert frontend config to backend TrainingTaskCreate format
    const taskRequest = {
      name: config.name,
      task_type: 'yolo',
      description: config.description,
      config: {
        document_ids: config.document_ids,
        epochs: config.epochs,
        batch_size: config.batch_size,
        base_model: config.model_base,
      },
    }
    const { data } = await apiClient.post('/api/v1/admin/training/tasks', taskRequest)
    return data
  },

  cancelTask: async (taskId: string) => {
    const { data } = await apiClient.post(
      `/api/v1/admin/training/tasks/${taskId}/cancel`
    )
    return data
  },

  downloadModel: async (taskId: string): Promise<Blob> => {
    const { data } = await apiClient.get(
      `/api/v1/admin/training/models/${taskId}/download`,
      {
        responseType: 'blob',
      }
    )
    return data
  },
}