ui/samples.py

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

Provides 4 pre-configured sample scenarios with complete project data,
room information, images, and qualitative observations.
"""

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

from PIL import Image

from ui.state import (
    SessionState,
    ProjectFormData,
    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
    project_data: dict
    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",
        project_data={
            "project_name": "Sample: Bar & Dining Fire Assessment",
            "address": "1234 Main Street",
            "city": "Springfield",
            "state": "IL",
            "zip_code": "62701",
            "client_name": "Sample Test Client",
            "fire_date": "2024-11-15",
            "assessment_date": "2024-12-01",
            "facility_classification": "non-operational",
            "construction_era": "pre-1980",
            "assessor_name": "Test Assessor",
            "assessor_credentials": ["CIH"],
        },
        room_data={
            "name": "Bar & Dining Area",
            "floor": "Ground Floor",
            "length_ft": 40.0,
            "width_ft": 30.0,
            "ceiling_height_ft": 12.0,
        },
        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",
        project_data={
            "project_name": "Sample: Bar Area Fire Assessment",
            "address": "1234 Main Street",
            "city": "Springfield",
            "state": "IL",
            "zip_code": "62701",
            "client_name": "Sample Test Client",
            "fire_date": "2024-11-15",
            "assessment_date": "2024-12-01",
            "facility_classification": "non-operational",
            "construction_era": "pre-1980",
            "assessor_name": "Test Assessor",
            "assessor_credentials": ["CIH"],
        },
        room_data={
            "name": "Bar Area",
            "floor": "Ground Floor",
            "length_ft": 25.0,
            "width_ft": 20.0,
            "ceiling_height_ft": 14.0,
        },
        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",
        project_data={
            "project_name": "Sample: Kitchen Fire Assessment",
            "address": "5678 Industrial Blvd",
            "city": "Chicago",
            "state": "IL",
            "zip_code": "60601",
            "client_name": "Sample Test Client",
            "fire_date": "2024-10-20",
            "assessment_date": "2024-11-05",
            "facility_classification": "non-operational",
            "construction_era": "1980-2000",
            "assessor_name": "Test Assessor",
            "assessor_credentials": ["CIH", "CSP"],
        },
        room_data={
            "name": "Commercial Kitchen",
            "floor": "Ground Floor",
            "length_ft": 30.0,
            "width_ft": 25.0,
            "ceiling_height_ft": 10.0,
        },
        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",
        project_data={
            "project_name": "Sample: Factory Fire Assessment",
            "address": "9999 Factory Way",
            "city": "Detroit",
            "state": "MI",
            "zip_code": "48201",
            "client_name": "Industrial Test Corp",
            "fire_date": "2024-09-01",
            "assessment_date": "2024-09-15",
            "facility_classification": "operational",
            "construction_era": "pre-1980",
            "assessor_name": "Test Assessor",
            "assessor_credentials": ["CIH", "PE"],
        },
        room_data={
            "name": "Factory Production Area",
            "floor": "Ground Floor",
            "length_ft": 80.0,
            "width_ft": 60.0,
            "ceiling_height_ft": 25.0,
        },
        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
    room_id = f"room-{uuid.uuid4().hex[:8]}"
    room = RoomFormData(
        id=room_id,
        name=scenario.room_data["name"],
        floor=scenario.room_data.get("floor", ""),
        length_ft=scenario.room_data["length_ft"],
        width_ft=scenario.room_data["width_ft"],
        ceiling_height_ft=scenario.room_data["ceiling_height_ft"],
    )

    # Load images
    images = load_sample_images(scenario, room_id)

    # Create session
    session = SessionState(
        project=ProjectFormData(**scenario.project_data),
        rooms=[room],
        images=images,
        observations=ObservationsFormData(**scenario.observations_data),
        name=scenario.project_data["project_name"],
    )

    # Mark tabs as complete since we have all data
    session.tab1_complete = True
    session.tab2_complete = True
    session.tab3_complete = len(images) > 0
    session.tab4_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)