YG_FT/backend/app/modules/data_convert/router.py

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from fastapi import APIRouter, Body, Depends, File, UploadFile
from fastapi.responses import FileResponse

from app.api.v1.endpoints.platform import ok, fail
from app.core.auth import get_current_user
from app.db.platform_store import get_platform_store, new_id


router = APIRouter(prefix="/data-convert", tags=["data-convert"])

# 存储根目录
STORAGE_ROOT = Path(__file__).resolve().parents[3] / "storage" / "data-convert"


def _safe_output_filename(value: Any) -> str:
    """输出文件名白名单校验：仅允许普通文件名，阻断 ``../``、``/``、``\\`` 等路径穿越。

    转换结果始终写入 ``STORAGE_ROOT/<task_id>/output/<output_filename>``，
    若文件名可被注入路径分隔符，将导致任意文件读写/删除。
    """
    name = str(value or "converted-data.jsonl").strip()
    if (
        not name
        or name in {".", ".."}
        or name != Path(name).name
        or "/" in name
        or "\\" in name
        or any(ord(character) < 32 or ord(character) == 127 for character in name)
    ):
        raise fail(400, "output filename must be a plain file name")
    return name


def _task_output_path(task: dict[str, Any]) -> Path:
    """返回经过白名单校验的转换输出文件路径（始终位于任务 output 目录内）。"""
    return _output_dir(task["id"]) / _safe_output_filename(task.get("output_filename"))


def _task_dir(task_id: str) -> Path:
    return STORAGE_ROOT / task_id


def _input_dir(task_id: str) -> Path:
    return _task_dir(task_id) / "input"


def _output_dir(task_id: str) -> Path:
    return _task_dir(task_id) / "output"


@router.get("")
def list_tasks(
    page: int = 1,
    page_size: int = 20,
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    store = get_platform_store()
    with store.connect() as conn:
        rows = conn.execute(
            "SELECT * FROM data_convert_tasks WHERE deleted_at IS NULL "
            "ORDER BY create_time DESC LIMIT %s OFFSET %s",
            (page_size, (page - 1) * page_size),
        ).fetchall()
        total = conn.execute(
            "SELECT COUNT(*) FROM data_convert_tasks WHERE deleted_at IS NULL"
        ).fetchone()[0]
    return ok({"items": [dict(r) for r in rows], "total": total})


@router.post("")
def create_task(
    payload: dict[str, Any] = Body(...),
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    name = str(payload.get("name") or "").strip()
    if not name:
        raise fail(400, "name is required")
    task_id = new_id("dct")
    output_filename = _safe_output_filename(payload.get("output_filename"))
    description = str(payload.get("description") or "").strip()
    store = get_platform_store()
    with store.connect() as conn:
        conn.execute(
            "INSERT INTO data_convert_tasks (id, name, description, output_filename) "
            "VALUES (%s, %s, %s, %s)",
            (task_id, name, description, output_filename),
        )
    # 创建目录
    _input_dir(task_id).mkdir(parents=True, exist_ok=True)
    _output_dir(task_id).mkdir(parents=True, exist_ok=True)
    return ok(_get_task(task_id))


@router.get("/{task_id}")
def get_task(
    task_id: str,
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    task = _get_task(task_id)
    if not task:
        raise fail(404, "task not found")
    # 附加输入文件列表
    input_dir = _input_dir(task_id)
    files = []
    if input_dir.exists():
        for f in sorted(input_dir.iterdir()):
            if f.is_file():
                files.append({"name": f.name, "size": f.stat().st_size})
    task["input_files"] = files
    return ok(task)


@router.post("/{task_id}/source-files")
async def upload_source_files(
    task_id: str,
    files: list[UploadFile] = File(...),
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    task = _get_task(task_id)
    if not task:
        raise fail(404, "task not found")
    if task["status"] not in ("pending", "uploaded"):
        raise fail(400, "task is not editable")
    input_dir = _input_dir(task_id)
    input_dir.mkdir(parents=True, exist_ok=True)
    staged = []
    for upload in files:
        name = Path(upload.filename or "input.json").name
        if not name.lower().endswith(".json"):
            raise fail(415, f"only JSON files are supported: {name}")
        target = input_dir / name
        content = await upload.read()
        target.write_bytes(content)
        staged.append({"name": name, "size": len(content)})
    store = get_platform_store()
    # 标记上传完成
    with store.connect() as conn:
        conn.execute(
            "UPDATE data_convert_tasks SET status='uploaded', update_time=NOW() WHERE id=%s",
            (task_id,),
        )
    # 自动转换并导入数据集
    try:
        output_dir = _output_dir(task_id)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = _task_output_path(task)
        # 清空旧输出（如果重新上传）
        if output_path.exists():
            output_path.unlink()
        input_count = 0
        output_count = 0
        for json_file in sorted(input_dir.iterdir()):
            if not json_file.is_file() or not json_file.name.lower().endswith(".json"):
                continue
            input_count += 1
            with open(json_file, "r", encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, list):
                records = data
            elif isinstance(data, dict):
                records = [data]
            else:
                raise ValueError(f"JSON must be object or array: {json_file.name}")
            with open(output_path, "a", encoding="utf-8") as f:
                for record in records:
                    f.write(json.dumps(record, ensure_ascii=False) + "\n")
                    output_count += 1
        with store.connect() as conn:
            conn.execute(
                "UPDATE data_convert_tasks SET status='completed', "
                "input_count=%s, output_count=%s, update_time=NOW() WHERE id=%s",
                (input_count, output_count, task_id),
            )
        # 自动导入数据集
        content = output_path.read_text(encoding="utf-8")
        size_bytes = len(content.encode("utf-8"))
        dataset = store.create_dataset({
            "name": task["name"],
            "type": "train",
            "storage_type": "local",
            "source": "upload",
            "task_id": task_id,
            "size": f"{size_bytes} B",
            "count": output_count,
            "description": f"由数据类型转换任务 {task_id} 自动导入",
        })
        dataset_id = dataset["id"]
        with store.connect() as conn:
            store.add_dataset_file(conn, dataset_id, _safe_output_filename(task.get("output_filename")), content)
        return ok({
            "staged_files": staged,
            "auto_converted": True,
            "dataset_id": dataset_id,
            "input_count": input_count,
            "output_count": output_count,
        })
    except Exception as exc:
        with store.connect() as conn:
            conn.execute(
                "UPDATE data_convert_tasks SET status='failed', error_message=%s, update_time=NOW() WHERE id=%s",
                (str(exc)[:500], task_id),
            )
        return ok({"staged_files": staged, "auto_converted": False, "error": str(exc)[:500]})


@router.post("/{task_id}/run")
def run_convert(
    task_id: str,
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    task = _get_task(task_id)
    if not task:
        raise fail(404, "task not found")
    if task["status"] not in ("uploaded", "completed", "failed"):
        raise fail(400, "please upload source files first")
    # 标记运行中
    store = get_platform_store()
    with store.connect() as conn:
        conn.execute(
            "UPDATE data_convert_tasks SET status='running', error_message='', update_time=NOW() WHERE id=%s",
            (task_id,),
        )
    try:
        input_dir = _input_dir(task_id)
        output_dir = _output_dir(task_id)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = _task_output_path(task)
        input_count = 0
        output_count = 0
        for json_file in sorted(input_dir.iterdir()):
            if not json_file.is_file() or not json_file.name.lower().endswith(".json"):
                continue
            input_count += 1
            with open(json_file, "r", encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, list):
                records = data
            elif isinstance(data, dict):
                records = [data]
            else:
                raise ValueError(f"JSON must be object or array: {json_file.name}")
            with open(output_path, "a", encoding="utf-8") as f:
                for record in records:
                    f.write(json.dumps(record, ensure_ascii=False) + "\n")
                    output_count += 1
        # 更新任务状态
        with store.connect() as conn:
            conn.execute(
                "UPDATE data_convert_tasks SET status='completed', "
                "input_count=%s, output_count=%s, update_time=NOW() WHERE id=%s",
                (input_count, output_count, task_id),
            )
    except Exception as exc:
        with store.connect() as conn:
            conn.execute(
                "UPDATE data_convert_tasks SET status='failed', error_message=%s, update_time=NOW() WHERE id=%s",
                (str(exc)[:500], task_id),
            )
        raise fail(500, f"convert failed: {exc}")
    return ok(_get_task(task_id))


@router.get("/{task_id}/download")
def download_result(
    task_id: str,
    current_user: dict = Depends(get_current_user),
):
    task = _get_task(task_id)
    if not task:
        raise fail(404, "task not found")
    if task["status"] != "completed":
        raise fail(400, "task is not completed")
    output_path = _task_output_path(task)
    if not output_path.exists():
        raise fail(404, "output file not found")
    return FileResponse(
        str(output_path),
        media_type="application/octet-stream",
        filename=_safe_output_filename(task.get("output_filename")),
    )


@router.post("/{task_id}/import-as-dataset")
def import_as_dataset(
    task_id: str,
    payload: dict[str, Any] = Body(default={}),
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    """把已转换的 JSONL 文件导入为数据集管理中的上传任务记录（source='task'）。"""
    task = _get_task(task_id)
    if not task:
        raise fail(404, "task not found")
    if task["status"] != "completed":
        raise fail(400, "task is not completed")
    output_path = _task_output_path(task)
    if not output_path.exists():
        raise fail(404, "output file not found")
    content = output_path.read_text(encoding="utf-8")
    dataset_name = str(payload.get("name") or task["name"]).strip()
    description = str(payload.get("description") or f"由数据类型转换任务 {task_id} 导入").strip()
    size_bytes = len(content.encode("utf-8"))
    store = get_platform_store()
    # 用 store 提供的接口创建数据集与文件
    dataset = store.create_dataset({
        "name": dataset_name,
        "type": "train",
        "storage_type": "local",
        "source": "upload",
        "task_id": task_id,
        "size": f"{size_bytes} B",
        "count": task["output_count"],
        "description": description,
    })
    dataset_id = dataset["id"]
    with store.connect() as conn:
        store.add_dataset_file(conn, dataset_id, _safe_output_filename(task.get("output_filename")), content)
    return ok({"dataset_id": dataset_id, "name": dataset_name})


@router.delete("/{task_id}")
def delete_task(
    task_id: str,
    current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
    task = _get_task(task_id)
    if not task:
        raise fail(404, "task not found")
    store = get_platform_store()
    with store.connect() as conn:
        conn.execute(
            "UPDATE data_convert_tasks SET deleted_at=NOW() WHERE id=%s",
            (task_id,),
        )
    # 清理文件
    import shutil
    task_dir = _task_dir(task_id)
    if task_dir.exists():
        shutil.rmtree(task_dir, ignore_errors=True)
    return ok({"deleted": task_id})


def _get_task(task_id: str) -> dict[str, Any] | None:
    store = get_platform_store()
    with store.connect() as conn:
        row = conn.execute(
            "SELECT * FROM data_convert_tasks WHERE id=%s AND deleted_at IS NULL",
            (task_id,),
        ).fetchone()
    return dict(row) if row else None