import os
import pickle
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from io import BytesIO
import subprocess

app = FastAPI()

@app.post("/container_task")
async def container_task(file: UploadFile = File(...)):
    # Save the uploaded file to disk
    input_filepath = "input.pkl"
    with open(input_filepath, "wb") as f:
        f.write(await file.read())

    # Process the unpickle_param from the file
    try:
        with open(input_filepath, 'rb') as f:
            unpickle_param = pickle.load(f)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to unpickle the input file: {str(e)}")

    # Execute the Docker command
    command = ["docker", "run", "--rm", "bbdown", str(unpickle_param)]
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    # Stream the output of the command
    stdout_stream = BytesIO()
    stderr_stream = BytesIO()

    # Create a generator to stream output
    def stream_output():
        while True:
            output = process.stdout.readline()
            if output == b"" and process.poll() is not None:
                break
            if output:
                stdout_stream.write(output)

        stderr_output = process.stderr.read()
        stderr_stream.write(stderr_output)
        yield ""

    # Return the StreamingResponse for the current output
    async def response_stream():
        for _ in stream_output():
            yield stdout_stream.getvalue()
            stdout_stream.seek(0)  # Rewind for next read
            stdout_stream.truncate()  # Clear for next fill

    # Run the process and wait for completion
    process.wait()

    # Check for errors
    if process.returncode != 0:
        raise HTTPException(status_code=500, detail=f"Docker command failed with error: {stderr_stream.getvalue().decode()}")

    # Create a new pickle file as output
    output_filepath = "output.pkl"
    with open(output_filepath, 'wb') as f:
        f.write(b"Your output data here.")  # Replace this with actual output data

    # Return the output file
    return StreamingResponse(open(output_filepath, "rb"), media_type='application/octet-stream',
                             headers={"Content-Disposition": f"attachment; filename={os.path.basename(output_filepath)}"})

# To run the application, use: uvicorn your_file_name:app --reload
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import time
import asyncio

app = FastAPI()

async def stream_generator():
    for i in range(10):
        yield f"Data chunk {i}\n"
        await asyncio.sleep(1)  # Simulating some delay

@app.get("/stream")
async def stream_response():
    return StreamingResponse(stream_generator(), media_type="text/plain")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)


def client_call(*args, **kwargs):

    result = execute(*args, **kwargs)

    result.text
    result.file_manifest
    result.files