import os
import aiohttp
import asyncio
import gradio as gr

# Função para dividir o arquivo em chunks
def chunk_file(file_path, chunk_size=50*1024*1024):
    chunks = []
    with open(file_path, 'rb') as f:
        while chunk := f.read(chunk_size):
            chunks.append(chunk)
    return chunks

# Função para fazer o upload de um chunk
async def upload_chunk(session, url, chunk, chunk_number):
    async with session.post(url, data={'file': chunk, 'chunk_number': chunk_number}) as response:
        return await response.text()

# Função principal para fazer o upload dos chunks em paralelo
async def upload_file_chunks(file_path):
    upload_url = "https://<your_space_name>.hf.space/upload"
    chunks = chunk_file(file_path)
    async with aiohttp.ClientSession() as session:
        tasks = [upload_chunk(session, upload_url, chunk, i) for i, chunk in enumerate(chunks)]
        results = await asyncio.gather(*tasks)
    return results

# Função chamada pelo Gradio
def upload_file(file):
    file_path = os.path.join("uploads", file.name)
    with open(file_path, "wb") as f:
        f.write(file.read())
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    result = loop.run_until_complete(upload_file_chunks(file_path))
    return f"Upload completed: {result}"

# Interface do Gradio
demo = gr.Interface(
    fn=upload_file,
    inputs=gr.inputs.File(label="Upload a file"),
    outputs="text",
    live=True,
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch(share=True)