import httpx
import numpy as np
from args import parse_arguments
import io

def process_yuv_frame(data: np.ndarray, resolution: str, format: str, url:str, model: str = "hd24k_full") -> np.ndarray:
    """
    Uploads a YUV file to the /upload endpoint for processing.

    Args:
        file_path (str): Path to the YUV file to upload.
        resolution (str): Video resolution (e.g., '720p', '1080p').
        format (str): YUV format (e.g., 'yuv420p', 'yuv422p').
        model (str, optional): The processing model to use (default is 'hd24k_full'). 
            Possible values: 'hd24k_full', 'hd24k_standard', 'sd2hd_full', 'sd2hd_standard'.

    Returns:
        str: The path to the processed file or an error message.
    """
    RES_MAP = {
        "720p": (1280, 720),
        "1080p": (1920, 1080)
    }

    if resolution not in RES_MAP.keys():
        return "Invalid resolution only accepts 720p and 1080p."

    in_width, in_height = RES_MAP[resolution]

    buffer = io.BytesIO()
    np.save(buffer, data)  # Binary .npy format
    buffer.seek(0)

    # Prepare the data and file for the request
    files = {
        'yuv_file': ('frame.yuv', buffer, 'application/octet-stream')
    }
    data = {
        "resolution": resolution,
        "format": format,
        "model": model,
    }
    # Important! need long timeout
    timeout = httpx.Timeout(60.0) 
    with httpx.Client(timeout=timeout) as client:
        # Send the request to the API endpoint
        response = client.post(f"{url}api/upload", files=files, data=data)
        
        if response.status_code != 200:
            print (f"Error during processing: {response.content}")
            exit()

        yuv_frame = np.frombuffer(response.content, dtype=np.uint16)
        return yuv_frame

# Example usage:
if __name__ == "__main__":
    args = parse_arguments()

    yuv_file = np.fromfile(args.input_file, dtype=np.uint16)
    result = process_yuv_frame(
        data=yuv_file,
        resolution=args.size,
        format=args.pixel_format,
        model=args.model,
        url=args.url
    )

    result.tofile(args.output_file)