<script setup>
import { ref, onMounted, onBeforeUnmount, computed } from 'vue'
import { useRouter } from 'vue-router'
import { ElMessage } from 'element-plus'
import { Upload, VideoPlay, Refresh, Document } from '@element-plus/icons-vue'
import request from '../utils/request'
import { useUserStore } from '../stores/userStore'
import SidebarMenu from '../components/SidebarMenu.vue'
import BreadcrumbNav from '../components/BreadcrumbNav.vue'
import HudPageHero from '../components/shell/HudPageHero.vue'
import { formatDateTime } from '../utils/dateTimeFormatter'

const router = useRouter()
const { userInfo, getUserInfo } = useUserStore()

const datasets = ref([])
const jobs = ref([])
const baseModels = ref([])
const uploadLoading = ref(false)
const startLoading = ref(false)
const logVisible = ref(false)
const currentJob = ref(null)
let pollTimer = null

const uploadForm = ref({ name: '', file: null })
const trainForm = ref({
  dataset_id: '',
  base_model: 'yolov8n-seg.pt',
  disease_category: '',
  output_model_name: '',
  epochs: 30,
  imgsz: 640,
  batch: 4,
  register_to_library: true,
})

const canTrain = computed(() => {
  const role = String(userInfo.value?.role || '').toUpperCase()
  return role === 'DEVELOPER' || role === 'ADMIN'
})

const statusType = {
  pending: 'info',
  running: 'warning',
  completed: 'success',
  failed: 'danger',
}

const statusLabel = {
  pending: '排队中',
  running: '训练中',
  completed: '已完成',
  failed: '失败',
}

const hasRunning = computed(() =>
  jobs.value.some((j) => j.status === 'pending' || j.status === 'running'),
)

const resumePendingJobs = async () => {
  const pending = jobs.value.filter((j) => j.status === 'pending')
  for (const job of pending) {
    try {
      await request.post(`/training/jobs/${job.job_id}/resume`)
    } catch (e) {
      console.warn('【恢复训练任务】', job.job_id, e?.message || e)
    }
  }
}

const fmt = (v) => (v ? formatDateTime(v) : '—')

const assertTrainingPayload = (res, key) => {
  if (res?.error) {
    throw new Error('训练接口不可用，请重启后端（需包含 /training 路由）')
  }
  if (!Array.isArray(res?.[key])) {
    const msg =
      res?.failure_message ||
      res?.message ||
      res?.operation?.failure_message ||
      `训练接口响应异常（缺少 ${key}）`
    throw new Error(msg)
  }
  return res[key]
}

const loadBaseModels = async () => {
  const res = await request.get('/training/base-models')
  baseModels.value = assertTrainingPayload(res, 'base_models')
}

const loadDatasets = async () => {
  const res = await request.get('/training/datasets')
  datasets.value = assertTrainingPayload(res, 'datasets')
  if (!trainForm.value.dataset_id && datasets.value.length) {
    trainForm.value.dataset_id = datasets.value[0].dataset_id
  }
}

const loadJobs = async () => {
  const res = await request.get('/training/jobs')
  if (Array.isArray(res?.jobs)) {
    jobs.value = res.jobs
    return
  }
  if (res?.error || res?.message) {
    console.warn('【训练任务列表】', res.message || res.error)
  }
  jobs.value = []
}

const refreshAll = async () => {
  await loadJobs()
  await resumePendingJobs()
  await loadJobs()
  await loadDatasets()
  if (hasRunning.value) startPolling()
}

const onFileChange = (file) => {
  uploadForm.value.file = file?.raw || null
}

const submitUpload = async () => {
  if (!uploadForm.value.file) {
    ElMessage.warning('请选择 ZIP 数据集文件')
    return
  }
  const fd = new FormData()
  fd.append('dataset_zip', uploadForm.value.file)
  if (uploadForm.value.name) fd.append('name', uploadForm.value.name)
  uploadLoading.value = true
  try {
    await request.post('/training/datasets/upload', fd, {
      headers: { 'Content-Type': 'multipart/form-data' },
    })
    ElMessage.success('数据集上传成功')
    uploadForm.value.file = null
    uploadForm.value.name = ''
    await loadDatasets()
  } catch (e) {
    ElMessage.error(e?.failure_message || e?.message || '上传失败')
  } finally {
    uploadLoading.value = false
  }
}

const submitTrain = async () => {
  if (!trainForm.value.dataset_id) {
    ElMessage.warning('请先上传并选择数据集')
    return
  }
  if (!trainForm.value.disease_category?.trim()) {
    ElMessage.warning('请填写隐患类别')
    return
  }
  if (!trainForm.value.output_model_name?.trim()) {
    ElMessage.warning('请填写输出模型文件名')
    return
  }
  startLoading.value = true
  try {
    await request.post('/training/jobs', { ...trainForm.value })
    ElMessage.success('训练任务已启动，请在下方的任务列表查看进度')
    await loadJobs()
    await resumePendingJobs()
    await loadJobs()
    startPolling()
  } catch (e) {
    ElMessage.error(e?.failure_message || e?.message || '启动失败')
  } finally {
    startLoading.value = false
  }
}

const showLogs = (row) => {
  currentJob.value = row
  logVisible.value = true
}

const openJobDetail = async (row) => {
  try {
    const res = await request.get(`/training/jobs/${row.job_id}`)
    currentJob.value = res.job
    logVisible.value = true
  } catch {
    currentJob.value = row
    logVisible.value = true
  }
}

const startPolling = () => {
  if (pollTimer) return
  pollTimer = setInterval(async () => {
    await loadJobs()
    if (!hasRunning.value) {
      clearInterval(pollTimer)
      pollTimer = null
    }
  }, 3000)
}

onMounted(() => {
  getUserInfo().then(async () => {
    if (!userInfo.value) {
      router.push('/login')
      return
    }
    if (!canTrain.value) {
      ElMessage.warning('仅管理员或开发人员可使用模型训练')
      router.push('/home')
      return
    }
    try {
      await loadBaseModels()
      await loadJobs()
      await resumePendingJobs()
      await loadDatasets()
      if (hasRunning.value) startPolling()
    } catch (e) {
      const msg = e?.failure_message || e?.message || '请确认后端已启动（http://127.0.0.1:5000）'
      ElMessage.error(`加载训练模块失败：${msg}`)
      console.error('【模型训练】初始化失败', e)
    }
  })
})

onBeforeUnmount(() => {
  if (pollTimer) clearInterval(pollTimer)
})
</script>

<template>
  <div class="iot-page shell-layout-root">
    <div class="main-container">
      <div class="sidebar"><SidebarMenu /></div>
      <div class="content-area">
        <BreadcrumbNav />
        <HudPageHero
          title="模型训练"
          description="基于 Ultralytics YOLOv8 分割模型，上传 YOLO 格式数据集并训练桥梁隐患检测权重，完成后可自动注册到模型库。"
        />

        <el-alert
          type="info"
          :closable="false"
          show-icon
          class="train-tip"
          title="数据集格式说明"
        >
          <p>请上传包含 <code>data.yaml</code> 的 ZIP 包，目录需符合 Ultralytics 规范，例如：</p>
          <pre class="yaml-sample">bridge_hazard_demo/
  data.yaml
  train/images  train/labels
  val/images    val/labels</pre>
          <p>
            完整说明见后端文档
            <code>BridgeDiseaseBackend-main/docs/TRAINING_DATASET.md</code>；
            模板文件 <code>training_templates/data.yaml.example</code>。
          </p>
          <p>
            <strong>尚无标注？</strong> 在后端目录执行
            <code>python scripts/create_demo_dataset.py</code>，
            会生成 <code>training_templates/bridge_hazard_demo.zip</code>（占位框，仅试跑流程）。
          </p>
          <p>训练在后台执行，CPU 建议先试 <code>yolov8n-seg</code>、<code>epochs=5</code>、<code>batch=2</code>。</p>
        </el-alert>

        <el-row :gutter="16">
          <el-col :span="12">
            <el-card class="iot-card">
              <template #header>
                <div class="card-header"><h2>上传训练数据集</h2></div>
              </template>
              <el-form label-width="100px">
                <el-form-item label="数据集名称">
                  <el-input v-model="uploadForm.name" placeholder="可选，便于识别" clearable />
                </el-form-item>
                <el-form-item label="ZIP 文件">
                  <el-upload
                    :auto-upload="false"
                    :limit="1"
                    accept=".zip"
                    :on-change="onFileChange"
                  >
                    <el-button type="primary" plain>
                      <el-icon><Upload /></el-icon> 选择 ZIP
                    </el-button>
                  </el-upload>
                </el-form-item>
                <el-form-item>
                  <el-button type="primary" :loading="uploadLoading" @click="submitUpload">
                    上传并校验
                  </el-button>
                </el-form-item>
              </el-form>
              <el-table :data="datasets" size="small" empty-text="暂无数据集，请先上传">
                <el-table-column prop="name" label="名称" min-width="120" />
                <el-table-column prop="owner_username" label="上传人" width="90" />
                <el-table-column label="上传时间" width="165">
                  <template #default="{ row }">{{ fmt(row.created_at) }}</template>
                </el-table-column>
              </el-table>
            </el-card>
          </el-col>

          <el-col :span="12">
            <el-card class="iot-card">
              <template #header>
                <div class="card-header"><h2>新建训练任务</h2></div>
              </template>
              <el-form label-width="110px">
                <el-form-item label="数据集" required>
                  <el-select v-model="trainForm.dataset_id" placeholder="选择数据集" style="width: 100%">
                    <el-option
                      v-for="d in datasets"
                      :key="d.dataset_id"
                      :label="d.name"
                      :value="d.dataset_id"
                    />
                  </el-select>
                </el-form-item>
                <el-form-item label="基线模型" required>
                  <el-select v-model="trainForm.base_model" style="width: 100%">
                    <el-option
                      v-for="m in baseModels"
                      :key="m.id"
                      :label="m.label"
                      :value="m.id"
                    />
                  </el-select>
                </el-form-item>
                <el-form-item label="隐患类别" required>
                  <el-input v-model="trainForm.disease_category" placeholder="如：钢构件锈蚀" />
                </el-form-item>
                <el-form-item label="输出文件名" required>
                  <el-input
                    v-model="trainForm.output_model_name"
                    placeholder="如：steel_corrosion.pt"
                  />
                </el-form-item>
                <el-form-item label="训练轮次">
                  <el-input-number v-model="trainForm.epochs" :min="1" :max="500" />
                </el-form-item>
                <el-form-item label="图像尺寸">
                  <el-input-number v-model="trainForm.imgsz" :min="320" :max="1280" :step="32" />
                </el-form-item>
                <el-form-item label="Batch">
                  <el-input-number v-model="trainForm.batch" :min="1" :max="32" />
                </el-form-item>
                <el-form-item label="写入模型库">
                  <el-switch v-model="trainForm.register_to_library" />
                </el-form-item>
                <el-form-item>
                  <el-button type="primary" :loading="startLoading" @click="submitTrain">
                    <el-icon><VideoPlay /></el-icon> 开始训练
                  </el-button>
                </el-form-item>
              </el-form>
            </el-card>
          </el-col>
        </el-row>

        <el-card class="iot-card">
          <template #header>
            <div class="card-header">
              <h2>训练任务</h2>
              <el-button @click="refreshAll"><el-icon><Refresh /></el-icon> 刷新</el-button>
            </div>
          </template>
          <el-table :data="jobs">
            <el-table-column prop="job_id" label="任务ID" width="110" show-overflow-tooltip />
            <el-table-column prop="dataset_name" label="数据集" width="120" show-overflow-tooltip />
            <el-table-column prop="disease_category" label="隐患类别" width="130" />
            <el-table-column prop="base_model" label="基线" width="130" />
            <el-table-column prop="output_model_name" label="输出模型" width="150" show-overflow-tooltip />
            <el-table-column label="状态" width="90">
              <template #default="{ row }">
                <el-tag :type="statusType[row.status]" size="small">{{ statusLabel[row.status] || row.status }}</el-tag>
              </template>
            </el-table-column>
            <el-table-column label="进度" width="140">
              <template #default="{ row }">
                <el-progress
                  :percentage="row.progress || 0"
                  :status="row.status === 'failed' ? 'exception' : row.status === 'completed' ? 'success' : undefined"
                />
              </template>
            </el-table-column>
            <el-table-column label="创建时间" width="165">
              <template #default="{ row }">{{ fmt(row.created_at) }}</template>
            </el-table-column>
            <el-table-column label="操作" width="100" fixed="right">
              <template #default="{ row }">
                <el-button link type="primary" @click="openJobDetail(row)">
                  <el-icon><Document /></el-icon> 日志
                </el-button>
              </template>
            </el-table-column>
          </el-table>
        </el-card>
      </div>
    </div>

    <el-dialog v-model="logVisible" title="训练日志" width="640px" destroy-on-close>
      <div v-if="currentJob" class="log-panel">
        <p><strong>任务</strong> {{ currentJob.job_id }} · {{ currentJob.disease_category }}</p>
        <p v-if="currentJob.error_message" class="log-error">{{ currentJob.error_message }}</p>
        <pre class="log-pre">{{ (currentJob.logs || []).join('\n') || '暂无日志' }}</pre>
      </div>
      <template #footer>
        <el-button @click="logVisible = false">关闭</el-button>
      </template>
    </el-dialog>
  </div>
</template>

<style scoped>
@import '../styles/iot-page.css';
.train-tip {
  margin-bottom: 16px;
}
.train-tip p {
  margin: 4px 0;
  font-size: 13px;
  line-height: 1.5;
}
.yaml-sample {
  margin: 8px 0;
  padding: 8px 12px;
  background: var(--card);
  border-radius: 6px;
  font-size: 12px;
}
.log-panel p {
  margin: 0 0 8px;
  font-size: 13px;
}
.log-error {
  color: var(--el-color-danger);
}
.log-pre {
  max-height: 360px;
  overflow: auto;
  padding: 12px;
  background: #1e1e1e;
  color: #d4d4d4;
  border-radius: 8px;
  font-size: 12px;
  line-height: 1.45;
  white-space: pre-wrap;
  word-break: break-all;
}
</style>
