import sys
root_dir = __file__.rsplit("/", 1)[0]
if root_dir not in sys.path:
    sys.path.append(root_dir)

import gradio as gr
import asyncio
import os

from demo.modules.search import build_search_module
from demo.modules.compute_score import build_score_computation
from demo.modules.tmalign import build_TMalign


# Build demo
with gr.Blocks() as demo:
    build_search_module()
    build_score_computation()
    build_TMalign()


# Asynchronously download the required index files
async def run_cmd(cmd):
    process = await asyncio.create_subprocess_shell(
        cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await process.communicate()
    print(f'[{cmd!r} exited with {process.returncode}]')
    if stdout:
        print(f'[stdout]\n{stdout.decode()}')
    if stderr:
        print(f'[stderr]\n{stderr.decode()}')


async def download_faiss_index():
    await asyncio.gather(
        run_cmd("echo Hello World!"),
        run_cmd('huggingface-cli download westlake-repl/ProTrek-faiss-index  --repo-type dataset  --local-dir /data/ProTrek-faiss-index --include "ProTrek_650M_UniRef50/UniRef50/*"'),
    )


if __name__ == '__main__':
    # Run the async function to download the index files
    asyncio.run(download_faiss_index())
    
    # Run the demo
    demo.launch()