# app.py

import os
from typing import List, Any

import gradio as gr
from PIL import Image

from pipeline import SmartCBC

# -------------------------------------------------
# Initialize pipeline ONCE (cached in Spaces)
# -------------------------------------------------
cbc = SmartCBC()   # loads YOLO + classifier once


# -------------------------------------------------
# Helper: convert uploaded files to PIL Images
# -------------------------------------------------
def files_to_pil_list(files: List[Any]) -> List[Image.Image]:
    """
    Gradio Files (file_count='multiple') returns a list of file objects or dicts.

    Each item commonly looks like:
      - {"name": "/tmp/....png", "orig_name": "...", ...}
      - or a file-like object with .name

    This helper normalizes them into a list of RGB PIL Images.
    """
    pil_list: List[Image.Image] = []

    if files is None:
        return pil_list

    for f in files:
        # Newer gradio often returns dicts with "name"
        if isinstance(f, dict) and "name" in f:
            path = f["name"]
        # Older style: File object with .name
        elif hasattr(f, "name"):
            path = f.name
        # Fallback: assume it's already a path-like
        else:
            path = str(f)

        if not os.path.isfile(path):
            raise FileNotFoundError(f"Uploaded file not found on disk: {path}")

        img = Image.open(path).convert("RGB")
        pil_list.append(img)

    return pil_list


# -------------------------------------------------
# Core Gradio wrapper
# -------------------------------------------------
def analyze_images(files, age, gender, output_mode):
    """
    Wrapper for SmartCBC.analyze().

    - Accepts one or multiple images from a Gradio Files input.
    - If a single image -> sends a single PIL.Image
      If multiple -> sends a list[Image.Image] (SmartCBC can route to analyze_batch).
    - Returns either a human-readable text report or full JSON.
    """
    if files is None or len(files) == 0:
        return "Please upload at least one image.", None

    pil_images = files_to_pil_list(files)

    if len(pil_images) == 1:
        image_input = pil_images[0]
    else:
        image_input = pil_images

    # Run SmartCBC pipeline
    result = cbc.analyze(
        image=image_input,
        age=age,
        gender=gender,
    )

    # Choose output mode
    if output_mode == "Text Report":
        return result.get("report_text", "No report generated."), None
    else:
        return None, result


# -------------------------------------------------
# Gradio UI Layout (compatible with Gradio 4.0.0)
# -------------------------------------------------
with gr.Blocks(title="SmartCBC - Multimodal Blood Analysis") as demo:

    gr.Markdown(
        """
    # 🩸 SmartCBC — Multimodal AI Blood Smear Analysis

    Upload **one or multiple** peripheral smear FOV images and get:
    - RBC / WBC / Platelet counts  
    - WBC subtype classification  
    - Aggregated multi-FOV differential  
    - Age-specific reference comparisons  
    - Clinical insights (non-diagnostic)  
    """
    )

    with gr.Row():
        # Use Files for multi-image upload (works on Gradio 4.0.0)
        img_in = gr.Files(
            label="Upload 1 or Multiple Blood Smear Images (FOVs)",
            file_count="multiple",
            file_types=["image"],
        )

        with gr.Column():
            age_in = gr.Number(label="Age (years)", value=30)
            gender_in = gr.Dropdown(
                ["", "M", "F"],
                label="Gender (optional)",
                value=""
            )
            output_mode = gr.Radio(
                ["Text Report", "Structured JSON"],
                value="Text Report",
                label="Output Format",
            )
            btn = gr.Button("Analyze")

    # OUTPUT AREAS
    txt_out = gr.Textbox(
        label="Report (Human Readable)",
        visible=True,
        lines=30,
        interactive=False,
    )

    json_out = gr.JSON(
        label="Structured Output (JSON)",
        visible=False,
    )

    # Toggle visibility based on output mode
    def toggle_output(mode):
        return (
            gr.update(visible=(mode == "Text Report")),
            gr.update(visible=(mode == "Structured JSON")),
        )

    output_mode.change(toggle_output, [output_mode], [txt_out, json_out])

    # Button Binding
    btn.click(
        analyze_images,
        inputs=[img_in, age_in, gender_in, output_mode],
        outputs=[txt_out, json_out],
    )


# -------------------------------------------------
# HF Spaces entrypoint
# -------------------------------------------------
if __name__ == "__main__":
    demo.launch()