load("@dictation_pip//:requirements.bzl", "requirement")
load("@rules_python//python:pip.bzl", "compile_pip_requirements")
load("@rules_python//python:py_binary.bzl", "py_binary")
load("@rules_python//python:py_library.bzl", "py_library")
load("@rules_python//python:py_test.bzl", "py_test")
load("@rules_shell//shell:sh_binary.bzl", "sh_binary")

compile_pip_requirements(
    name = "requirements",
    src = "requirements.in",
    extra_args = ["--allow-unsafe"],
    requirements_txt = "requirements_lock.txt",
)

py_library(
    name = "dictation_lib",
    srcs = [
        "__init__.py",
        "audio.py",
        "config.py",
        "server.py",
        "session.py",
        "transcriber.py",
    ],
    data = glob(["web/**"]),
    deps = [
        requirement("aiortc"),
        requirement("av"),
        requirement("fastapi"),
        requirement("faster-whisper"),
        requirement("huggingface-hub"),
        requirement("numpy"),
        requirement("uvicorn"),
    ],
)

py_binary(
    name = "dictation_bin",
    srcs = ["main.py"],
    data = glob(["web/**"]),
    main = "main.py",
    deps = [
        ":dictation_lib",
        requirement("ctranslate2"),
        requirement("nvidia-cublas-cu12"),
        requirement("nvidia-cudnn-cu12"),
    ],
)

py_binary(
    name = "webrtc_smoke",
    srcs = ["webrtc_smoke.py"],
    main = "webrtc_smoke.py",
    deps = [
        requirement("aiortc"),
        requirement("httpx"),
    ],
)

[
    sh_binary(
        name = command,
        srcs = ["launcher.sh"],
        args = [
            "$(rootpath :dictation_bin)",
            command,
        ],
        data = [
            ":dictation_bin",
            requirement("nvidia-cublas-cu12"),
            requirement("nvidia-cudnn-cu12"),
        ],
        visibility = ["//infinite_canvas:__pkg__"],
    )
    for command in [
        "server",
        "preflight",
        "download_model",
        "transcribe",
    ]
]

py_test(
    name = "audio_test",
    srcs = ["audio_test.py"],
    deps = [
        ":dictation_lib",
        requirement("numpy"),
    ],
)

py_test(
    name = "config_test",
    srcs = ["config_test.py"],
    deps = [":dictation_lib"],
)

py_test(
    name = "server_test",
    srcs = ["server_test.py"],
    data = glob(["web/**"]),
    deps = [
        ":dictation_lib",
        requirement("httpx"),
    ],
)

py_test(
    name = "webrtc_test",
    srcs = ["webrtc_test.py"],
    deps = [
        ":dictation_lib",
        requirement("aiortc"),
        requirement("av"),
        requirement("numpy"),
    ],
)

test_suite(
    name = "all_tests",
    tests = [
        ":audio_test",
        ":config_test",
        ":server_test",
        ":webrtc_test",
    ],
)