import gradio as gr
import io
import sys
import pandas as pd
from PIL import Image

from rayt_rust._core import Camera, Point3, Vec3
from rayt.gpu_utils import is_cuda_available
from rayt.scene import random_scene
from rayt.cuda_renderer import render_with_cuda
from rayt.numba_renderer import render_with_numba
from rayt.rust_renderer import render_with_rust
from numba import cuda

IS_CUDA_AVAILABE, _ = is_cuda_available()


def render_image(image_width, samples_per_pixel, aspect_ratio_str, max_depth, engine):
    """
    Renders an image using the rayt library and returns it.
    """
    # --- Capture stdout ---
    old_stdout = sys.stdout
    sys.stdout = captured_output = io.StringIO()

    # --- Run the rendering logic ---
    aspect_ratio = eval(aspect_ratio_str)
    image_height = int(image_width / aspect_ratio)
    world = random_scene()
    camera = Camera(
        lookfrom=Point3(13, 2, 3),
        lookat=Point3(0, 0, 0),
        vup=Vec3(0, 1, 0),
        vfov=20.0,
        aspect_ratio=aspect_ratio,
        aperture=0.1,
        focus_dist=10.0,
    )

    render_funcs = {
        "rust": render_with_rust,
        "numba": render_with_numba,
        "cuda": render_with_cuda,
    }

    if engine == "cuda" and not IS_CUDA_AVAILABE:
        gr.Warning(
            "CUDA renderer is not available in this server. Falling back to numba."
        )
        engine = "numba"

    render_func = render_funcs[engine]

    render_func(world, camera, image_width, image_height, samples_per_pixel, max_depth)

    # --- Restore stdout and get the PPM data ---
    sys.stdout = old_stdout
    ppm_data = captured_output.getvalue()

    # --- Convert PPM to Image ---
    # The PPM data is in a string, so we need to wrap it in a file-like object
    ppm_file = io.BytesIO(ppm_data.encode())
    image = Image.open(ppm_file)

    return image


# --- Create the Gradio Interface ---
with gr.Blocks() as iface:
    gr.Markdown("# Ray Tracing in One Weekend")

    with gr.Row():
        with gr.Column():
            gr.Markdown(
                "A Python implementation of the book by Peter Shirley, rendered with Numba."
            )

            if IS_CUDA_AVAILABE:
                device = cuda.get_current_device()
                cc = device.compute_capability
                cuda_info_dataframe = pd.DataFrame(
                    {
                        "CUDA support": ["AVAILABLE"],
                        "Found GPU": [device.name.decode("utf-8")],
                        "Compute Capability": [f"{cc[0]}.{cc[1]}"],
                    }
                )
            else:
                cuda_info_dataframe = pd.DataFrame(
                    {
                        "CUDA support": ["NOT AVAILABLE"],
                        "Found GPU": ["-"],
                        "Compute Capability": ["-"],
                    }
                )

            gr.DataFrame(cuda_info_dataframe)
            gr.Markdown(
                "If you have an NVIDIA GPU on your local machine, you can clone the project and "
                "run it locally. You can also deploy it to a server with CUDA support for faster "
                "rendering."
            )

            gr.Image(
                label="Sample Output",
                value="image.webp",
                show_download_button=False,
                show_fullscreen_button=False,
                show_label=True,
            )

        with gr.Column():
            gr.Markdown("## PERFORMANCE")
            gr.Markdown("RAYT tested on this machine:")
            test_info_dataframe = pd.DataFrame(
                {
                    "CPU": ["AMD Radeon 9 9900X"],
                    "GPU": ["NVIDIA RTX 5080"],
                    "Image Width": ["600"],
                    "Samples Per Pixel": ["100"],
                }
            )
            gr.DataFrame(test_info_dataframe)
            gr.Image(
                label="Performance",
                value="performance.png",
                show_download_button=False,
                show_fullscreen_button=False,
                show_label=True,
            )

    with gr.Row():
        with gr.Column():
            gr.Markdown("## INPUT")
            image_width = gr.Slider(
                label="Image Width", minimum=100, maximum=1200, value=300, step=10
            )
            samples_per_pixel = gr.Slider(
                label="Samples Per Pixel", minimum=1, maximum=200, value=20, step=1
            )
            max_depth = gr.Slider(
                label="Max Ray Depth", minimum=1, maximum=100, value=50, step=1
            )
            aspect_ratio = gr.Dropdown(
                label="Aspect Ratio", choices=["16/9", "3/2", "1/1"], value="16/9"
            )
            engine = gr.Dropdown(
                label="Engine", choices=["rust", "numba", "cuda"], value="rust"
            )
            render_button = gr.Button("Render Image")
        with gr.Column():
            gr.Markdown("## OUTPUT")
            output_image = gr.Image(label="Rendered Image", type="pil")

    render_button.click(
        fn=render_image,
        inputs=[image_width, samples_per_pixel, aspect_ratio, max_depth, engine],
        outputs=output_image,
    )

if __name__ == "__main__":
    iface.launch()