<script setup>
import { Refresh } from '@element-plus/icons-vue'
import { ElMessage, ElMessageBox } from 'element-plus'
import { computed, onBeforeUnmount, onMounted, reactive, ref } from 'vue'
import {
  cancelTrainTask,
  createTrainTask,
  deleteTrainTask,
  getTrainPackages,
  getTrainTasks,
  restartTrainTask,
  saveTrainModel,
} from '@/api/trainManagement'

// ── form ──────────────────────────────────────────────────────────────────────
const form = reactive({
  model_name: '',
  package_id: '',
  params: {
    seq_len: 20,
    hidden_size: 64,
    num_layers: 2,
    epochs: 50,
    batch_size: 32,
    learning_rate: 0.001,
    train_ratio: 0.8,
  },
})

const packages = ref([])
const submitting = ref(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()
    } else {
      stopPolling()
    }
  } catch {
    // ignore poll errors silently
  }
}

const handleStartTraining = async () => {
  if (!form.model_name.trim()) {
    ElMessage.warning('请输入模型名称')
    return
  }
  if (!form.package_id) {
    ElMessage.warning('请选择数据包')
    return
  }

  submitting.value = true
  try {
    await createTrainTask({
      model_name: form.model_name.trim(),
      package_id: form.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 restartTrainTask(task.id)
    ElMessage.success('重新训练已启动')
    await loadTasks()
  } catch (e) {
    ElMessage.error(e?.message || '重启失败')
  }
}

const handleSave = async (task) => {
  try {
    await saveTrainModel(task.id)
    ElMessage.success('模型已保存，可在模型列表中查看')
    await loadTasks()
  } catch (e) {
    ElMessage.error(e?.message || '保存失败')
  }
}

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

// ── 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}`,
    `训练比 ${params.train_ratio}`,
  ].join(' / ')
}

const formatLoss = (task) => {
  if (task.train_loss == null) return '-'
  const train = `训练: ${Number(task.train_loss).toFixed(5)}`
  const val = task.val_loss != null ? ` / 验证: ${Number(task.val_loss).toFixed(5)}` : ''
  return train + val
}

onMounted(async () => {
  loadingTasks.value = true
  try {
    await Promise.all([
      getTrainPackages().then((d) => (packages.value = d)),
      loadTasks(),
    ])
  } finally {
    loadingTasks.value = false
  }
})

onBeforeUnmount(stopPolling)
</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: package + model name -->
        <div class="form-row">
          <el-form-item label="选择数据包" required class="form-item-wide">
            <el-select
              v-model="form.package_id"
              placeholder="请选择数据包"
              filterable
              style="width: 100%"
            >
              <el-option
                v-for="pkg in packages"
                :key="pkg.id"
                :label="`${pkg.name}（${pkg.data_count} 条）`"
                :value="pkg.id"
              />
            </el-select>
          </el-form-item>

          <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: 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>
            <el-form-item label="训练集比例">
              <el-input-number
                v-model="form.params.train_ratio"
                :min="0.5"
                :max="0.99"
                :step="0.05"
                :precision="2"
                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>

    <!-- ── tasks table ──────────────────────────────────────────────────── -->
    <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
        :height="tableHeight"
      >
        <el-table-column prop="model_name" label="模型名称" min-width="140" show-overflow-tooltip />
        <el-table-column prop="package_name" label="数据包" min-width="130" show-overflow-tooltip />

        <el-table-column label="参数" min-width="180" show-overflow-tooltip>
          <template #default="{ row }">
            <el-tooltip :content="formatParams(row.params)" placement="top">
              <span class="params-cell">{{ formatParams(row.params) }}</span>
            </el-tooltip>
          </template>
        </el-table-column>

        <el-table-column label="状态" width="100" 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="150" align="center">
          <template #default="{ row }">
            <template v-if="row.status === 'running'">
              <el-progress
                :percentage="row.progress"
                :stroke-width="6"
                :show-text="true"
                style="width: 120px"
              />
            </template>
            <template v-else-if="row.status === 'completed'">
              <el-progress
                :percentage="100"
                status="success"
                :stroke-width="6"
                style="width: 120px"
              />
            </template>
            <span v-else class="muted">—</span>
          </template>
        </el-table-column>

        <el-table-column label="损失" min-width="200" 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="160" />

        <el-table-column label="操作" width="210" fixed="right">
          <template #default="{ row }">
            <!-- running / pending -->
            <template v-if="row.status === 'running' || row.status === 'pending'">
              <el-button link type="warning" @click="handleCancel(row)">取消</el-button>
            </template>

            <!-- completed -->
            <template v-else-if="row.status === 'completed'">
              <el-button
                v-if="!row.is_saved"
                link
                type="primary"
                @click="handleSave(row)"
              >
                保存模型
              </el-button>
              <el-tag v-else type="success" size="small" style="margin-right: 6px">已保存</el-tag>
              <el-button link type="info" @click="handleRestart(row)">重新训练</el-button>
              <el-button link type="danger" @click="handleDelete(row)">删除</el-button>
            </template>

            <!-- failed / cancelled -->
            <template v-else>
              <el-button link type="info" @click="handleRestart(row)">重新训练</el-button>
              <el-button link type="danger" @click="handleDelete(row)">删除</el-button>
            </template>
          </template>
        </el-table-column>
      </el-table>
    </el-card>
  </div>
</template>

<script>
// tableHeight is a non-reactive calculation; compute once
const tableHeight = 'calc(100vh - 520px)'
export default {}
</script>

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

.config-card,
.tasks-card {
  flex-shrink: 0;
  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;
  }
}

.tasks-card {
  flex: 1;
  min-height: 0;
  overflow: hidden;
}

.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;
    }
  }
}

.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;
}

.params-cell {
  font-size: 12px;
  color: var(--text-secondary);
  white-space: nowrap;
  overflow: hidden;
  text-overflow: ellipsis;
  display: block;
}

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

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