ui/tabs/images.py

"""Tab 2: Images.

Upload and manage fire damage images for AI analysis.
MVP Simplification: Single room - images auto-assigned to the room.
"""

import uuid
import gradio as gr
from typing import Any
from PIL import Image
import io

from ui.state import SessionState, ImageFormData
from ui.components import image_store
from config.settings import settings


def create_tab() -> dict[str, Any]:
    """Create Tab 2 UI components.

    Returns:
        Dictionary of component references for event wiring.
    """
    gr.Markdown("### Fire Damage Images")
    gr.Markdown(
        f"*Upload up to {settings.max_images_per_assessment} images for AI analysis.*"
    )

    with gr.Row():
        with gr.Column(scale=2):
            image_upload = gr.Files(
                label="Upload Images (select multiple)",
                file_count="multiple",
                file_types=["image"],
                elem_id="image_upload",
            )
            image_description = gr.Textbox(
                label="Description (optional)",
                placeholder="e.g., View of ceiling deck from center aisle",
                elem_id="image_description",
                info="Applied to all images in batch",
            )

            with gr.Row():
                add_image_btn = gr.Button("Add Images", variant="primary")
                clear_upload_btn = gr.Button("Clear", variant="secondary")

        with gr.Column(scale=3):
            images_gallery = gr.Gallery(
                label="Images Added",
                columns=3,
                height="auto",
                elem_id="images_gallery",
            )
            with gr.Row():
                remove_last_btn = gr.Button("Remove Last Image", variant="secondary")
                clear_all_btn = gr.Button("Clear All Images", variant="stop")

    # Image count and status
    with gr.Row():
        image_count = gr.Textbox(
            label="Images Added",
            value="0 / 20",
            interactive=False,
        )

    # Validation status
    with gr.Row():
        validation_status = gr.HTML(
            value="",
            elem_id="tab2_validation",
        )

    # Resume warning (shown when images need re-upload)
    with gr.Row():
        resume_warning = gr.HTML(
            value="",
            elem_id="resume_warning",
            visible=False,
        )

    with gr.Row():
        back_btn = gr.Button("← Back to Room")
        validate_btn = gr.Button(
            "Validate & Continue to Observations →",
            variant="primary",
        )

    return {
        "image_upload": image_upload,
        "image_description": image_description,
        "add_image_btn": add_image_btn,
        "clear_upload_btn": clear_upload_btn,
        "images_gallery": images_gallery,
        "remove_last_btn": remove_last_btn,
        "clear_all_btn": clear_all_btn,
        "image_count": image_count,
        "validation_status": validation_status,
        "resume_warning": resume_warning,
        "back_btn": back_btn,
        "validate_btn": validate_btn,
    }


def add_image(
    session: SessionState,
    files: list | None,
    description: str,
) -> tuple[SessionState, list[tuple], str, str, None, str]:
    """Add one or more images to the session (batch upload).

    Images are automatically associated with the single room.

    Args:
        session: Current session state.
        files: List of uploaded file objects from gr.Files, each with a `name` attribute.
        description: Optional description applied to all images.

    Returns:
        Tuple of (session, gallery_data, validation_html, image_count,
                  cleared_files, cleared_description).
    """
    validation_html = ""

    # Validate input
    errors = []
    if not files or len(files) == 0:
        errors.append("Please upload at least one image")

    # Check capacity
    current_count = len(session.images)
    max_allowed = settings.max_images_per_assessment
    if files and current_count + len(files) > max_allowed:
        remaining = max_allowed - current_count
        if remaining <= 0:
            errors.append(f"Maximum of {max_allowed} images allowed (already at limit)")
        else:
            errors.append(f"Can only add {remaining} more image(s) (limit: {max_allowed})")

    if errors:
        error_items = "".join(f"<li>{e}</li>" for e in errors)
        validation_html = f"""
        <div style="background: #ffebee; border: 1px solid #ef5350; border-radius: 4px; padding: 10px;">
            <ul style="margin: 0; padding-left: 20px; color: #c62828;">
                {error_items}
            </ul>
        </div>
        """
        gallery_data = _get_gallery_data(session)
        count_str = f"{len(session.images)} / {max_allowed}"
        return session, gallery_data, validation_html, count_str, files, description

    # Get room info from session (single room)
    room_id = session.room.id
    room_name = session.room.name.replace(" ", "_")[:20] if session.room.name else "room"

    # Process each uploaded file
    added_count = 0
    for file_obj in files:
        # Check if we've hit the limit
        if len(session.images) >= max_allowed:
            break

        try:
            # Open image from file path
            img = Image.open(file_obj.name)

            # Generate image ID
            image_id = f"img-{uuid.uuid4().hex[:8]}"

            # Store image bytes in memory
            img_bytes = io.BytesIO()
            img.save(img_bytes, format="PNG")
            image_store.store(image_id, img_bytes.getvalue())

            # Add image metadata to session
            img_meta = ImageFormData(
                id=image_id,
                filename=f"{room_name}_{image_id}.png",
                room_id=room_id,
                description=description.strip() if description else "",
            )
            session.images.append(img_meta)
            added_count += 1
        except Exception:
            # Skip files that can't be opened as images
            continue

    session.update_timestamp()

    # Success message
    if added_count > 0:
        validation_html = f"""
        <div style="background: #e8f5e9; border: 1px solid #66bb6a; border-radius: 4px; padding: 10px;">
            <span style="color: #2e7d32;">✓ {added_count} image(s) added for {room_name}</span>
        </div>
        """
    else:
        validation_html = """
        <div style="background: #fff3e0; border: 1px solid #ffb74d; border-radius: 4px; padding: 10px;">
            <span style="color: #e65100;">No images could be processed</span>
        </div>
        """

    gallery_data = _get_gallery_data(session)
    count_str = f"{len(session.images)} / {max_allowed}"

    # Clear form
    return session, gallery_data, validation_html, count_str, None, ""


def remove_last_image(session: SessionState) -> tuple[SessionState, list[tuple], str, str]:
    """Remove the last image from the session."""
    validation_html = ""

    if session.images:
        removed = session.images.pop()
        image_store.remove(removed.id)
        session.update_timestamp()
        validation_html = f"""
        <div style="background: #fff3e0; border: 1px solid #ffb74d; border-radius: 4px; padding: 10px;">
            <span style="color: #e65100;">Removed image: {removed.filename}</span>
        </div>
        """

    gallery_data = _get_gallery_data(session)
    count_str = f"{len(session.images)} / {settings.max_images_per_assessment}"
    return session, gallery_data, validation_html, count_str


def clear_all_images(session: SessionState) -> tuple[SessionState, list[tuple], str, str]:
    """Clear all images from the session."""
    count = len(session.images)

    # Clear from store
    for img in session.images:
        image_store.remove(img.id)

    session.images = []
    session.update_timestamp()

    validation_html = ""
    if count > 0:
        validation_html = f"""
        <div style="background: #fff3e0; border: 1px solid #ffb74d; border-radius: 4px; padding: 10px;">
            <span style="color: #e65100;">Cleared {count} image(s)</span>
        </div>
        """

    count_str = f"0 / {settings.max_images_per_assessment}"
    return session, [], validation_html, count_str


def validate_and_continue(session: SessionState) -> tuple[SessionState, str, int]:
    """Validate Tab 2 and proceed to Tab 3.

    Returns:
        Tuple of (session, validation_html, next_tab_index).
    """
    # Check if images need re-upload (session restored but images not in memory)
    expected_ids = [img.id for img in session.images]
    missing_ids = image_store.get_missing_ids(expected_ids)

    if missing_ids:
        missing_count = len(missing_ids)
        html = f"""
        <div style="background: #fff3e0; border: 1px solid #ffb74d; border-radius: 4px; padding: 10px;">
            <strong style="color: #e65100;">⚠ {missing_count} image(s) need to be re-uploaded</strong>
            <p style="color: #e65100; margin: 5px 0 0 0;">
                Images are not stored in browser storage. Please re-upload the missing images
                or clear the image list and start fresh.
            </p>
        </div>
        """
        return session, html, gr.update(selected=1)  # Stay on Images tab (index 1)

    is_valid, errors = session.validate_tab2()

    if is_valid:
        session.tab2_complete = True
        session.update_timestamp()
        html = """
        <div style="background: #e8f5e9; border: 1px solid #66bb6a; border-radius: 4px; padding: 10px;">
            <span style="color: #2e7d32;">✓ Images complete. Proceeding to Observations tab...</span>
        </div>
        """
        return session, html, gr.update(selected=2)  # Go to tab index 2 (Observations)
    else:
        session.tab2_complete = False
        error_items = "".join(f"<li>{e}</li>" for e in errors)
        html = f"""
        <div style="background: #ffebee; border: 1px solid #ef5350; border-radius: 4px; padding: 10px;">
            <strong style="color: #c62828;">Please fix the following:</strong>
            <ul style="margin: 5px 0 0 0; padding-left: 20px; color: #c62828;">
                {error_items}
            </ul>
        </div>
        """
        return session, html, gr.update(selected=1)  # Stay on current tab


def load_from_session(session: SessionState) -> tuple[list[tuple], str, str]:
    """Load gallery data and count from session.

    Returns:
        Tuple of (gallery_data, image_count, resume_warning_html).
    """
    gallery_data = _get_gallery_data(session)
    count_str = f"{len(session.images)} / {settings.max_images_per_assessment}"

    # Check for missing images
    expected_ids = [img.id for img in session.images]
    missing_ids = image_store.get_missing_ids(expected_ids)

    resume_html = ""
    if missing_ids and session.images:
        resume_html = f"""
        <div style="background: #fff3e0; border: 1px solid #ffb74d; border-radius: 4px; padding: 10px;">
            <strong style="color: #e65100;">⚠ {len(missing_ids)} image(s) need to be re-uploaded</strong>
            <p style="color: #e65100; margin: 5px 0 0 0;">
                Session restored, but images must be re-uploaded as they are not stored in browser storage.
            </p>
        </div>
        """

    return gallery_data, count_str, resume_html


def _get_gallery_data(session: SessionState) -> list[tuple]:
    """Get gallery data from session images.

    Returns:
        List of (image, caption) tuples for gallery.
    """
    gallery_data = []
    for img_meta in session.images:
        img_bytes = image_store.get(img_meta.id)
        if img_bytes:
            # Convert bytes to PIL Image for gallery
            pil_image = Image.open(io.BytesIO(img_bytes))
            caption = img_meta.description or img_meta.filename
            gallery_data.append((pil_image, caption))
    return gallery_data