ui/tabs/input_tab.py

"""Tab 1: Input - Combined Room, Images, and Observations.

Consolidated input tab using Accordions for collapsible sections.
Room and Images are open by default, Observations is collapsed.
"""

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

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


# Facility classification options
FACILITY_OPTIONS = [
    ("Operational", "operational"),
    ("Non-Operational", "non-operational"),
    ("Public/Childcare", "public-childcare"),
]

# Construction era options
CONSTRUCTION_ERA_OPTIONS = [
    ("Pre-1980 (potential LBP/ACM)", "pre-1980"),
    ("1980-2000", "1980-2000"),
    ("Post-2000", "post-2000"),
]

# Odor mapping
ODOR_MAP = {
    "None": "none",
    "Faint": "faint",
    "Moderate": "moderate",
    "Strong": "strong",
}
ODOR_MAP_REVERSE = {v: k for k, v in ODOR_MAP.items()}

# Char density mapping
CHAR_DENSITY_MAP = {
    "None": None,
    "Sparse": "sparse",
    "Moderate": "moderate",
    "Dense": "dense",
}
CHAR_DENSITY_MAP_REVERSE = {v: k for k, v in CHAR_DENSITY_MAP.items()}


def create_tab() -> dict[str, Any]:
    """Create combined Input tab with accordions.

    Returns:
        Dictionary of component references for event wiring.
    """
    # --- Room Details Accordion (OPEN by default) ---
    with gr.Accordion("Room Details", open=True):
        room_name = gr.Textbox(
            label="Room/Area Name *",
            placeholder="e.g., Warehouse Bay A, Office 101",
            elem_id="room_name",
        )

        with gr.Row():
            room_length = gr.Number(
                label="Length (ft) *",
                minimum=1,
                value=None,
                elem_id="room_length",
            )
            room_width = gr.Number(
                label="Width (ft) *",
                minimum=1,
                value=None,
                elem_id="room_width",
            )

        with gr.Row():
            room_height_preset = gr.Dropdown(
                label="Ceiling Height *",
                choices=CEILING_HEIGHT_PRESETS,
                elem_id="room_height_preset",
                info="Select preset or choose Custom",
            )
            room_height_custom = gr.Number(
                label="Custom Height (ft)",
                minimum=1,
                value=None,
                visible=False,
                elem_id="room_height_custom",
            )

        with gr.Row():
            floor_area = gr.Textbox(
                label="Floor Area (SF)",
                value="0",
                interactive=False,
            )
            room_volume = gr.Textbox(
                label="Volume (CF)",
                value="0",
                interactive=False,
            )

        facility_classification = gr.Radio(
            label="Facility Classification *",
            choices=FACILITY_OPTIONS,
            value="non-operational",
            elem_id="facility_classification",
            info="Affects clearance thresholds",
        )

        construction_era = gr.Radio(
            label="Construction Era *",
            choices=CONSTRUCTION_ERA_OPTIONS,
            value="post-2000",
            elem_id="construction_era",
            info="Pre-1980 triggers LBP/ACM flags",
        )

    # --- Images Accordion (OPEN by default) ---
    with gr.Accordion("Images", open=True):
        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", variant="secondary")
                    clear_all_btn = gr.Button("Clear All", variant="stop")

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

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

    # --- Observations Accordion (COLLAPSED by default) ---
    with gr.Accordion("Field Observations (Optional)", open=False):
        gr.Markdown("*Document observations per FDAM §2.3. All fields optional.*")

        with gr.Row():
            with gr.Column():
                smoke_odor = gr.Checkbox(
                    label="Smoke/fire odor present?",
                    elem_id="smoke_odor",
                )
                odor_intensity = gr.Radio(
                    choices=["None", "Faint", "Moderate", "Strong"],
                    label="Odor Intensity",
                    value="None",
                    elem_id="odor_intensity",
                )

                visible_soot = gr.Checkbox(
                    label="Visible soot deposits?",
                    elem_id="visible_soot",
                )
                soot_description = gr.Textbox(
                    label="Soot Pattern Description",
                    placeholder="e.g., Heavy deposits on ceiling",
                    elem_id="soot_description",
                )

                large_char = gr.Checkbox(
                    label="Large char particles observed?",
                    elem_id="large_char",
                )
                char_density = gr.Radio(
                    choices=["None", "Sparse", "Moderate", "Dense"],
                    label="Char Density",
                    value="None",
                    elem_id="char_density",
                )

                ash_residue = gr.Checkbox(
                    label="Ash-like residue present?",
                    elem_id="ash_residue",
                )
                ash_description = gr.Textbox(
                    label="Ash Color/Texture",
                    placeholder="e.g., Gray powdery residue",
                    elem_id="ash_description",
                )

            with gr.Column():
                surface_discoloration = gr.Checkbox(
                    label="Surface discoloration?",
                    elem_id="surface_discoloration",
                )
                discoloration_description = gr.Textbox(
                    label="Discoloration Description",
                    placeholder="e.g., Yellowing on painted surfaces",
                    elem_id="discoloration_description",
                )

                dust_interference = gr.Checkbox(
                    label="Dust loading or interference?",
                    info="Pre-existing dust may affect samples",
                    elem_id="dust_interference",
                )
                dust_notes = gr.Textbox(
                    label="Dust Notes",
                    placeholder="e.g., Heavy ambient dust",
                    elem_id="dust_notes",
                )

                wildfire_indicators = gr.Checkbox(
                    label="Wildfire indicators (vegetation/pollen)?",
                    info="May indicate wildfire vs structural fire",
                    elem_id="wildfire_indicators",
                )
                wildfire_notes = gr.Textbox(
                    label="Wildfire Notes",
                    placeholder="e.g., Burned pine pollen visible",
                    elem_id="wildfire_notes",
                )

                additional_notes = gr.Textbox(
                    label="Additional Observations",
                    lines=3,
                    placeholder="Any other relevant observations...",
                    elem_id="additional_notes",
                )

    # --- Generate Button and Validation ---
    validation_status = gr.HTML(
        value="",
        elem_id="input_validation",
    )

    generate_btn = gr.Button(
        "Generate Assessment →",
        variant="primary",
        size="lg",
    )

    return {
        # Room components
        "room_name": room_name,
        "room_length": room_length,
        "room_width": room_width,
        "room_height_preset": room_height_preset,
        "room_height_custom": room_height_custom,
        "floor_area": floor_area,
        "room_volume": room_volume,
        "facility_classification": facility_classification,
        "construction_era": construction_era,
        # Image components
        "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,
        "resume_warning": resume_warning,
        # Observation components
        "smoke_odor": smoke_odor,
        "odor_intensity": odor_intensity,
        "visible_soot": visible_soot,
        "soot_description": soot_description,
        "large_char": large_char,
        "char_density": char_density,
        "ash_residue": ash_residue,
        "ash_description": ash_description,
        "surface_discoloration": surface_discoloration,
        "discoloration_description": discoloration_description,
        "dust_interference": dust_interference,
        "dust_notes": dust_notes,
        "wildfire_indicators": wildfire_indicators,
        "wildfire_notes": wildfire_notes,
        "additional_notes": additional_notes,
        # Validation and generation
        "validation_status": validation_status,
        "generate_btn": generate_btn,
    }


# --- Room Functions ---


def on_height_preset_change(preset_value: int | None) -> dict:
    """Show/hide custom height input based on preset selection."""
    return gr.update(visible=(preset_value is None))


def update_calculated_values(
    length: float | None,
    width: float | None,
    height_preset: int | None,
    height_custom: float | None,
) -> tuple[str, str]:
    """Calculate and return floor area and volume."""
    length_val = float(length) if length and length > 0 else 0
    width_val = float(width) if width and width > 0 else 0

    if height_preset is not None:
        height_val = float(height_preset)
    elif height_custom is not None and height_custom > 0:
        height_val = float(height_custom)
    else:
        height_val = 0

    area = length_val * width_val
    volume = area * height_val

    return f"{area:,.0f}", f"{volume:,.0f}"


def save_room_to_session(
    session: SessionState,
    name: str,
    length: float | None,
    width: float | None,
    height_preset: int | None,
    height_custom: float | None,
    facility_classification: str,
    construction_era: str,
) -> SessionState:
    """Save room data to session."""
    if height_preset is not None:
        height = float(height_preset)
    elif height_custom is not None and height_custom > 0:
        height = float(height_custom)
    else:
        height = 0

    session.room.name = name.strip() if name else ""
    session.room.length_ft = float(length) if length and length > 0 else 0
    session.room.width_ft = float(width) if width and width > 0 else 0
    session.room.ceiling_height_ft = height
    session.room.facility_classification = facility_classification
    session.room.construction_era = construction_era
    session.update_timestamp()

    return session


def load_room_from_session(
    session: SessionState,
) -> tuple[str, float | None, float | None, int | None, float | None, str, str, str, str]:
    """Load room data from session."""
    r = session.room

    height_preset = None
    height_custom = None
    preset_values = [p[1] for p in CEILING_HEIGHT_PRESETS if p[1] is not None]
    if r.ceiling_height_ft in preset_values:
        height_preset = int(r.ceiling_height_ft)
    elif r.ceiling_height_ft > 0:
        height_custom = r.ceiling_height_ft

    area = r.length_ft * r.width_ft
    volume = area * r.ceiling_height_ft

    return (
        r.name,
        r.length_ft if r.length_ft > 0 else None,
        r.width_ft if r.width_ft > 0 else None,
        height_preset,
        height_custom,
        f"{area:,.0f}",
        f"{volume:,.0f}",
        r.facility_classification,
        r.construction_era,
    )


# --- Image Functions ---


def add_image(
    session: SessionState,
    files: list | None,
    description: str,
) -> tuple[SessionState, list[tuple], str, str, None, str]:
    """Add images to the session."""
    validation_html = ""

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

    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")
        else:
            errors.append(f"Can only add {remaining} more image(s)")

    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

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

    added_count = 0
    for file_obj in files:
        if len(session.images) >= max_allowed:
            break

        try:
            img = Image.open(file_obj.name)
            image_id = f"img-{uuid.uuid4().hex[:8]}"

            img_bytes = io.BytesIO()
            img.save(img_bytes, format="PNG")
            image_store.store(image_id, img_bytes.getvalue())

            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:
            continue

    session.update_timestamp()

    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</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}"
    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: {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)

    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 load_images_from_session(session: SessionState) -> tuple[list[tuple], str, str]:
    """Load gallery data and count from session."""
    gallery_data = _get_gallery_data(session)
    count_str = f"{len(session.images)} / {settings.max_images_per_assessment}"

    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 re-upload</strong>
        </div>
        """

    return gallery_data, count_str, resume_html


def _get_gallery_data(session: SessionState) -> list[tuple]:
    """Get gallery data from session images."""
    gallery_data = []
    for img_meta in session.images:
        img_bytes = image_store.get(img_meta.id)
        if img_bytes:
            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


# --- Observations Functions ---


def save_observations_to_session(
    session: SessionState,
    smoke_odor: bool,
    odor_intensity: str,
    visible_soot: bool,
    soot_description: str,
    large_char: bool,
    char_density: str,
    ash_residue: bool,
    ash_description: str,
    surface_discoloration: bool,
    discoloration_description: str,
    dust_interference: bool,
    dust_notes: str,
    wildfire_indicators: bool,
    wildfire_notes: str,
    additional_notes: str,
) -> SessionState:
    """Update session state from observation form values."""
    session.observations = ObservationsFormData(
        smoke_fire_odor=smoke_odor or False,
        odor_intensity=ODOR_MAP.get(odor_intensity, "none"),
        visible_soot_deposits=visible_soot or False,
        soot_pattern_description=soot_description or "",
        large_char_particles=large_char or False,
        char_density_estimate=CHAR_DENSITY_MAP.get(char_density),
        ash_like_residue=ash_residue or False,
        ash_color_texture=ash_description or "",
        surface_discoloration=surface_discoloration or False,
        discoloration_description=discoloration_description or "",
        dust_loading_interference=dust_interference or False,
        dust_notes=dust_notes or "",
        wildfire_indicators=wildfire_indicators or False,
        wildfire_notes=wildfire_notes or "",
        additional_notes=additional_notes or "",
    )
    session.update_timestamp()
    return session


def load_observations_from_session(session: SessionState) -> tuple:
    """Load observation form values from session state."""
    obs = session.observations
    return (
        obs.smoke_fire_odor,
        ODOR_MAP_REVERSE.get(obs.odor_intensity, "None"),
        obs.visible_soot_deposits,
        obs.soot_pattern_description,
        obs.large_char_particles,
        CHAR_DENSITY_MAP_REVERSE.get(obs.char_density_estimate, "None"),
        obs.ash_like_residue,
        obs.ash_color_texture,
        obs.surface_discoloration,
        obs.discoloration_description,
        obs.dust_loading_interference,
        obs.dust_notes,
        obs.wildfire_indicators,
        obs.wildfire_notes,
        obs.additional_notes,
    )


# --- Validation and Generation ---


def validate_input(session: SessionState) -> tuple[bool, list[str]]:
    """Validate all input sections."""
    errors = []

    # Room validation
    r = session.room
    if not r.name:
        errors.append("Room name is required")
    if r.length_ft <= 0:
        errors.append("Length must be greater than 0")
    if r.width_ft <= 0:
        errors.append("Width must be greater than 0")
    if r.ceiling_height_ft <= 0:
        errors.append("Ceiling height must be greater than 0")

    # Image validation
    if not session.images:
        errors.append("At least one image is required")

    # Check for missing images in memory
    expected_ids = [img.id for img in session.images]
    missing_ids = image_store.get_missing_ids(expected_ids)
    if missing_ids:
        errors.append(f"{len(missing_ids)} image(s) need to be re-uploaded")

    return len(errors) == 0, errors


def validate_and_generate(session: SessionState) -> tuple[SessionState, str, dict]:
    """Validate input and switch to Results tab if valid.

    Returns:
        Tuple of (session, validation_html, tabs_update).
    """
    is_valid, errors = validate_input(session)

    if is_valid:
        session.input_complete = True
        session.update_timestamp()
        html = """
        <div style="background: #e8f5e9; border: 1px solid #66bb6a; border-radius: 4px; padding: 10px;">
            <span style="color: #2e7d32;">✓ All inputs valid. Switching to Results...</span>
        </div>
        """
        return session, html, gr.update(selected=1)  # Go to Results tab (index 1)
    else:
        session.input_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=0)  # Stay on Input tab