<script setup lang="ts">
import { ref, reactive, computed, onMounted } from 'vue'
import { useRouter } from 'vue-router'
import { ElMessage, type FormInstance, type FormRules } from 'element-plus'
import PageCard from '@/components/PageCard.vue'
import ModelSelectDialog from '@/components/ModelSelectDialog.vue'
import {
  createFineTune,
  startFineTune,
  updateFineTune,
  checkFineTuneName,
} from '@/api/modules/fineTune'
import { getModelList } from '@/api/modules/model'
import { getDatasetList } from '@/api/modules/dataset'
import { getSystemInfo } from '@/api/modules/system'
import { TEMPLATE_GROUPS, LR_SCHEDULER_OPTIONS } from '@/constants'
import type { ModelItem, DatasetItem, GpuInfo } from '@/types'

const router = useRouter()
const formRef = ref<FormInstance>()
const submitting = ref(false)

const models = ref<ModelItem[]>([])
const datasets = ref<DatasetItem[]>([])
const gpus = ref<GpuInfo[]>([])
const selectedGpus = ref<number[]>([])
const modelDialogVisible = ref(false)

const form = reactive({
  name: '',
  description: '',
  train_type: 'SFT' as 'SFT' | 'DPO' | 'CPT',
  base_model: '' as string | number,
  template: 'qwen',
  train_method: 'lora' as 'lora' | 'full',
  train_dataset_id: '' as string | number,
  auto_merge: false,
  // 训练参数
  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 参数
  lora_alpha: 16, // 修复原项目 lora_alpha 默认值不一致 bug
  lora_dropout: 0.1,
  lora_rank: 8,
})

const rules: FormRules = {
  name: [
    { required: true, message: '请输入任务名称', trigger: 'blur' },
    {
      pattern: /^[a-zA-Z0-9_]+$/,
      message: '仅支持字母、数字、下划线',
      trigger: 'blur',
    },
    { max: 50, message: '不超过 50 字符', trigger: 'blur' },
  ],
  base_model: [{ required: true, message: '请选择模型', trigger: 'change' }],
  template: [{ required: true, message: '请选择训练模板', trigger: 'change' }],
  train_dataset_id: [{ required: true, message: '请选择训练数据集', trigger: 'change' }],
}

const showLoraParams = computed(() => form.train_method === 'lora')

const selectedModel = computed(() => models.value.find((model) => model.id === form.base_model))

const modelDialogTitle = computed(() => selectedModel.value?.name || '')

/** 训练命令实时预览 */
const commandPreview = computed(() => {
  const gpuIds = selectedGpus.value.length ? selectedGpus.value.join(',') : '0'
  let cmd = `CUDA_VISIBLE_DEVICES=${gpuIds} llamafactory-cli train \\\n`
  cmd += `  --stage ${form.train_type === 'DPO' ? 'dpo' : form.train_type === 'CPT' ? 'cpt' : 'sft'} \\\n`
  cmd += `  --do_train \\\n`
  cmd += `  --model_name_or_path <base_model_path> \\\n`
  cmd += `  --dataset <dataset> \\\n`
  cmd += `  --template ${form.template} \\\n`
  cmd += `  --finetuning_type ${form.train_method} \\\n`
  cmd += `  --output_dir ./saves/${form.name || 'output'} \\\n`
  cmd += `  --per_device_train_batch_size ${form.batch_size} \\\n`
  cmd += `  --learning_rate ${form.learning_rate} \\\n`
  cmd += `  --num_train_epochs ${form.n_epochs} \\\n`
  cmd += `  --save_steps ${form.save_steps} \\\n`
  cmd += `  --lr_scheduler_type ${form.lr_scheduler_type} \\\n`
  cmd += `  --cutoff_len ${form.max_length} \\\n`
  cmd += `  --warmup_ratio ${form.warmup_ratio} \\\n`
  cmd += `  --weight_decay ${form.weight_decay}`
  if (showLoraParams.value) {
    cmd += ` \\\n  --lora_alpha ${form.lora_alpha}`
    cmd += ` \\\n  --lora_dropout ${form.lora_dropout}`
    cmd += ` \\\n  --lora_rank ${form.lora_rank}`
  }
  return cmd
})

/** GPU 多选切换 */
function toggleGpu(index: number) {
  const idx = selectedGpus.value.indexOf(index)
  if (idx === -1) selectedGpus.value.push(index)
  else selectedGpus.value.splice(idx, 1)
}

function gpuUsageWidth(percent: number) {
  return `${Math.max(0, Math.min(percent, 100))}%`
}

function openModelDialog() {
  modelDialogVisible.value = true
}

/** 模型选择弹窗确认 */
function handleModelConfirm(modelId: string | number) {
  form.base_model = modelId
  formRef.value?.validateField('base_model')
}

function resetParams() {
  Object.assign(form, {
    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,
  })
}

const isParamsExpanded = ref(false)

const allParams = computed(() => {
  const params = [
    { key: 'batch_size', name: 'batch_size', desc: '批次大小，代表模型训练过程中，模型更新一次参数所需要的数据样本数。', hint: '[1, 64], step:1', type: 'number', min: 1, max: 64, step: 1 },
    { key: 'learning_rate', name: 'learning_rate', desc: '学习率，代表每次更新数据的增量参数权重比例。', hint: '[0.000001, 1]', type: 'number', min: 0.000001, max: 1, step: 0.00001, precision: 6 },
    { key: 'n_epochs', name: 'n_epochs', desc: '循环次数，代表模型训练过程中模型学习数据集的次数，可理解为看几遍数据，一般建议的范围是 1-3 遍即可，可依据需求进行调整', hint: '[1, 100], step:1', type: 'number', min: 1, max: 100, step: 1 },
    { key: 'save_steps', name: 'save_steps', desc: '保存步数，训练阶段模型保存的间隔步长。', hint: '[10, 10000]', type: 'number', min: 10, max: 10000, step: 1 },
    { key: 'lr_scheduler_type', name: 'lr_scheduler_type', desc: '学习率调整策略，选择不同的学习率策略。', hint: '', type: 'select', options: LR_SCHEDULER_OPTIONS },
    { key: 'max_length', name: 'max_length', desc: '序列长度，单个训练数据样本的最大长度。', hint: '[64, 4096]', type: 'number', min: 64, max: 4096, step: 1 },
    { key: 'warmup_ratio', name: 'warmup_ratio', desc: '学习率预热比例，学习率预热阶段占总训练步数的比例。', hint: '[0, 1]', type: 'number', min: 0, max: 1, step: 0.01, precision: 2 },
    { key: 'weight_decay', name: 'weight_decay', desc: '权重衰减，用于在优化过程中对模型参数施加惩罚，防止过拟合。', hint: '[0, 1]', type: 'number', min: 0, max: 1, step: 0.01, precision: 2 },
  ]
  if (showLoraParams.value) {
    params.push(
      { key: 'lora_alpha', name: 'lora_alpha', desc: 'LoRA 缩放系数。', hint: '16/32/64/128', type: 'select', options: [{label:'16',value:16},{label:'32',value:32},{label:'64',value:64},{label:'128',value:128}] },
      { key: 'lora_rank', name: 'lora_rank', desc: 'LoRA 秩大小，控制低秩矩阵的维度。', hint: '8/16/32/64', type: 'select', options: [{label:'8',value:8},{label:'16',value:16},{label:'32',value:32},{label:'64',value:64}] },
      { key: 'lora_dropout', name: 'lora_dropout', desc: 'LoRA 层的 dropout 比例。', hint: '[0, 1]', type: 'number', min: 0, max: 1, step: 0.05, precision: 2 }
    )
  }
  return params
})

const visibleParams = computed(() => {
  return isParamsExpanded.value ? allParams.value : allParams.value.slice(0, 3)
})

async function loadModels() {
  try {
    models.value = (await getModelList()) || []
  } catch {
    models.value = []
  }
}

async function loadDatasets() {
  try {
    datasets.value = (await getDatasetList()) || []
  } catch {
    datasets.value = []
  }
}

async function loadGpus() {
  try {
    const sys = await getSystemInfo()
    gpus.value = sys?.gpu || []
    // 默认选中第一个
    if (gpus.value.length > 0) selectedGpus.value = [0]
  } catch {
    gpus.value = []
  }
}

async function handleSubmit() {
  if (!formRef.value) return
  await formRef.value.validate(async (valid) => {
    if (!valid) return
    if (selectedGpus.value.length === 0) {
      ElMessage.warning('请至少选择一个 GPU')
      return
    }
    submitting.value = true
    try {
      // 任务名查重
      const check = await checkFineTuneName(form.name).catch(() => ({ exists: false }))
      if ((check as any).exists) {
        ElMessage.error('任务名称已存在，请更换')
        submitting.value = false
        return
      }

      // 第一步：创建任务记录
      const taskData = {
        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: selectedGpus.value,
        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,
        status: 'pending',
        progress: 0,
      }
      const createRes: any = await createFineTune(taskData)
      const taskId = createRes?.id || createRes

      // 第二步：启动训练
      try {
        await startFineTune({
          task_id: taskId,
          name: form.name,
          base_model: form.base_model,
          template: form.template,
          train_type: form.train_type,
          train_method: form.train_method,
          train_dataset_id: form.train_dataset_id,
          auto_merge: form.train_type === 'SFT' && form.auto_merge,
          output_model_name: form.name,
          gpus: selectedGpus.value,
          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,
        })
        ElMessage.success('训练任务已创建并启动')
      } catch (e) {
        // 启动失败，回写状态
        await updateFineTune(taskId, { status: 'failed' })
      }
      router.push('/fine-tune')
    } catch {
      // ignore
    } finally {
      submitting.value = false
    }
  })
}

function handleCancel() {
  router.back()
}

onMounted(() => {
  loadModels()
  loadDatasets()
  loadGpus()
})
</script>

<template>
  <div class="fine-tune-create page-card-host has-fixed-footer">
    <PageCard title="创建训练任务">
      <el-form ref="formRef" :model="form" :rules="rules" label-width="140px" label-position="left">
        <!-- 基本信息 -->
        <el-divider content-position="left">基本信息</el-divider>
        <el-form-item label="任务名称" prop="name">
          <el-input v-model="form.name" placeholder="字母、数字、下划线" maxlength="50" show-word-limit style="width: 420px;" />
        </el-form-item>
        <el-form-item label="任务描述">
          <el-input v-model="form.description" type="textarea" :rows="4" maxlength="200" show-word-limit style="width: 600px;" />
        </el-form-item>

        <!-- 训练配置 -->
        <el-divider content-position="left">训练配置</el-divider>
        <el-form-item label="GPU 硬件">
          <div class="gpu-list">
            <div
              v-for="(gpu, idx) in gpus"
              :key="idx"
              class="gpu-card"
              :class="{ active: selectedGpus.includes(idx), 'is-busy': gpu.gpu_percent > 80 }"
              @click="toggleGpu(idx)"
            >
              <div class="gpu-card-top">
                <div class="gpu-title">
                  <span class="gpu-index">GPU-{{ idx }}</span>
                  <span class="gpu-name">{{ gpu.name }}</span>
                </div>
                <span class="gpu-usage">{{ gpu.gpu_percent }}%</span>
              </div>
              <div class="gpu-usage-bar">
                <span :style="{ width: gpuUsageWidth(gpu.gpu_percent) }" />
              </div>
              <div class="gpu-meta">
                <span>显存 {{ gpu.memory_used_gb }}/{{ gpu.memory_total_gb }}GB</span>
                <span>{{ gpu.temperature }}°C</span>
                <span>{{ gpu.power_w }}W</span>
              </div>
            </div>
            <div v-if="!gpus.length" class="gpu-empty">暂无 GPU 信息</div>
          </div>
        </el-form-item>

        <el-form-item label="训练方式">
          <el-radio-group v-model="form.train_type">
            <el-radio-button value="SFT">SFT 微调训练</el-radio-button>
            <el-radio-button value="DPO">DPO 偏好训练</el-radio-button>
            <el-radio-button value="CPT">CPT 继续预训练</el-radio-button>
          </el-radio-group>
        </el-form-item>

        <el-form-item label="选择模型" prop="base_model">
          <el-input
            class="model-picker-input"
            :model-value="modelDialogTitle"
            placeholder="请选择基座模型"
            readonly
            @click="openModelDialog"
            style="width: 420px;"
          >
            <template #suffix>
              <i class="fa fa-angle-right" />
            </template>
          </el-input>
        </el-form-item>

        <el-form-item label="训练模板" prop="template">
          <el-select v-model="form.template" placeholder="请选择训练模板" filterable style="width: 420px;">
            <el-option-group v-for="group in TEMPLATE_GROUPS" :key="group.label" :label="group.label">
              <el-option v-for="opt in group.options" :key="opt.value" :label="opt.label" :value="opt.value" />
            </el-option-group>
          </el-select>
        </el-form-item>

        <el-form-item label="训练方法">
          <el-radio-group v-model="form.train_method">
            <el-radio value="lora">LoRA（高效微调）</el-radio>
            <el-radio value="full">全参微调</el-radio>
          </el-radio-group>
        </el-form-item>

        <!-- 超参配置 -->
        <el-divider content-position="left">
          <span style="font-size: 15px; font-weight: 600; color: #1f2937; margin-right: 12px;">超参配置</span>
          <el-button link type="primary" size="small" @click="resetParams">
            <i class="fa fa-refresh" style="margin-right: 4px;" /> 恢复默认配置
          </el-button>
        </el-divider>

        <div class="hyperparam-section">
          <div class="param-table-wrapper">
            <div class="param-table">
              <div class="param-header">
                <div class="param-col name">参数名称</div>
                <div class="param-col config">配置</div>
                <div class="param-col desc">说明</div>
              </div>
              <div class="param-row" v-for="param in visibleParams" :key="param.key">
                <div class="param-col name">{{ param.name }}</div>
                <div class="param-col config">
                  <el-input-number 
                    v-if="param.type === 'number'" 
                    v-model="(form as any)[param.key]" 
                    :min="param.min" :max="param.max" :step="param.step" :precision="param.precision" 
                    controls-position="right"
                    style="width: 200px"
                  />
                  <el-select 
                    v-else-if="param.type === 'select'" 
                    v-model="(form as any)[param.key]" 
                    style="width: 200px"
                  >
                    <el-option v-for="o in param.options" :key="o.value" :label="o.label" :value="o.value" />
                  </el-select>
                  <span class="param-hint" v-if="param.hint">{{ param.hint }}</span>
                </div>
                <div class="param-col desc">
                  <el-tooltip :content="param.desc" placement="top" effect="dark" :show-after="200">
                    <span class="desc-text">{{ param.desc }}</span>
                  </el-tooltip>
                </div>
              </div>
            </div>
            <div class="param-footer">
              <el-button link type="primary" @click="isParamsExpanded = !isParamsExpanded">
                <i :class="isParamsExpanded ? 'fa fa-angle-up' : 'fa fa-angle-down'" style="margin-right: 4px;" />
                {{ isParamsExpanded ? '收起配置' : '展开配置' }}
              </el-button>
            </div>
          </div>
        </div>

        <!-- 数据配置 -->
        <el-divider content-position="left">数据配置</el-divider>
        <el-form-item label="训练数据集" prop="train_dataset_id">
          <el-select v-model="form.train_dataset_id" placeholder="请选择训练数据集" filterable style="width: 420px;">
            <el-option v-for="d in datasets" :key="d.id" :label="d.name" :value="d.id" />
          </el-select>
        </el-form-item>

        <template v-if="form.train_type === 'SFT'">
          <el-divider content-position="left">合并模型</el-divider>
          <el-form-item label="自动合并权重并保存">
            <el-select v-model="form.auto_merge" style="width: 420px;">
              <el-option label="否" :value="false" />
              <el-option label="是" :value="true" />
            </el-select>
          </el-form-item>
        </template>

        <!-- 训练命令预览 -->
        <el-divider content-position="left">训练命令预览</el-divider>
        <div class="command-preview-wrapper">
          <pre class="command-preview">{{ commandPreview }}</pre>
        </div>

        <div class="form-actions-wrapper">
          <el-button type="primary" :loading="submitting" @click="handleSubmit">创建并启动训练</el-button>
          <el-button @click="handleCancel">取消</el-button>
        </div>
      </el-form>
    </PageCard>

    <ModelSelectDialog
      v-model="modelDialogVisible"
      :models="models"
      :current-model-id="form.base_model"
      @confirm="handleModelConfirm"
    />
  </div>
</template>

<style scoped lang="scss">
.fine-tune-create {
  :deep(.page-card) {
    margin-bottom: 0;
  }
}

.gpu-list {
  display: grid;
  grid-template-columns: repeat(auto-fit, minmax(180px, 1fr));
  gap: 8px;
  width: 100%;
}

.gpu-card {
  border: 1px solid #e5e7eb;
  border-radius: 6px;
  padding: 9px 12px;
  cursor: pointer;
  transition: border-color 0.2s, box-shadow 0.2s, background 0.2s;
  background: #fff;

  &:hover {
    border-color: #94a3b8;
    box-shadow: 0 4px 12px rgba(15, 23, 42, 0.06);
  }

  &.active {
    border-color: #2563eb;
    background: #f8fbff;
    box-shadow: inset 0 0 0 1px #2563eb;
  }

  &.is-busy {
    border-color: #f59e0b;
    background: #fffaf0;

    .gpu-usage {
      color: #b45309;
      background: #fffbeb;
      border-color: #fcd34d;
    }

    .gpu-usage-bar span {
      background: #f59e0b;
    }
  }

  &.is-busy.active {
    border-color: #dc2626;
    background: #fff7f7;
    box-shadow: inset 0 0 0 1px #dc2626;

    .gpu-usage {
      color: #dc2626;
      background: #fef2f2;
      border-color: #fecaca;
    }

    .gpu-usage-bar span {
      background: #dc2626;
    }
  }
}

.gpu-card-top {
  display: flex;
  align-items: flex-start;
  justify-content: space-between;
  gap: 10px;
}

.gpu-title {
  min-width: 0;
  display: grid;
  gap: 2px;
}

.gpu-index {
  font-size: 11px;
  line-height: 1;
  font-weight: 600;
  letter-spacing: 0.04em;
  color: #64748b;
}

.gpu-name {
  font-size: 13px;
  font-weight: 600;
  color: #1f2937;
  overflow: hidden;
  text-overflow: ellipsis;
  white-space: nowrap;
}

.gpu-usage {
  height: 20px;
  min-width: 40px;
  padding: 0 6px;
  display: inline-flex;
  align-items: center;
  justify-content: center;
  border: 1px solid #bfdbfe;
  border-radius: 4px;
  background: #eff6ff;
  color: #2563eb;
  font-size: 11px;
  font-weight: 700;
  white-space: nowrap;
}

.gpu-usage-bar {
  height: 3px;
  margin: 7px 0 6px;
  overflow: hidden;
  border-radius: 999px;
  background: #eef2f7;

  span {
    display: block;
    height: 100%;
    border-radius: inherit;
    background: #2563eb;
  }
}

.gpu-meta {
  display: flex;
  flex-wrap: nowrap;
  gap: 10px;
  font-size: 11px;
  color: #64748b;
  line-height: 1.4;
  white-space: nowrap;
  overflow: hidden;
}

.gpu-empty {
  color: #909399;
  font-size: 13px;
}

.command-preview-wrapper {
  margin-left: 80px;
  margin-bottom: 22px;
}

.form-actions-wrapper {
  position: fixed;
  bottom: 0;
  left: 240px;
  right: 0;
  height: 56px;
  background: #fff;
  border-top: 1px solid #eef0f5;
  display: flex;
  align-items: center;
  padding-left: 112px;
  z-index: 1000;
  box-shadow: 0 -4px 12px rgba(0, 0, 0, 0.05);
}

.command-preview {
  background: #1e1e1e;
  color: #d4d4d4;
  padding: 12px 16px;
  border-radius: 6px;
  font-family: 'SFMono-Regular', Consolas, monospace;
  font-size: 12px;
  line-height: 1.6;
  overflow-x: auto;
  width: 100%;
  margin: 0;
}

.model-picker-input {
  cursor: pointer;

  :deep(.el-input__wrapper) {
    cursor: pointer;
  }

  :deep(.el-input__inner) {
    cursor: pointer;
  }
}


/* 训练参数表格样式 */
.hyperparam-section {
  margin: 32px 0 20px;
}
.hyperparam-title {
  font-size: 16px;
  font-weight: 700;
  color: #1f2937;
  margin-bottom: 24px;
}
.param-config-header {
  display: flex;
  justify-content: space-between;
  align-items: center;
  margin-bottom: 12px;
}
.param-config-title {
  font-size: 14px;
  color: #475569;
  font-weight: 600;
}
.param-table-wrapper {
  margin-left: 80px;
  border: 1px solid #eef0f5;
  border-radius: 6px;
  overflow: hidden;
  background: #fff;
}
:deep(.el-form-item) {
  margin-left: 80px;

  .el-form-item__label {
    position: relative;
    padding-left: 0; /* ensure it starts at 0 */
  }

  &.is-required:not(.is-no-asterisk) > .el-form-item__label::before {
    position: absolute;
    left: -10px;
    margin-right: 0;
  }
}
.param-header, .param-row {
  display: grid;
  grid-template-columns: 140px 420px 1fr;
  align-items: center;
}
.param-header {
  background: #f8fafc;
  border-bottom: 1px solid #eef0f5;
  color: #64748b;
  font-size: 13px;
  font-weight: 600;
}
.param-row {
  border-bottom: 1px solid #eef0f5;
}
.param-col {
  padding: 16px 20px;
}
.param-col.name {
  font-weight: 600;
  color: #475569;
  font-size: 14px;
  text-align: left;
}
.param-col.config {
  display: flex;
  align-items: center;
  gap: 12px;
}
.param-col.desc {
  color: #64748b;
  font-size: 13px;
  min-width: 0;
}
.desc-text {
  display: block;
  white-space: nowrap;
  overflow: hidden;
  text-overflow: ellipsis;
  cursor: default;
}
.param-hint {
  color: #94a3b8;
  font-size: 13px;
}
.param-footer {
  padding: 12px 20px;
}

@media (min-width: 1280px) {
  .gpu-list {
    grid-template-columns: repeat(4, minmax(0, 1fr));
  }
}
</style>
