import gradio as gr
import tempfile
from pathlib import Path
from wrapper import run_pipeline_on_image
import numpy as np
from PIL import Image
from itertools import product
import logging

def process(image_path):
    if not image_path:
        return None, None, None, None, [], ""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Copy uploaded file as-is to preserve format/bit-depth
        src = Path(image_path)
        ext = src.suffix.lstrip('.') or 'tif'
        img_path = Path(tmpdir) / f"input.{ext}"
        try:
            try:
                # Log original file stats
                size_bytes = src.stat().st_size if src.exists() else 0
                try:
                    with Image.open(src) as im_src:
                        frames = getattr(im_src, 'n_frames', 1)
                        logging.info(f"Uploaded file: path={src}, size_bytes={size_bytes}, mode={im_src.mode}, size={im_src.size}, frames={frames}")
                except Exception:
                    logging.info(f"Uploaded file: path={src}, size_bytes={size_bytes} (PIL open failed)")
            except Exception:
                pass
            img_bytes = src.read_bytes()
            img_path.write_bytes(img_bytes)
            # Log copied file as read by PIL
            try:
                with Image.open(img_path) as im_tmp:
                    frames_tmp = getattr(im_tmp, 'n_frames', 1)
                    logging.info(f"Temp input: path={img_path}, mode={im_tmp.mode}, size={im_tmp.size}, frames={frames_tmp}")
            except Exception:
                logging.info(f"Temp input: path={img_path} (PIL open failed)")
        except Exception:
            # Fallback: save via PIL if direct copy fails
            Image.open(src).save(img_path)
        outputs = run_pipeline_on_image(str(img_path), tmpdir, save_artifacts=True)

        def load_pil(path_str):
            try:
                if not path_str:
                    return None
                im = Image.open(path_str)
                im = im.convert('RGB')
                copied = im.copy()
                im.close()
                return copied
            except Exception:
                return None

        composite = load_pil(outputs.get('Composite'))
        overlay = load_pil(outputs.get('Overlay'))
        mask = load_pil(outputs.get('Mask'))
        size_img = load_pil(str(Path(tmpdir) / 'results/size.size_analysis.png'))
        # Texture LBP green path
        lbp_path = Path(tmpdir) / 'texture_output/lbp_green.png'
        texture_img = load_pil(str(lbp_path)) if lbp_path.exists() else None
        order = ['NDVI', 'GNDVI', 'SAVI']
        gallery_items = [load_pil(outputs[k]) for k in order if k in outputs]
        stats_text = outputs.get('StatsText', '')
        return size_img, composite, mask, overlay, texture_img, gallery_items, stats_text

with gr.Blocks() as demo:
    gr.Markdown("# 🌿 Automated Plant Analysis Demo")
    gr.Markdown("Upload a sorghum plant image to compute and visualize composite, mask, overlay, texture (LBP), vegetation indices, and statistics.")

    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="filepath", label="Upload Image")
            run = gr.Button("Run Pipeline", variant="primary")

    with gr.Row():
        size_img = gr.Image(type="pil", label="Morphology Size", interactive=False)
        composite_img = gr.Image(type="pil", label="Composite (Segmentation Input)", interactive=False)
        mask_img = gr.Image(type="pil", label="Mask", interactive=False)
        overlay_img = gr.Image(type="pil", label="Segmentation Overlay", interactive=False)

    with gr.Row():
        texture_img = gr.Image(type="pil", label="Texture LBP (Green Band)", interactive=False)

    gallery = gr.Gallery(label="Vegetation Indices", columns=3, height="auto")
    stats = gr.Textbox(label="Statistics", lines=4)

    run.click(process, inputs=inp, outputs=[size_img, composite_img, mask_img, overlay_img, texture_img, gallery, stats])

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