from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse
from fastapi import HTTPException
from colab_easy_ui.Functions import ProgressFunctionCallResponse
import os
import sys
from fastapi.responses import FileResponse
from colab_easy_ui.const import sanitaize_path

sys.path.append("../easy-vc-dev")
from easy_vc_dev.onnx.export_onnx import export_onnx as _export_onnx  # noqa


def onnx_download(
    project_name: str,
    torch_file: str,
):
    if sanitaize_path(project_name) is False or sanitaize_path(torch_file) is False:
        data = ProgressFunctionCallResponse(
            status="NG",
            uuid="",
            message=f"not valid project name {project_name} or torch_file {torch_file}",
        )
        json_compatible_item_data = jsonable_encoder(data)
        return JSONResponse(content=json_compatible_item_data)

    file_path = os.path.join("trainer", project_name, "logs", os.path.splitext(torch_file)[0] + ".onnx")

    if not os.path.exists(file_path) or not os.path.isfile(file_path):
        raise HTTPException(status_code=404, detail="File not found")

    return FileResponse(path=file_path, filename=os.path.basename(file_path), media_type="application/octet-stream")
