ui/state.py

"""Session state management for FDAM AI Pipeline.

Provides Pydantic models for session state and localStorage persistence.
Images are stored separately (not in localStorage due to size limits).

MVP Simplification: Single room assessment (no project-level fields).
"""

import json
import uuid
from datetime import datetime
from typing import Optional

from pydantic import BaseModel, Field

from schemas.input import (
    ConstructionEra,
    FacilityClassification,
    OdorIntensity,
    CharDensity,
)


# --- Form Data Models (for localStorage) ---


class RoomFormData(BaseModel):
    """Form data for Tab 1: Room Assessment.

    Includes facility_classification and construction_era moved from
    the removed ProjectFormData (required for calculations).
    """

    id: str = Field(default_factory=lambda: f"room-{uuid.uuid4().hex[:8]}")
    name: str = ""
    length_ft: float = 0
    width_ft: float = 0
    ceiling_height_ft: float = 0
    # Moved from ProjectFormData (required for calculations)
    facility_classification: FacilityClassification = "non-operational"
    construction_era: ConstructionEra = "post-2000"


class ImageFormData(BaseModel):
    """Form data for a single image (metadata only, not bytes)."""

    id: str = Field(default_factory=lambda: f"img-{uuid.uuid4().hex[:8]}")
    filename: str = ""
    room_id: str = ""
    description: str = ""
    # Image bytes stored separately, referenced by id


class ObservationsFormData(BaseModel):
    """Form data for Tab 3: Observations."""

    smoke_fire_odor: bool = False
    odor_intensity: OdorIntensity = "none"
    visible_soot_deposits: bool = False
    soot_pattern_description: str = ""
    large_char_particles: bool = False
    char_density_estimate: Optional[CharDensity] = None
    ash_like_residue: bool = False
    ash_color_texture: str = ""
    surface_discoloration: bool = False
    discoloration_description: str = ""
    dust_loading_interference: bool = False
    dust_notes: str = ""
    wildfire_indicators: bool = False
    wildfire_notes: str = ""
    additional_notes: str = ""


class SessionState(BaseModel):
    """Complete session state for an assessment.

    This model is serialized to localStorage for persistence.
    Images are stored separately and referenced by ID.

    MVP Simplification: Single room, 2 tabs (Input + Results/Chat).
    """

    # Session metadata
    session_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    name: str = ""  # Display name for history list

    # Input completion status (single flag replaces 3 tab flags)
    input_complete: bool = False

    # Form data - single room (not list)
    room: RoomFormData = Field(default_factory=RoomFormData)
    images: list[ImageFormData] = Field(default_factory=list)
    observations: ObservationsFormData = Field(default_factory=ObservationsFormData)

    # Results (after generation)
    has_results: bool = False
    results_generated_at: Optional[str] = None

    # Chat history (Gradio 6 messages format)
    chat_history: list[dict] = Field(default_factory=list)

    # Serializable subset of PipelineResult (excludes PIL images)
    pipeline_result_json: Optional[str] = None

    # Document state for modifications
    generated_document: Optional[str] = None
    original_document: Optional[str] = None

    def update_timestamp(self) -> None:
        """Update the updated_at timestamp."""
        self.updated_at = datetime.now().isoformat()

    def get_display_name(self) -> str:
        """Get a display name for the history list."""
        if self.name:
            return self.name
        if self.room.name:
            return self.room.name
        return f"Assessment {self.session_id[:8]}"

    def validate_tab1(self) -> tuple[bool, list[str]]:
        """Validate Tab 1 (Room Assessment) is complete."""
        errors = []
        r = self.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")
        return len(errors) == 0, errors

    def validate_tab2(self) -> tuple[bool, list[str]]:
        """Validate Tab 2 (Images) is complete."""
        errors = []
        if not self.images:
            errors.append("At least one image is required")
        for img in self.images:
            if not img.room_id:
                errors.append(f"Image '{img.filename}': Must be associated with the room")
        return len(errors) == 0, errors

    def validate_tab3(self) -> tuple[bool, list[str]]:
        """Validate Tab 3 (Observations) is complete."""
        # Tab 3 has no required fields - all checkboxes default to False
        return True, []

    def can_generate(self) -> tuple[bool, list[str]]:
        """Check if assessment can be generated."""
        all_errors = []

        valid1, errors1 = self.validate_tab1()
        if not valid1:
            all_errors.extend(errors1)

        valid2, errors2 = self.validate_tab2()
        if not valid2:
            all_errors.extend(errors2)

        valid3, errors3 = self.validate_tab3()
        if not valid3:
            all_errors.extend(errors3)

        return len(all_errors) == 0, all_errors


class AssessmentHistory(BaseModel):
    """Collection of saved assessments for history list."""

    assessments: list[SessionState] = Field(default_factory=list)
    current_session_id: Optional[str] = None

    def add_assessment(self, session: SessionState) -> None:
        """Add or update an assessment in history."""
        session.update_timestamp()
        # Remove existing if present
        self.assessments = [a for a in self.assessments if a.session_id != session.session_id]
        # Add to front of list
        self.assessments.insert(0, session)
        # Keep only last 20 assessments
        self.assessments = self.assessments[:20]

    def get_assessment(self, session_id: str) -> Optional[SessionState]:
        """Get an assessment by ID."""
        for a in self.assessments:
            if a.session_id == session_id:
                return a
        return None

    def remove_assessment(self, session_id: str) -> None:
        """Remove an assessment from history."""
        self.assessments = [a for a in self.assessments if a.session_id != session_id]

    def get_history_items(self) -> list[dict]:
        """Get history items for display in dropdown."""
        return [
            {
                "id": a.session_id,
                "name": a.get_display_name(),
                "updated": a.updated_at,
                "has_results": a.has_results,
            }
            for a in self.assessments
        ]


# --- Gradio State Helpers ---


def create_new_session() -> SessionState:
    """Create a new empty session."""
    return SessionState()


def session_to_json(session: SessionState) -> str:
    """Serialize session to JSON for localStorage."""
    return session.model_dump_json()


def session_from_json(json_str: str) -> SessionState:
    """Deserialize session from JSON.

    Includes migration from old multi-room format to single room.
    """
    try:
        data = json.loads(json_str)

        # Migration: Convert old multi-room format to single room
        if "rooms" in data and isinstance(data["rooms"], list):
            # Use first room if available, otherwise create empty
            if data["rooms"]:
                data["room"] = data["rooms"][0]
            else:
                data["room"] = {}
            del data["rooms"]

        # Migration: Move facility_classification and construction_era from project to room
        if "project" in data:
            project = data["project"]
            if "room" not in data:
                data["room"] = {}
            # Move fields if they exist in project
            if "facility_classification" in project:
                data["room"]["facility_classification"] = project["facility_classification"]
            if "construction_era" in project:
                data["room"]["construction_era"] = project["construction_era"]
            del data["project"]

        # Migration: Remove old tab4_complete (now only 3 tabs)
        if "tab4_complete" in data:
            del data["tab4_complete"]

        # Migration: Convert old tab1/2/3_complete to input_complete
        if "tab1_complete" in data or "tab2_complete" in data or "tab3_complete" in data:
            # Input is complete if all three old tabs were complete
            tab1 = data.pop("tab1_complete", False)
            tab2 = data.pop("tab2_complete", False)
            tab3 = data.pop("tab3_complete", False)
            data["input_complete"] = tab1 and tab2 and tab3

        return SessionState.model_validate(data)
    except Exception:
        return create_new_session()


def history_to_json(history: AssessmentHistory) -> str:
    """Serialize history to JSON for localStorage."""
    return history.model_dump_json()


def history_from_json(json_str: str) -> AssessmentHistory:
    """Deserialize history from JSON."""
    try:
        return AssessmentHistory.model_validate_json(json_str)
    except Exception:
        return AssessmentHistory()