<script setup>
import * as echarts from 'echarts'
import { Refresh, Search } from '@element-plus/icons-vue'
import { ElMessage, ElMessageBox } from 'element-plus'
import { computed, nextTick, onBeforeUnmount, onMounted, reactive, ref, watch } from 'vue'
import {
  cancelTrainTask,
  createTrainTask,
  deleteTrainTask,
  getTrainTasks,
  restartTrainTask,
  saveTrainModel,
} from '@/api/trainManagement'
import { getPackages, getPkgCategoryTree } from '@/api/packageManagement'

// ── form ──────────────────────────────────────────────────────────────────────
const form = reactive({
  model_name: '',
  train_package_id: null,
  train_package_name: '',
  val_package_id: null,
  val_package_name: '',
  params: {
    seq_len: 20,
    hidden_size: 64,
    num_layers: 2,
    epochs: 50,
    batch_size: 32,
    learning_rate: 0.001,
  },
})

const submitting = ref(false)

// ── package picker modal ──────────────────────────────────────────────────────
const pickerVisible = ref(false)
const pickerMode = ref('train')  // 'train' | 'val'
const pickerSearch = ref('')
const pickerCategoryId = ref('')
const pickerCategories = ref([])
const pickerPackages = ref([])
const pickerLoading = ref(false)

const openPicker = async (mode) => {
  pickerMode.value = mode
  pickerSearch.value = ''
  pickerCategoryId.value = ''
  pickerVisible.value = true
  if (!pickerCategories.value.length) {
    try {
      const tree = await getPkgCategoryTree()
      const flatten = (nodes, result = []) => {
        for (const n of nodes) {
          result.push({ value: String(n.id), label: n.name })
          if (n.children?.length) flatten(n.children, result)
        }
        return result
      }
      pickerCategories.value = flatten(tree)
    } catch {}
  }
  await loadPickerPackages()
}

const loadPickerPackages = async () => {
  pickerLoading.value = true
  try {
    const data = await getPackages({
      category_id: pickerCategoryId.value || '',
      name: pickerSearch.value.trim(),
    })
    pickerPackages.value = data
  } finally {
    pickerLoading.value = false
  }
}

const selectPackage = (pkg) => {
  if (pickerMode.value === 'train') {
    form.train_package_id = pkg.id
    form.train_package_name = pkg.name
  } else {
    form.val_package_id = pkg.id
    form.val_package_name = pkg.name
  }
  pickerVisible.value = false
}

// ── tasks table ───────────────────────────────────────────────────────────────
const tasks = ref([])
const loadingTasks = ref(false)
let pollTimer = null

const hasActiveTasks = computed(() =>
  tasks.value.some((t) => t.status === 'pending' || t.status === 'running'),
)

const startPolling = () => {
  if (pollTimer) return
  pollTimer = setInterval(loadTasks, 3000)
}

const stopPolling = () => {
  if (pollTimer) {
    clearInterval(pollTimer)
    pollTimer = null
  }
}

const loadTasks = async () => {
  try {
    tasks.value = await getTrainTasks()
    if (hasActiveTasks.value) {
      startPolling()
      // refresh selected task detail if still active
      if (selectedTask.value && ['pending', 'running'].includes(selectedTask.value.status)) {
        const fresh = tasks.value.find((t) => t.id === selectedTask.value.id)
        if (fresh) selectedTask.value = fresh
      }
    } else {
      stopPolling()
    }
  } catch {
    // ignore poll errors silently
  }
}

const handleStartTraining = async () => {
  if (!form.model_name.trim()) return ElMessage.warning('请输入模型名称')
  if (!form.train_package_id) return ElMessage.warning('请选择训练集数据包')
  if (!form.val_package_id) return ElMessage.warning('请选择验证集数据包')

  submitting.value = true
  try {
    await createTrainTask({
      model_name: form.model_name.trim(),
      train_package_id: form.train_package_id,
      val_package_id: form.val_package_id,
      params: { ...form.params },
    })
    ElMessage.success('训练任务已提交')
    form.model_name = ''
    await loadTasks()
  } finally {
    submitting.value = false
  }
}

const handleCancel = async (task) => {
  try {
    await cancelTrainTask(task.id)
    ElMessage.success('已发送取消请求')
    await loadTasks()
  } catch (e) {
    ElMessage.error(e?.message || '取消失败')
  }
}

const handleRestart = async (task) => {
  try {
    await ElMessageBox.confirm(
      `确定重新训练"${task.model_name}"？将使用相同配置和数据包新建一个训练任务。`,
      '确认重新训练',
      { type: 'warning', confirmButtonText: '确认', cancelButtonText: '取消' },
    )
    await restartTrainTask(task.id)
    ElMessage.success('重新训练已启动')
    await loadTasks()
  } catch {
    // user cancelled
  }
}

// ── save model modal ──────────────────────────────────────────────────────────
const saveModalVisible = ref(false)
const saveModalTask = ref(null)
const saveForm = reactive({ model_name: '', description: '' })
const saving = ref(false)

const handleSave = (task) => {
  saveModalTask.value = task
  saveForm.model_name = task.model_name
  saveForm.description = ''
  saveModalVisible.value = true
}

const confirmSave = async () => {
  if (!saveForm.model_name.trim()) return ElMessage.warning('请输入模型名称')
  saving.value = true
  try {
    await saveTrainModel(saveModalTask.value.id, {
      model_name: saveForm.model_name.trim(),
      description: saveForm.description.trim(),
    })
    ElMessage.success('模型已保存，可在模型列表中查看')
    saveModalVisible.value = false
    await loadTasks()
  } catch (e) {
    ElMessage.error(e?.message || '保存失败')
  } finally {
    saving.value = false
  }
}

const handleDelete = async (task) => {
  try {
    await ElMessageBox.confirm(`确定删除训练任务"${task.model_name}"吗？`, '提示', {
      type: 'warning',
    })
    await deleteTrainTask(task.id)
    ElMessage.success('已删除')
    if (selectedTask.value?.id === task.id) selectedTask.value = null
    await loadTasks()
  } catch {
    // user cancelled
  }
}

// ── task detail (epoch logs) ──────────────────────────────────────────────────
const selectedTask = ref(null)
const detailMode = ref('table')  // 'table' | 'chart'
const chartRef = ref(null)
let chartInstance = null

const epochLogs = computed(() => selectedTask.value?.epoch_logs || [])

const handleRowClick = (row) => {
  selectedTask.value = row
  detailMode.value = 'table'
}

const renderEpochChart = () => {
  if (!chartRef.value || detailMode.value !== 'chart') return
  if (!chartInstance) chartInstance = echarts.init(chartRef.value)
  const logs = epochLogs.value
  const epochs = logs.map((l) => l.epoch)
  const trainLoss = logs.map((l) => l.train_loss)
  const valLoss = logs.map((l) => l.val_loss)
  chartInstance.setOption(
    {
      animation: false,
      color: ['#409EFF', '#F56C6C'],
      tooltip: {
        trigger: 'axis',
        formatter(params) {
          const e = params[0]?.axisValue
          const lines = [`<b>Epoch ${e}</b>`]
          params.forEach((p) => {
            if (p.data != null)
              lines.push(
                `<span style="color:${p.color}">● </span>${p.seriesName}: ${Number(p.data).toExponential(4)}`,
              )
          })
          return lines.join('<br/>')
        },
      },
      legend: { data: ['训练损失', '验证损失'], top: 4 },
      grid: { left: 70, right: 20, top: 36, bottom: 36 },
      xAxis: { type: 'category', data: epochs, name: 'Epoch', nameLocation: 'end' },
      yAxis: { type: 'value', name: 'Loss', scale: true },
      series: [
        { name: '训练损失', type: 'line', data: trainLoss, smooth: true, symbol: 'none' },
        { name: '验证损失', type: 'line', data: valLoss, smooth: true, symbol: 'none' },
      ],
    },
    true,
  )
}

watch(detailMode, (val) => {
  if (val === 'chart') nextTick(() => { renderEpochChart(); chartInstance?.resize() })
})

watch(epochLogs, () => {
  if (detailMode.value === 'chart') nextTick(renderEpochChart)
})

// ── display helpers ───────────────────────────────────────────────────────────
const STATUS_MAP = {
  pending: { type: 'info', text: '等待中' },
  running: { type: 'warning', text: '训练中' },
  completed: { type: 'success', text: '已完成' },
  failed: { type: 'danger', text: '失败' },
  cancelled: { type: '', text: '已取消' },
}

const getStatusTag = (status) => STATUS_MAP[status] || { type: 'info', text: status }

const formatParams = (params) => {
  if (!params) return '-'
  return [
    `序列长度 ${params.seq_len}`,
    `隐藏层 ${params.hidden_size}`,
    `层数 ${params.num_layers}`,
    `轮数 ${params.epochs}`,
    `批次 ${params.batch_size}`,
    `学习率 ${params.learning_rate}`,
  ].join(' / ')
}

const formatLoss = (task) => {
  if (task.train_loss == null) return '-'
  const parts = [`训练: ${Number(task.train_loss).toFixed(5)}`]
  if (task.val_loss != null) parts.push(`验证: ${Number(task.val_loss).toFixed(5)}`)
  return parts.join(' / ')
}

onMounted(async () => {
  loadingTasks.value = true
  try {
    await loadTasks()
  } finally {
    loadingTasks.value = false
  }
})

onBeforeUnmount(() => {
  stopPolling()
  chartInstance?.dispose()
})
</script>

<template>
  <div class="train-page">
    <!-- ── config card ──────────────────────────────────────────────────── -->
    <el-card class="config-card" shadow="hover">
      <template #header>
        <span class="card-title">训练配置</span>
      </template>

      <el-form :model="form" label-position="top" class="train-form">
        <!-- row 1: model name -->
        <div class="form-row">
          <el-form-item label="模型名称" required class="form-item-wide">
            <el-input
              v-model="form.model_name"
              placeholder="请输入模型名称"
              maxlength="100"
              show-word-limit
            />
          </el-form-item>
        </div>

        <!-- row 2: package pickers -->
        <div class="form-row">
          <el-form-item label="训练集数据包" required class="form-item-wide">
            <div class="pkg-picker-row">
              <el-input
                :model-value="form.train_package_name || ''"
                placeholder="点击右侧按钮选择训练集"
                readonly
                class="pkg-input"
              />
              <el-button @click="openPicker('train')">选择</el-button>
            </div>
          </el-form-item>
          <el-form-item label="验证集数据包" required class="form-item-wide">
            <div class="pkg-picker-row">
              <el-input
                :model-value="form.val_package_name || ''"
                placeholder="点击右侧按钮选择验证集"
                readonly
                class="pkg-input"
              />
              <el-button @click="openPicker('val')">选择</el-button>
            </div>
          </el-form-item>
        </div>

        <!-- row 3: LSTM params -->
        <div class="params-section">
          <span class="params-label">LSTM 超参数</span>
          <div class="params-grid">
            <el-form-item label="序列长度">
              <el-input-number v-model="form.params.seq_len" :min="5" :max="500" :step="5" controls-position="right" style="width:100%" />
            </el-form-item>
            <el-form-item label="隐藏层大小">
              <el-input-number v-model="form.params.hidden_size" :min="8" :max="1024" :step="8" controls-position="right" style="width:100%" />
            </el-form-item>
            <el-form-item label="LSTM 层数">
              <el-input-number v-model="form.params.num_layers" :min="1" :max="8" controls-position="right" style="width:100%" />
            </el-form-item>
            <el-form-item label="训练轮数 (Epochs)">
              <el-input-number v-model="form.params.epochs" :min="1" :max="2000" :step="10" controls-position="right" style="width:100%" />
            </el-form-item>
            <el-form-item label="批次大小 (Batch)">
              <el-input-number v-model="form.params.batch_size" :min="1" :max="512" :step="8" controls-position="right" style="width:100%" />
            </el-form-item>
            <el-form-item label="学习率">
              <el-input-number v-model="form.params.learning_rate" :min="0.00001" :max="1" :step="0.0001" :precision="5" controls-position="right" style="width:100%" />
            </el-form-item>
          </div>
        </div>

        <!-- action -->
        <div class="form-action">
          <el-button type="primary" size="large" :loading="submitting" @click="handleStartTraining">
            开始训练
          </el-button>
        </div>
      </el-form>
    </el-card>

    <!-- ── bottom: tasks list + epoch detail ────────────────────────────── -->
    <div class="bottom-area">
      <!-- tasks card -->
      <el-card class="tasks-card" shadow="hover">
        <template #header>
          <div class="card-header-row">
            <span class="card-title">训练记录</span>
            <el-button :icon="Refresh" size="small" plain :loading="loadingTasks" @click="loadTasks">刷新</el-button>
          </div>
        </template>

        <el-table
          :data="tasks"
          v-loading="loadingTasks"
          border
          stripe
          highlight-current-row
          height="100%"
          style="cursor:pointer"
          @row-click="handleRowClick"
        >
          <el-table-column prop="model_name" label="模型名称" min-width="130" show-overflow-tooltip />
          <el-table-column label="训练集" min-width="120" show-overflow-tooltip>
            <template #default="{ row }">{{ row.package_name }}</template>
          </el-table-column>
          <el-table-column label="验证集" min-width="120" show-overflow-tooltip>
            <template #default="{ row }">{{ row.val_package_name || '-' }}</template>
          </el-table-column>
          <el-table-column label="状态" width="90" align="center">
            <template #default="{ row }">
              <el-tag :type="getStatusTag(row.status).type" size="small">
                {{ getStatusTag(row.status).text }}
              </el-tag>
            </template>
          </el-table-column>
          <el-table-column label="进度" width="140" align="center">
            <template #default="{ row }">
              <el-progress v-if="row.status === 'running'" :percentage="row.progress" :stroke-width="6" style="width:110px" />
              <el-progress v-else-if="row.status === 'completed'" :percentage="100" status="success" :stroke-width="6" style="width:110px" />
              <span v-else class="muted">—</span>
            </template>
          </el-table-column>
          <el-table-column label="损失" min-width="180" show-overflow-tooltip>
            <template #default="{ row }">
              <span v-if="row.status === 'failed'" class="err-text">{{ row.error_msg || '未知错误' }}</span>
              <span v-else>{{ formatLoss(row) }}</span>
            </template>
          </el-table-column>
          <el-table-column prop="created_at" label="创建时间" width="155" />
          <el-table-column label="操作" width="190" fixed="right">
            <template #default="{ row }">
              <template v-if="row.status === 'running' || row.status === 'pending'">
                <el-button link type="warning" @click.stop="handleCancel(row)">取消</el-button>
              </template>
              <template v-else-if="row.status === 'completed'">
                <el-button v-if="!row.is_saved" link type="primary" @click.stop="handleSave(row)">保存模型</el-button>
                <el-tag v-else type="success" size="small" style="margin-right:6px">已保存</el-tag>
                <el-button link type="info" @click.stop="handleRestart(row)">重新训练</el-button>
                <el-button link type="danger" @click.stop="handleDelete(row)">删除</el-button>
              </template>
              <template v-else>
                <el-button link type="info" @click.stop="handleRestart(row)">重新训练</el-button>
                <el-button link type="danger" @click.stop="handleDelete(row)">删除</el-button>
              </template>
            </template>
          </el-table-column>
        </el-table>
      </el-card>

      <!-- epoch detail panel -->
      <el-card v-if="selectedTask" class="detail-card" shadow="hover">
        <template #header>
          <div class="card-header-row">
            <span class="card-title">训练过程 — {{ selectedTask.model_name }}</span>
            <el-radio-group v-model="detailMode" size="small">
              <el-radio-button value="table">表格</el-radio-button>
              <el-radio-button value="chart">曲线</el-radio-button>
            </el-radio-group>
          </div>
        </template>

        <div class="detail-body">
          <!-- table mode -->
          <el-table
            v-if="detailMode === 'table'"
            :data="epochLogs"
            border
            stripe
            height="100%"
            size="small"
          >
            <el-table-column prop="epoch" label="轮次" width="70" align="center" />
            <el-table-column label="训练损失" align="right">
              <template #default="{ row }">
                {{ row.train_loss != null ? Number(row.train_loss).toExponential(4) : '-' }}
              </template>
            </el-table-column>
            <el-table-column label="验证损失" align="right">
              <template #default="{ row }">
                {{ row.val_loss != null ? Number(row.val_loss).toExponential(4) : '-' }}
              </template>
            </el-table-column>
          </el-table>

          <!-- chart mode -->
          <div v-show="detailMode === 'chart'" ref="chartRef" class="epoch-chart" />
        </div>
      </el-card>
    </div>
  </div>

  <!-- ── package picker dialog ────────────────────────────────────────── -->
  <el-dialog
    v-model="pickerVisible"
    :title="pickerMode === 'train' ? '选择训练集数据包' : '选择验证集数据包'"
    width="640px"
    :close-on-click-modal="false"
  >
    <div class="picker-filters">
      <el-select
        v-model="pickerCategoryId"
        placeholder="所有分类"
        clearable
        style="width:180px"
        @change="loadPickerPackages"
      >
        <el-option
          v-for="opt in pickerCategories"
          :key="opt.value"
          :label="opt.label"
          :value="opt.value"
        />
      </el-select>
      <el-input
        v-model="pickerSearch"
        placeholder="数据包名称"
        :prefix-icon="Search"
        clearable
        style="flex:1"
        @keyup.enter="loadPickerPackages"
        @clear="loadPickerPackages"
      />
      <el-button type="primary" @click="loadPickerPackages">搜索</el-button>
    </div>

    <el-table
      :data="pickerPackages"
      v-loading="pickerLoading"
      border
      stripe
      highlight-current-row
      height="360px"
      style="cursor:pointer;margin-top:10px"
      @row-click="selectPackage"
    >
      <el-table-column prop="name" label="数据包名称" show-overflow-tooltip />
      <el-table-column prop="data_count" label="数据量" width="90" align="center" />
      <el-table-column prop="created_at" label="创建时间" width="160" />
    </el-table>
  </el-dialog>

  <!-- ── save model dialog ────────────────────────────────────────── -->
  <el-dialog
    v-model="saveModalVisible"
    title="保存模型"
    width="480px"
    :close-on-click-modal="false"
  >
    <el-form :model="saveForm" label-position="top">
      <el-form-item label="模型名称" required>
        <el-input v-model="saveForm.model_name" placeholder="请输入模型名称" maxlength="100" show-word-limit />
      </el-form-item>
      <el-form-item label="说明">
        <el-input
          v-model="saveForm.description"
          type="textarea"
          :rows="3"
          placeholder="可选，填写模型用途、训练说明等"
          maxlength="500"
          show-word-limit
        />
      </el-form-item>
    </el-form>
    <template #footer>
      <el-button @click="saveModalVisible = false">取消</el-button>
      <el-button type="primary" :loading="saving" @click="confirmSave">保存</el-button>
    </template>
  </el-dialog>
</template>

<style lang="scss" scoped>
.train-page {
  height: 100%;
  display: flex;
  flex-direction: column;
  gap: 12px;
  overflow: hidden;
  background: var(--bg-page);
}

.config-card,
.tasks-card,
.detail-card {
  border: 1px solid var(--border-color);
  box-shadow: var(--shadow-card) !important;
  border-radius: 4px;

  :deep(.el-card__header) {
    padding: 0 20px;
    height: 44px;
    display: flex;
    align-items: center;
    border-bottom: 1px solid var(--border-color);
    background: #F7F8FA;
  }

  :deep(.el-card__body) {
    padding: 16px 20px;
    height: calc(100% - 44px);
    overflow: hidden;
  }
}

.config-card {
  flex-shrink: 0;
}

.bottom-area {
  flex: 1;
  min-height: 0;
  display: flex;
  gap: 12px;
}

.tasks-card {
  flex: 1.4;
  min-width: 0;
  overflow: hidden;
}

.detail-card {
  flex: 1;
  min-width: 320px;
  overflow: hidden;
}

.detail-body {
  height: 100%;
  display: flex;
  flex-direction: column;
}

.epoch-chart {
  flex: 1;
  min-height: 0;
}

.card-title {
  font-size: 14px;
  font-weight: 600;
  color: var(--text-primary);
}

.card-header-row {
  display: flex;
  align-items: center;
  justify-content: space-between;
  width: 100%;
}

.train-form {
  .form-row {
    display: flex;
    gap: 24px;
    flex-wrap: wrap;
    margin-bottom: 4px;
  }

  .form-item-wide {
    flex: 1;
    min-width: 220px;

    :deep(.el-form-item__content) {
      display: block;
    }
  }
}

.pkg-picker-row {
  display: flex;
  gap: 8px;
  align-items: center;

  .pkg-input {
    flex: 1;
  }
}

.params-section {
  border-top: 1px solid var(--border-color);
  padding-top: 12px;
  margin-bottom: 8px;
}

.params-label {
  display: block;
  font-size: 13px;
  font-weight: 600;
  color: var(--text-secondary);
  margin-bottom: 12px;
}

.params-grid {
  display: grid;
  grid-template-columns: repeat(auto-fill, minmax(160px, 1fr));
  gap: 12px 24px;

  :deep(.el-form-item) {
    margin-bottom: 0;
  }
}

.form-action {
  border-top: 1px solid var(--border-color);
  padding-top: 16px;
  text-align: right;
}

.picker-filters {
  display: flex;
  gap: 8px;
  align-items: center;
}

.muted {
  color: var(--text-tertiary);
}

.err-text {
  color: #F53F3F;
  font-size: 12px;
  white-space: nowrap;
  overflow: hidden;
  text-overflow: ellipsis;
  display: block;
}
</style>
