ui/samples.py

"""Sample room data for testing the FDAM AI Pipeline.

Provides 4 pre-configured sample scenarios with complete room data,
images, and qualitative observations.
MVP Simplification: Single room, no project-level data.
"""

import uuid
import io
from pathlib import Path
from dataclasses import dataclass, field

from PIL import Image

from ui.state import (
    SessionState,
    RoomFormData,
    ImageFormData,
    ObservationsFormData,
)
from ui.components import image_store


# Path to sample images directory
SAMPLE_IMAGES_DIR = Path(__file__).parent.parent / "sample_images"


@dataclass
class SampleScenario:
    """Definition of a sample fire damage scenario."""

    id: str
    name: str
    description: str
    room_data: dict
    observations_data: dict
    image_files: list[str] = field(default_factory=list)


# --- Sample Scenario Definitions ---

SAMPLE_SCENARIOS = [
    # 1. Bar & Dining Area
    SampleScenario(
        id="bar_dining",
        name="Bar & Dining Area",
        description="3 images",
        room_data={
            "name": "Bar & Dining Area",
            "length_ft": 40.0,
            "width_ft": 30.0,
            "ceiling_height_ft": 12.0,
            "facility_classification": "non-operational",
            "construction_era": "pre-1980",
        },
        observations_data={
            "smoke_fire_odor": True,
            "odor_intensity": "strong",
            "visible_soot_deposits": True,
            "soot_pattern_description": "Heavy soot deposits on corrugated metal ceiling, moderate wall discoloration",
            "large_char_particles": True,
            "char_density_estimate": "moderate",
            "ash_like_residue": True,
            "ash_color_texture": "Ash deposits on horizontal surfaces and upholstered furniture",
            "surface_discoloration": True,
            "discoloration_description": "Tan/brown soot staining on walls, yellowing on decorative elements",
            "dust_loading_interference": False,
            "dust_notes": "",
            "wildfire_indicators": False,
            "wildfire_notes": "",
            "additional_notes": "",
        },
        image_files=[
            "Bar and dining area1.jpg",
            "Bar and dining area2.jpg",
            "Bar and dining area3.jpg",
        ],
    ),
    # 2. Bar Area
    SampleScenario(
        id="bar_area",
        name="Bar Area",
        description="3 images",
        room_data={
            "name": "Bar Area",
            "length_ft": 25.0,
            "width_ft": 20.0,
            "ceiling_height_ft": 14.0,
            "facility_classification": "non-operational",
            "construction_era": "pre-1980",
        },
        observations_data={
            "smoke_fire_odor": True,
            "odor_intensity": "strong",
            "visible_soot_deposits": True,
            "soot_pattern_description": "Dense black coating on ceiling/ductwork, severe overhead damage",
            "large_char_particles": True,
            "char_density_estimate": "dense",
            "ash_like_residue": True,
            "ash_color_texture": "Heavy ash on shelving and bottled goods",
            "surface_discoloration": True,
            "discoloration_description": "Metal oxidation, melted plastic signage, deformed ductwork",
            "dust_loading_interference": False,
            "dust_notes": "",
            "wildfire_indicators": False,
            "wildfire_notes": "",
            "additional_notes": "",
        },
        image_files=[
            "Bar area1.jpg",
            "Bar area2.jpg",
            "Bar area3.jpg",
        ],
    ),
    # 3. Kitchen
    SampleScenario(
        id="kitchen",
        name="Kitchen",
        description="6 images",
        room_data={
            "name": "Commercial Kitchen",
            "length_ft": 30.0,
            "width_ft": 25.0,
            "ceiling_height_ft": 10.0,
            "facility_classification": "non-operational",
            "construction_era": "1980-2000",
        },
        observations_data={
            "smoke_fire_odor": True,
            "odor_intensity": "strong",
            "visible_soot_deposits": True,
            "soot_pattern_description": "Heavy soot on all surfaces, ceiling collapse debris",
            "large_char_particles": True,
            "char_density_estimate": "dense",
            "ash_like_residue": True,
            "ash_color_texture": "Thick ash deposits on work surfaces, equipment heavily coated",
            "surface_discoloration": True,
            "discoloration_description": "Charred drywall, oxidized metal equipment, concrete staining",
            "dust_loading_interference": False,
            "dust_notes": "",
            "wildfire_indicators": False,
            "wildfire_notes": "",
            "additional_notes": "",
        },
        image_files=[
            "Kitchen 1.jpg",
            "Kitchen 2.jpg",
            "Kitchen 3.jpg",
            "Kitchen 4.jpg",
            "Kitchen 5.jpg",
            "Kitchen 6.jpg",
        ],
    ),
    # 4. Factory Area
    SampleScenario(
        id="factory",
        name="Factory Area",
        description="1 image",
        room_data={
            "name": "Factory Production Area",
            "length_ft": 80.0,
            "width_ft": 60.0,
            "ceiling_height_ft": 25.0,
            "facility_classification": "operational",
            "construction_era": "pre-1980",
        },
        observations_data={
            "smoke_fire_odor": True,
            "odor_intensity": "strong",
            "visible_soot_deposits": True,
            "soot_pattern_description": "Complete structural compromise, deep char on all surfaces",
            "large_char_particles": True,
            "char_density_estimate": "dense",
            "ash_like_residue": True,
            "ash_color_texture": "Heavy ash coating throughout, debris accumulation",
            "surface_discoloration": True,
            "discoloration_description": "Extreme oxidation on metal framing, thermal spalling on concrete",
            "dust_loading_interference": False,
            "dust_notes": "",
            "wildfire_indicators": False,
            "wildfire_notes": "",
            "additional_notes": "",
        },
        image_files=[
            "factory_area.jpg",
        ],
    ),
]

# Create lookup dict for fast access
SAMPLE_SCENARIOS_BY_ID = {s.id: s for s in SAMPLE_SCENARIOS}


def get_sample_choices() -> list[tuple[str, str]]:
    """Get dropdown choices for sample selector.

    Returns:
        List of (label, value) tuples for Gradio dropdown.
    """
    choices = [("Select a sample scenario...", "")]
    for scenario in SAMPLE_SCENARIOS:
        label = f"{scenario.name} ({scenario.description})"
        choices.append((label, scenario.id))
    return choices


def load_sample_images(scenario: SampleScenario, room_id: str) -> list[ImageFormData]:
    """Load sample images from disk into image_store.

    Args:
        scenario: The sample scenario to load images for.
        room_id: The room ID to associate images with.

    Returns:
        List of ImageFormData objects for the loaded images.
    """
    image_metas = []

    for filename in scenario.image_files:
        filepath = SAMPLE_IMAGES_DIR / filename
        if filepath.exists():
            try:
                # Read and convert image to PNG bytes
                img = Image.open(filepath)
                img_bytes = io.BytesIO()
                img.save(img_bytes, format="PNG")

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

                # Store in image_store
                image_store.store(image_id, img_bytes.getvalue())

                # Create metadata
                image_metas.append(
                    ImageFormData(
                        id=image_id,
                        filename=filename,
                        room_id=room_id,
                        description=f"Sample image: {filename}",
                    )
                )
            except Exception:
                # Skip files that can't be opened as images
                continue

    return image_metas


def load_sample(scenario_id: str) -> SessionState | None:
    """Load a sample scenario into a new SessionState.

    Args:
        scenario_id: The ID of the scenario to load.

    Returns:
        A new SessionState populated with the scenario data, or None if not found.
    """
    scenario = SAMPLE_SCENARIOS_BY_ID.get(scenario_id)
    if not scenario:
        return None

    # Create room with unique ID and all fields from room_data
    room_id = f"room-{uuid.uuid4().hex[:8]}"
    room = RoomFormData(
        id=room_id,
        name=scenario.room_data["name"],
        length_ft=scenario.room_data["length_ft"],
        width_ft=scenario.room_data["width_ft"],
        ceiling_height_ft=scenario.room_data["ceiling_height_ft"],
        facility_classification=scenario.room_data.get("facility_classification", "non-operational"),
        construction_era=scenario.room_data.get("construction_era", "post-2000"),
    )

    # Load images
    images = load_sample_images(scenario, room_id)

    # Create session with single room
    session = SessionState(
        room=room,
        images=images,
        observations=ObservationsFormData(**scenario.observations_data),
        name=scenario.room_data["name"],
    )

    # Mark input as complete since sample has all required data
    session.input_complete = True

    return session


def get_scenario_by_id(scenario_id: str) -> SampleScenario | None:
    """Get a sample scenario by its ID.

    Args:
        scenario_id: The scenario ID.

    Returns:
        The SampleScenario object or None if not found.
    """
    return SAMPLE_SCENARIOS_BY_ID.get(scenario_id)