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

from __future__ import annotations

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

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

from app.api.v1.endpoints.platform import ok, fail
from app.core.auth import get_current_user, is_admin
from app.core.config import get_settings
from app.core.op_log import op_log, OpModule, OpAction
from app.db.platform_store import get_platform_store, new_id
from app.modules.storage.minio_store import get_object_storage
from app.modules.storage.policy import should_store_in_minio


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"


def _minio_enabled() -> bool:
    return bool(get_settings().minio_enabled)


def _input_object_key(task_id: str, name: str) -> str:
    return f"data-convert/{task_id}/input/{Path(name).name}"


def _output_object_key(task_id: str, name: str) -> str:
    return f"data-convert/{task_id}/output/{Path(name).name}"


def _task_objects(task_id: str) -> list[dict[str, Any]]:
    return get_platform_store().storage_objects_for_resource("data_convert", task_id)


def _register_object(
    task_id: str,
    *,
    version_id: str,
    object_key: str,
    file_name: str,
    content_type: str,
    content: bytes,
    created_by: str | None,
) -> dict[str, Any]:
    storage = get_object_storage()
    uploaded = storage.put_bytes(object_key, content, content_type)
    return get_platform_store().create_storage_object(
        {
            "resource_type": "data_convert",
            "resource_id": task_id,
            "version_id": version_id,
            "bucket": uploaded["bucket"],
            "object_key": object_key,
            "file_name": file_name,
            "content_type": content_type,
            "byte_size": len(content),
            "checksum_sha256": hashlib.sha256(content).hexdigest(),
            "status": "available",
            "created_by": created_by,
        }
    )


def _input_objects(task_id: str) -> list[dict[str, Any]]:
    prefix = f"data-convert/{task_id}/input/"
    return sorted(
        [item for item in _task_objects(task_id) if str(item.get("object_key") or "").startswith(prefix)],
        key=lambda item: str(item.get("file_name") or item.get("object_key") or ""),
    )


def _output_object(task: dict[str, Any]) -> dict[str, Any] | None:
    key = _output_object_key(task["id"], _safe_output_filename(task.get("output_filename")))
    return next((item for item in _task_objects(task["id"]) if item.get("object_key") == key), None)


def _read_output(task: dict[str, Any]) -> bytes | None:
    if _minio_enabled():
        item = _output_object(task)
        if item:
            return get_object_storage().get_bytes(item["object_key"])
        inline = task.get("output_content")
        return str(inline).encode("utf-8") if inline is not None else None
    path = _task_output_path(task)
    return path.read_bytes() if path.exists() else None


def _convert_from_minio(task: dict[str, Any], created_by: str | None) -> tuple[int, int, bytes]:
    output_name = _safe_output_filename(task.get("output_filename"))
    output_lines: list[str] = []
    input_count = 0
    output_count = 0
    for item in _input_objects(task["id"]):
        input_count += 1
        raw = get_object_storage().get_bytes(item["object_key"])
        data = json.loads(raw.decode("utf-8"))
        if isinstance(data, list):
            records = data
        elif isinstance(data, dict):
            records = [data]
        else:
            raise ValueError(f"JSON must be object or array: {item.get('file_name')}")
        for record in records:
            output_lines.append(json.dumps(record, ensure_ascii=False) + "\n")
            output_count += 1
    output = "".join(output_lines).encode("utf-8")
    store = get_platform_store()
    with store.connect() as conn:
        if should_store_in_minio(len(output)):
            output_object = _register_object(
                task["id"],
                version_id="output",
                object_key=_output_object_key(task["id"], output_name),
                file_name=output_name,
                content_type="application/jsonl",
                content=output,
                created_by=created_by,
            )
            conn.execute(
                "UPDATE data_convert_tasks SET output_storage_object_id=%s, output_content=NULL, storage_backend='minio' WHERE id=%s",
                (output_object["id"], task["id"]),
            )
        else:
            conn.execute(
                "UPDATE data_convert_tasks SET output_storage_object_id=NULL, output_content=%s, storage_backend='database' WHERE id=%s",
                (output.decode("utf-8"), task["id"]),
            )
    return input_count, output_count, output


def _convert_from_local(task: dict[str, Any]) -> tuple[int, int, bytes]:
    input_dir = _input_dir(task["id"])
    output_dir = _output_dir(task["id"])
    input_dir.mkdir(parents=True, exist_ok=True)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = _task_output_path(task)
    output_path.unlink(missing_ok=True)
    input_count = 0
    output_count = 0
    with output_path.open("w", encoding="utf-8") as output_file:
        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
            data = json.loads(json_file.read_text(encoding="utf-8"))
            records = data if isinstance(data, list) else [data] if isinstance(data, dict) else None
            if records is None:
                raise ValueError(f"JSON must be object or array: {json_file.name}")
            for record in records:
                output_file.write(json.dumps(record, ensure_ascii=False) + "\n")
                output_count += 1
    return input_count, output_count, output_path.read_bytes()


@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:
        if is_admin(current_user):
            # 管理员可见全部
            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]
        else:
            # 普通用户只能看到自己创建的
            user_id = current_user.get("id")
            rows = conn.execute(
                "SELECT * FROM data_convert_tasks WHERE deleted_at IS NULL AND created_by=%s "
                "ORDER BY create_time DESC LIMIT %s OFFSET %s",
                (user_id, page_size, (page - 1) * page_size),
            ).fetchall()
            total = conn.execute(
                "SELECT COUNT(*) FROM data_convert_tasks WHERE deleted_at IS NULL AND created_by=%s",
                (user_id,)
            ).fetchone()[0]
    return ok({"items": [dict(r) for r in rows], "total": total})


@router.post("")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.CREATE, target_type="convert_task", target_name_param="name")
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()
    user_id = current_user.get("id")
    store = get_platform_store()
    with store.connect() as conn:
        conn.execute(
            "INSERT INTO data_convert_tasks (id, name, description, output_filename, created_by) "
            "VALUES (%s, %s, %s, %s, %s)",
            (task_id, name, description, output_filename, user_id),
        )
    # MinIO 是正式存储；本地目录只在关闭 MinIO 的旧兼容模式下创建。
    if not _minio_enabled():
        _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")
    # 附加输入文件列表；旧任务没有对象记录时继续读取本地兼容目录。
    files = []
    if _minio_enabled():
        files = [
            {"name": item.get("file_name") or Path(item["object_key"]).name, "size": item.get("byte_size") or 0}
            for item in _input_objects(task_id)
        ]
    else:
        input_dir = _input_dir(task_id)
        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")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.UPLOAD, target_type="convert_task", target_name_param="task_id")
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")
    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}")
        content = await upload.read()
        if _minio_enabled():
            _register_object(
                task_id,
                version_id=f"input-{hashlib.sha256(name.encode('utf-8')).hexdigest()[:16]}",
                object_key=_input_object_key(task_id, name),
                file_name=name,
                content_type=upload.content_type or "application/json",
                content=content,
                created_by=task.get("created_by") or current_user.get("id"),
            )
        else:
            input_dir = _input_dir(task_id)
            input_dir.mkdir(parents=True, exist_ok=True)
            (input_dir / name).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:
        if _minio_enabled():
            input_count, output_count, output = _convert_from_minio(
                task, task.get("created_by") or current_user.get("id")
            )
        else:
            input_count, output_count, output = _convert_from_local(task)
        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.decode("utf-8")
        size_bytes = len(content.encode("utf-8"))
        dataset = store.create_dataset({
            "name": task["name"],
            "type": "train",
            "storage_type": "minio" if should_store_in_minio(size_bytes) else ("database" if _minio_enabled() else "local"),
            "source": "upload",
            "task_id": task_id,
            "size": f"{size_bytes} B",
            "count": output_count,
            "description": f"由数据类型转换任务 {task_id} 自动导入",
            "created_by": task.get("created_by") or current_user.get("id"),
        })
        dataset_id = dataset["id"]
        with store.connect() as conn:
            dataset_file = store.add_dataset_file(conn, dataset_id, _safe_output_filename(task.get("output_filename")), content)
        if should_store_in_minio(len(output)):
            output_name = _safe_output_filename(task.get("output_filename"))
            object_key = f"datasets/{dataset_id}/versions/{dataset_file.get('active_version_id') or dataset_file['id']}/{output_name}"
            uploaded = get_object_storage().put_bytes(object_key, output, "application/jsonl")
            storage_object = store.create_storage_object({
                "resource_type": "dataset", "resource_id": dataset_id,
                "version_id": dataset_file.get("active_version_id") or dataset_file["id"],
                "bucket": uploaded["bucket"], "object_key": object_key,
                "file_name": output_name, "content_type": "application/jsonl",
                "byte_size": len(output), "checksum_sha256": hashlib.sha256(output).hexdigest(),
                "status": "available", "created_by": task.get("created_by") or current_user.get("id"),
            })
            store.link_dataset_file_storage_object(dataset_file["id"], storage_object["id"])
        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")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.CONVERT, target_type="convert_task", target_name_param="task_id")
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:
        if _minio_enabled():
            input_count, output_count, _ = _convert_from_minio(task, task.get("created_by") or current_user.get("id"))
        else:
            input_count, output_count, _ = _convert_from_local(task)
        # 更新任务状态
        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 = _read_output(task)
    if output is None:
        raise fail(404, "output file not found")
    if _minio_enabled():
        return Response(content=output, media_type="application/octet-stream", headers={
            "Content-Disposition": f"attachment; filename={_safe_output_filename(task.get('output_filename'))}"
        })
    return FileResponse(
        str(_task_output_path(task)),
        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 = _read_output(task)
    if output is None:
        raise fail(404, "output file not found")
    content = output.decode("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": "minio" if should_store_in_minio(size_bytes) else ("database" if _minio_enabled() else "local"),
        "source": "upload",
        "task_id": task_id,
        "size": f"{size_bytes} B",
        "count": task["output_count"],
        "description": description,
        "created_by": task.get("created_by") or (current_user.get("id") if current_user else None),
    })
    dataset_id = dataset["id"]
    with store.connect() as conn:
        dataset_file = store.add_dataset_file(conn, dataset_id, _safe_output_filename(task.get("output_filename")), content)
    if should_store_in_minio(len(output)):
        output_name = _safe_output_filename(task.get("output_filename"))
        object_key = f"datasets/{dataset_id}/versions/{dataset_file.get('active_version_id') or dataset_file['id']}/{output_name}"
        uploaded = get_object_storage().put_bytes(object_key, output, "application/jsonl")
        storage_object = store.create_storage_object({
            "resource_type": "dataset", "resource_id": dataset_id,
            "version_id": dataset_file.get("active_version_id") or dataset_file["id"],
            "bucket": uploaded["bucket"], "object_key": object_key,
            "file_name": output_name, "content_type": "application/jsonl",
            "byte_size": len(output), "checksum_sha256": hashlib.sha256(output).hexdigest(),
            "status": "available", "created_by": task.get("created_by") or (current_user.get("id") if current_user else None),
        })
        store.link_dataset_file_storage_object(dataset_file["id"], storage_object["id"])
    return ok({"dataset_id": dataset_id, "name": dataset_name})


@router.delete("/{task_id}")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.DELETE, target_type="convert_task", target_name_param="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,),
        )
    if _minio_enabled():
        for item in _task_objects(task_id):
            try:
                get_object_storage().delete(item["object_key"])
                store.update_storage_object(item["id"], {"status": "deleted"})
            except Exception:
                pass
    else:
        # 旧兼容数据仍清理本地目录。
        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