YG_FT/backend/app/api/v1/endpoints/platform.py

from __future__ import annotations

from typing import Any

from fastapi import APIRouter, Body, File, HTTPException, Query, UploadFile
from fastapi.responses import PlainTextResponse

from app.db.platform_store import get_platform_store

router = APIRouter()


def ok(data: Any = None, message: str = "ok") -> dict[str, Any]:
    return {"code": 0, "message": message, "data": data}


def fail(status_code: int, message: str) -> HTTPException:
    return HTTPException(status_code=status_code, detail={"code": status_code, "message": message, "data": None})


@router.post("/login")
async def login(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    user = get_platform_store().login(payload.get("username", ""), payload.get("password", ""))
    if not user:
        raise fail(401, "invalid username or password")
    return ok({"token": f"platform-token-{user['id']}", "user": user})


@router.get("/me")
async def me() -> dict[str, Any]:
    return ok(get_platform_store().users()[0])


@router.get("/dashboard/overview")
async def dashboard_overview() -> dict[str, Any]:
    store = get_platform_store()
    tasks = store.tasks()
    return ok(
        {
            "models": len(store.models()),
            "datasets": len(store.datasets()),
            "fine_tune_tasks": len(tasks),
            "running_tasks": len([t for t in tasks if t["status"] in {"syncing", "queued", "running"}]),
            "compute_nodes": len(store.compute_nodes()),
            "gpus": len(store.gpus()),
        }
    )


@router.get("/system-info")
async def system_info() -> dict[str, Any]:
    return ok(get_platform_store().system_info())


@router.get("/users")
async def users() -> dict[str, Any]:
    return ok(get_platform_store().users())


@router.post("/users")
async def create_user(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    return ok(get_platform_store().create_user(payload))


@router.put("/users/{user_id}")
async def update_user(user_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().update_user(user_id, payload))
    except KeyError:
        raise fail(404, "user not found")


@router.delete("/users/{user_id}")
async def delete_user(user_id: str, current_username: str | None = Query(default=None)) -> dict[str, Any]:
    try:
        get_platform_store().delete_user(user_id)
        return ok({"deleted": user_id, "current_username": current_username})
    except KeyError:
        raise fail(404, "user not found")
    except ValueError as exc:
        raise fail(400, str(exc))


@router.get("/model-manage/local-models")
async def local_models() -> dict[str, Any]:
    models = [{"path": item.get("path") or "", "name": item["name"]} for item in get_platform_store().models()]
    return ok({"models": models})


@router.get("/model-manage/trained-models")
async def trained_models() -> dict[str, Any]:
    return ok({"models": get_platform_store().trained_models()})


@router.delete("/model-manage/trained-models/{model_id}")
async def delete_trained_model(model_id: str, type: str = Query(default="merged")) -> dict[str, Any]:
    return ok({"deleted": model_id, "type": type})


@router.get("/model-manage/name/{name}")
async def model_by_name(name: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().model_by_name(name))
    except KeyError:
        raise fail(404, "model not found")


@router.get("/model-manage")
async def model_list() -> dict[str, Any]:
    return ok(get_platform_store().models())


@router.post("/model-manage")
async def create_model(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    return ok(get_platform_store().create_model(payload))


@router.get("/model-manage/{model_id}")
async def model_detail(model_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().model(model_id))
    except KeyError:
        raise fail(404, "model not found")


@router.put("/model-manage/{model_id}")
async def update_model(model_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().update_model(model_id, payload))
    except KeyError:
        raise fail(404, "model not found")


@router.put("/model-manage/{model_id}/purpose")
async def update_model_purpose(model_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().update_model(model_id, {"purpose": payload.get("purpose", "training")}))
    except KeyError:
        raise fail(404, "model not found")


@router.delete("/model-manage/{model_id}")
async def delete_model(model_id: str) -> dict[str, Any]:
    get_platform_store().delete_model(model_id)
    return ok({"deleted": model_id})


@router.post("/model-manage/merge")
async def merge_model(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    return ok({"job_id": "merge-sample-job", "status": "queued", **payload})


@router.get("/dataset-manage/preview/{file_id}")
async def dataset_preview(file_id: str) -> dict[str, Any]:
    try:
        row = get_platform_store().dataset_file(file_id)
        return ok({"content": row["content"]})
    except KeyError:
        raise fail(404, "dataset file not found")


@router.get("/dataset-manage/versions/{file_id}")
async def dataset_versions(file_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().file_versions(file_id))
    except KeyError:
        raise fail(404, "dataset file not found")


@router.get("/dataset-manage/versions/{file_id}/{version_id}")
async def dataset_version_content(file_id: str, version_id: str) -> dict[str, Any]:
    try:
        row = get_platform_store().dataset_file(file_id)
        versions = get_platform_store().file_versions(file_id)["versions"]
        version = next((item for item in versions if item["id"] == version_id), None)
        if not version:
            raise KeyError(version_id)
        return ok({"version": version, "content": row["content"]})
    except KeyError:
        raise fail(404, "dataset version not found")


@router.post("/dataset-manage/versions/{file_id}")
async def create_dataset_version(file_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().create_file_version(file_id, payload))
    except KeyError:
        raise fail(404, "dataset file not found")


@router.put("/dataset-manage/versions/{file_id}/active")
async def activate_dataset_version(file_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().activate_file_version(file_id, payload["version_id"]))
    except KeyError:
        raise fail(404, "dataset version not found")


@router.delete("/dataset-manage/versions/{file_id}/{version_id}")
async def delete_dataset_version(file_id: str, version_id: str) -> dict[str, Any]:
    return ok(get_platform_store().file_versions(file_id))


@router.post("/dataset-manage/upload/{dataset_id}")
async def upload_dataset_files(dataset_id: str, files: list[UploadFile] = File(default=[])) -> dict[str, Any]:
    created: list[dict[str, Any]] = []
    store = get_platform_store()
    try:
        store.dataset(dataset_id)
    except KeyError:
        raise fail(404, "dataset not found")
    with store.connect() as conn:
        for file in files:
            raw = await file.read()
            content = raw.decode("utf-8", errors="replace")
            created.append(store.add_dataset_file(conn, dataset_id, file.filename or "upload.jsonl", content))
    return ok({"files": created})


@router.get("/dataset-manage/download/{dataset_id}")
async def download_dataset(dataset_id: str) -> PlainTextResponse:
    dataset = get_platform_store().dataset(dataset_id)
    content = "\n".join([f"{file['name']}" for file in dataset.get("files", [])])
    return PlainTextResponse(content, media_type="text/plain")


@router.get("/dataset-manage/download/{dataset_id}/{file_id}")
async def download_dataset_file(dataset_id: str, file_id: str, version_id: str | None = Query(default=None)) -> PlainTextResponse:
    row = get_platform_store().dataset_file(file_id)
    return PlainTextResponse(row["content"], media_type="text/plain")


@router.get("/dataset-manage")
async def dataset_list() -> dict[str, Any]:
    return ok(get_platform_store().datasets())


@router.post("/dataset-manage")
async def create_dataset(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    dataset = get_platform_store().create_dataset(payload)
    return ok({"id": dataset["id"]})


@router.get("/dataset-manage/{dataset_id}")
async def dataset_detail(dataset_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().dataset(dataset_id))
    except KeyError:
        raise fail(404, "dataset not found")


@router.put("/dataset-manage/{dataset_id}")
async def update_dataset(dataset_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().update_dataset(dataset_id, payload))
    except KeyError:
        raise fail(404, "dataset not found")


@router.delete("/dataset-manage/{dataset_id}")
async def delete_dataset(dataset_id: str) -> dict[str, Any]:
    get_platform_store().delete_dataset(dataset_id)
    return ok({"deleted": dataset_id})


@router.get("/fine-tune/check-name")
async def check_fine_tune_name(name: str = Query(...)) -> dict[str, Any]:
    exists = any(task["name"] == name for task in get_platform_store().tasks())
    return ok({"exists": exists})


@router.get("/fine-tune/progress/{task_id}")
async def fine_tune_progress(task_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().progress(task_id))
    except KeyError:
        raise fail(404, "fine tune task not found")


@router.post("/fine-tune/tensorboard/start")
async def tensorboard_start() -> dict[str, Any]:
    return ok({"status": "running", "url": "http://localhost:6006"})


@router.get("/fine-tune")
async def fine_tune_list() -> dict[str, Any]:
    return ok(get_platform_store().tasks())


@router.post("/fine-tune")
async def create_fine_tune(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    task = get_platform_store().create_task(payload)
    return ok({"id": task["id"]})


@router.post("/fine-tune/start")
async def start_fine_tune(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().start_task(payload))
    except KeyError:
        raise fail(404, "fine tune task not found")
    except RuntimeError as exc:
        raise fail(409, str(exc))


@router.get("/fine-tune/{task_id}")
async def fine_tune_detail(task_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().task(task_id))
    except KeyError:
        raise fail(404, "fine tune task not found")


@router.put("/fine-tune/{task_id}")
async def update_fine_tune(task_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().update_task(task_id, payload))
    except KeyError:
        raise fail(404, "fine tune task not found")


@router.post("/fine-tune/stop/{task_id}")
async def stop_fine_tune(task_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().stop_task(task_id))
    except KeyError:
        raise fail(404, "fine tune task not found")


@router.post("/fine-tune/{task_id}/stop")
async def stop_fine_tune_alt(task_id: str) -> dict[str, Any]:
    return await stop_fine_tune(task_id)


@router.delete("/fine-tune/{task_id}")
async def delete_fine_tune(task_id: str) -> dict[str, Any]:
    get_platform_store().delete_task(task_id)
    return ok({"deleted": task_id})


@router.get("/fine-tune/{task_id}/overview")
async def fine_tune_overview(task_id: str) -> dict[str, Any]:
    task = get_platform_store().task(task_id)
    return ok({"task": task, "progress": get_platform_store().progress(task_id)})


@router.get("/fine-tune/{task_id}/checkpoints")
async def fine_tune_checkpoints(task_id: str) -> dict[str, Any]:
    task = get_platform_store().task(task_id)
    checkpoints = []
    for step in [50, 100, 150]:
        if task.get("progress", 0) >= min(100, step // 2):
            checkpoints.append({"step": step, "path": f"/data/yg-ft/outputs/{task['name']}/checkpoint-{step}"})
    return ok(checkpoints)


@router.get("/compute/nodes")
async def compute_nodes() -> dict[str, Any]:
    return ok(get_platform_store().compute_nodes())


@router.post("/compute/nodes")
async def create_compute_node(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    return ok({"id": "node_sample_new", **payload, "status": "created"})


@router.put("/compute/nodes/{node_id}")
async def update_compute_node(node_id: str, payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().update_compute_node(node_id, payload))
    except KeyError:
        raise fail(404, "compute node not found")


@router.post("/compute/nodes/{node_id}/test-connection")
async def test_compute_node(node_id: str) -> dict[str, Any]:
    return ok({"node_id": node_id, "success": True, "latency_ms": 12})


@router.post("/compute/nodes/{node_id}/enable")
async def enable_compute_node(node_id: str) -> dict[str, Any]:
    return ok(get_platform_store().update_compute_node(node_id, {"enabled": True, "scheduler_status": "online"}))


@router.post("/compute/nodes/{node_id}/disable")
async def disable_compute_node(node_id: str) -> dict[str, Any]:
    return ok(get_platform_store().update_compute_node(node_id, {"enabled": False, "scheduler_status": "offline"}))


@router.post("/compute/nodes/{node_id}/drain")
async def drain_compute_node(node_id: str) -> dict[str, Any]:
    return ok(get_platform_store().update_compute_node(node_id, {"scheduler_status": "draining"}))


@router.get("/compute/nodes/{node_id}/replicas")
async def compute_node_replicas(node_id: str) -> dict[str, Any]:
    return ok(get_platform_store().replicas(node_id))


@router.get("/compute/gpus")
async def compute_gpus() -> dict[str, Any]:
    return ok(get_platform_store().gpus())


@router.get("/compute/queue")
async def compute_queue() -> dict[str, Any]:
    return ok(get_platform_store().queue())


@router.post("/internal/compute-sync/resources")
async def create_compute_sync(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    sync_id = get_platform_store().create_sync_job(payload.get("target_node_id", "node_01"), payload)
    return ok(get_platform_store().sync_job(sync_id))


@router.get("/internal/compute-sync/resources/{sync_id}")
async def compute_sync_detail(sync_id: str) -> dict[str, Any]:
    try:
        return ok(get_platform_store().sync_job(sync_id))
    except KeyError:
        raise fail(404, "sync job not found")


@router.get("/training-log-files")
async def training_log_files() -> dict[str, Any]:
    return ok(get_platform_store().training_log_files())


@router.get("/training-log-content")
async def training_log_content(file: str = Query(...)) -> dict[str, Any]:
    try:
        return ok(get_platform_store().training_log_content(file))
    except KeyError:
        raise fail(404, "training log not found")


@router.get("/log-files")
async def log_files(date: str | None = Query(default=None)) -> dict[str, Any]:
    return ok(get_platform_store().log_files(date))


@router.get("/log-content")
async def log_content(file: str = Query(...)) -> dict[str, Any]:
    return ok(get_platform_store().log_content(file))


@router.post("/web-log")
async def web_log(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
    return ok({"received": True, **payload})