import type { FineTuneStartPayload, FineTuneTask } from '@/types'

export type FineTuneFormModel = {
  name: string
  description: string
  train_type: 'SFT' | 'DPO' | 'CPT'
  base_model: string | number
  template: string
  train_method: 'lora' | 'full'
  train_dataset_id: string | number
  auto_merge: boolean
  batch_size: number
  learning_rate: number
  n_epochs: number
  save_steps: number
  lr_scheduler_type: string
  max_length: number
  warmup_ratio: number
  weight_decay: number
  lora_alpha: number
  lora_dropout: number
  lora_rank: number
  quantization_bit: number
  export_quantized: boolean
  quant_method: string
  quant_bits: number
  quant_group_size: number
  export_format: string
}

export const DEFAULT_TRAINING_PARAMS = {
  batch_size: 1,
  learning_rate: 0.0001,
  n_epochs: 1,
  save_steps: 100,
  lr_scheduler_type: 'cosine',
  max_length: 512,
  warmup_ratio: 0.05,
  weight_decay: 0.01,
  lora_alpha: 16,
  lora_dropout: 0.1,
  lora_rank: 8,
  quantization_bit: 0,
  export_quantized: false,
  quant_method: 'bnb',
  quant_bits: 4,
  quant_group_size: 128,
  export_format: 'Q4_K_M',
} as const

export function createDefaultFineTuneForm(): FineTuneFormModel {
  return {
    name: '',
    description: '',
    train_type: 'SFT',
    base_model: '',
    template: 'qwen',
    train_method: 'lora',
    train_dataset_id: '',
    auto_merge: false,
    ...DEFAULT_TRAINING_PARAMS,
  }
}

export function buildFineTunePayload(
  form: FineTuneFormModel,
  gpus: number[],
): Omit<FineTuneStartPayload, 'task_id'> {
  return {
    name: form.name,
    description: form.description,
    base_model: form.base_model,
    template: form.template,
    train_type: form.train_type,
    train_method: form.train_method,
    gpus: [...gpus],
    train_dataset_id: form.train_dataset_id,
    auto_merge: form.train_type === 'SFT' && form.auto_merge,
    output_model_name: form.name,
    batch_size: form.batch_size,
    learning_rate: form.learning_rate,
    n_epochs: form.n_epochs,
    save_steps: form.save_steps,
    lr_scheduler_type: form.lr_scheduler_type,
    max_length: form.max_length,
    warmup_ratio: form.warmup_ratio,
    weight_decay: form.weight_decay,
    lora_alpha: form.lora_alpha,
    lora_dropout: form.lora_dropout,
    lora_rank: form.lora_rank,
    quantization_bit: form.train_method === 'lora' ? form.quantization_bit : 0,
    export_quantized: form.export_quantized,
    quant_method: form.export_quantized ? form.quant_method : '',
    quant_bits: form.export_quantized ? form.quant_bits : 0,
    quant_group_size: form.export_quantized ? form.quant_group_size : 0,
    export_format: form.export_quantized && form.quant_method === 'gguf' ? form.export_format : '',
  }
}

export function buildFineTuneCommand(form: FineTuneFormModel, gpus: number[]) {
  const gpuIds = gpus.length ? gpus.join(',') : '0'
  const stage = form.train_type === 'DPO' ? 'dpo' : form.train_type === 'CPT' ? 'cpt' : 'sft'
  const lines = [
    `CUDA_VISIBLE_DEVICES=${gpuIds} llamafactory-cli train`,
    `  --stage ${stage}`,
    '  --do_train',
    '  --model_name_or_path <base_model_path>',
    '  --dataset <dataset>',
    `  --template ${form.template}`,
    `  --finetuning_type ${form.train_method}`,
    `  --output_dir ./saves/${form.name || 'output'}`,
    `  --per_device_train_batch_size ${form.batch_size}`,
    `  --learning_rate ${form.learning_rate}`,
    `  --num_train_epochs ${form.n_epochs}`,
    `  --save_steps ${form.save_steps}`,
    `  --lr_scheduler_type ${form.lr_scheduler_type}`,
    `  --cutoff_len ${form.max_length}`,
    `  --warmup_ratio ${form.warmup_ratio}`,
    `  --weight_decay ${form.weight_decay}`,
  ]

  if (form.train_method === 'lora') {
    lines.push(
      `  --lora_alpha ${form.lora_alpha}`,
      `  --lora_dropout ${form.lora_dropout}`,
      `  --lora_rank ${form.lora_rank}`,
    )
    if (form.quantization_bit) lines.push(`  --quantization_bit ${form.quantization_bit}`)
  }

  if (form.export_quantized) {
    lines.push(`  # 训练后导出量化模型：${form.quant_method} ${form.quant_bits}bit`)
    if (form.quant_method === 'gguf') lines.push(`  # export_format=${form.export_format}`)
    if (form.quant_method === 'gptq' || form.quant_method === 'awq') {
      lines.push(`  # group_size=${form.quant_group_size}`)
    }
  }

  return lines.join(' \\\n')
}

export function toCreateFineTunePayload(
  payload: Omit<FineTuneStartPayload, 'task_id'>,
): Partial<FineTuneTask> {
  return { ...payload, status: 'pending', progress: 0 }
}
