app.py

"""
WOD Analyzer - Clean and Refactored Version
A Gradio application for analyzing Work Order Documents with improved code structure.
"""

import os
import json
import time
from typing import Tuple, Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum

import gradio as gr
import pandas as pd
from python_request import process_wod_document
from dummy import output_test

PRODUCTION = True

# === CONSTANTS ===
class WODType(Enum):
    """Enum for Work Order Document types."""
    REPLACEMENT = "REPLACEMENT"
    THERMAL = "THERMAL"
    VISIT = "VISIT"
    PREVENTIVE_MAINTENANCE = "PREVENTIVE_MAINTENANCE"
    INSTALLATION = "INSTALLATION"
    WITHDRAWAL = "WITHDRAWAL"


class UIConstants:
    """UI-related constants."""
    DEFAULT_WOD_TYPE = "-- WOD type --"
    BUTTON_ANALYZE = "Analyze Document"
    BUTTON_PROCESSING = "Processing..."
    BUTTON_RESET = "Reset"
    
    # Styling
    MAX_WIDTH = "960px"
    TABLE_MAX_HEIGHT = 1250
    COLUMN_WIDTHS = [30, 60, 10]
    
    # Messages
    TITLE = "# WOD Analyzer"
    DESCRIPTION = "Upload a Work Order Document to automatically check for requirements."
    NO_WOD_TYPE_WARNING = "Please select a WOD type first!"
    NO_FILE_WARNING = "Please upload a PDF file first!"


class Config:
    """Application configuration."""
    USERNAME = "demo"
    PASSWORD_ENV_VAR = os.environ["PASSWORD"]
    DEBUG = True
    MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB


@dataclass
class AnalysisResult:
    """Data class for analysis results."""
    prediction: str
    dataframe: pd.DataFrame
    json_data: Dict[str, Any]
    accordion_visible: bool
    success: bool
    error_message: Optional[str] = None


# === CORE BUSINESS LOGIC ===
class WODAnalyzer:
    """Core business logic for WOD analysis."""
    
    @staticmethod
    def validate_inputs(file_obj: Optional[Any], wod_type: str) -> Tuple[bool, str]:
        """Validate user inputs."""
        if wod_type == UIConstants.DEFAULT_WOD_TYPE or wod_type is None:
            return False, UIConstants.NO_WOD_TYPE_WARNING
        
        if file_obj is None:
            return False, UIConstants.NO_FILE_WARNING
        
        return True, ""
    
    @staticmethod
    def get_file_path(file_obj: Any) -> str:
        """Extract file path from Gradio file object."""
        if hasattr(file_obj, 'name') and os.path.isfile(file_obj.name):
            return file_obj.name
        return str(file_obj)
    
    @staticmethod
    def process_api_response(api_response: Dict[str, Any]) -> AnalysisResult:
        """Process API response and convert to AnalysisResult."""
        if api_response.get("status") != "success":
            error_msg = api_response.get("message", "Unknown error occurred")
            return AnalysisResult(
                prediction="",
                dataframe=pd.DataFrame(),
                json_data={},
                accordion_visible=False,
                success=False,
                error_message=f"API Error: {error_msg}"
            )
        
        # Parse the API response
        results = api_response.get("results", {})
        summary = results.get("summary", {})
        extracted_data = api_response.get("extracted_data", {})
        
        # Convert summary to DataFrame
        df = WODAnalyzer._create_summary_dataframe(summary)
        
        # Get prediction
        prediction = results.get("prediction", "Unknown")
        
        return AnalysisResult(
            prediction=prediction,
            dataframe=df,
            json_data=extracted_data,
            accordion_visible=bool(extracted_data),
            success=True
        )
    
    @staticmethod
    def _create_summary_dataframe(summary: Dict[str, Any]) -> pd.DataFrame:
        """Create DataFrame from summary data."""
        requirements = []
        reasons = []
        statuses = []
        
        for requirement_name, details in summary.items():
            requirements.append(requirement_name)
            reasons.append(details.get("reasoning", ""))
            
            # Convert true/false to PASS/FAIL
            status_bool = details.get("status", "false")
            if isinstance(status_bool, str):
                status = "PASS" if status_bool.lower() == "true" else "FAIL"
            else:
                status = "PASS" if status_bool else "FAIL"
            statuses.append(status)
        
        return pd.DataFrame({
            "Requirement": requirements,
            "Reason": reasons,
            "Status": statuses
        })
    
    @classmethod
    def analyze_document(cls, file_obj: Any, wod_type: str) -> AnalysisResult:
        """Main analysis function."""
        # Validate inputs
        is_valid, error_msg = cls.validate_inputs(file_obj, wod_type)
        if not is_valid:
            return AnalysisResult(
                prediction="",
                dataframe=pd.DataFrame(),
                json_data={},
                accordion_visible=False,
                success=False,
                error_message=error_msg
            )
        
        try:
            print(f"Analyzing '{file_obj.name}' (Type: {wod_type})...")
            
            # Get file path
            file_path = cls.get_file_path(file_obj)

            if PRODUCTION:
                api_response = process_wod_document(file_path, wod_type)
            else:
                time.sleep(1)  # Simulate processing time
                api_response = json.loads(output_test)
            
            # Process the response
            return cls.process_api_response(api_response)
            
        except Exception as e:
            error_msg = f"Error processing document: {str(e)}"
            print(error_msg)
            return AnalysisResult(
                prediction="",
                dataframe=pd.DataFrame(),
                json_data={},
                accordion_visible=False,
                success=False,
                error_message=error_msg
            )


# === AUTHENTICATION ===
class AuthManager:
    """Handle user authentication."""
    
    @staticmethod
    def authenticate_user(username: str, password: str) -> bool:
        """
        Simple authentication function.
        In production, use more secure methods like hashed passwords.
        """
        expected_password = Config.PASSWORD_ENV_VAR
        if not expected_password:
            print("Warning: PASSWORD environment variable not set")
            return False
        
        return username == Config.USERNAME and password == expected_password


# === UI COMPONENTS ===
class UIBuilder:
    """Builds and manages UI components."""
    
    @staticmethod
    def get_wod_type_options() -> List[str]:
        """Get WOD type dropdown options."""
        return [UIConstants.DEFAULT_WOD_TYPE] + [wod_type.value for wod_type in WODType]
    
    @staticmethod
    def create_custom_css() -> str:
        """Create custom CSS for the application."""
        return f"""
        .gradio-container {{
            max-width: {UIConstants.MAX_WIDTH} !important;
            margin: auto !important;
        }}
        .progress-text {{
            display: none !important;
        }}
        """
    
    @staticmethod
    def create_theme() -> gr.Theme:
        """Create custom theme for the application."""
        return gr.themes.Default(primary_hue="blue", secondary_hue="sky")
    
    @staticmethod
    def format_prediction_display(prediction: str) -> str:
        """Format prediction for display."""
        if prediction and prediction != "Unknown":
            return f"<h1 style='text-align: center; color: #1f77b4;'>{prediction}</h1>"
        return ""


# === EVENT HANDLERS ===
class EventHandlers:
    """Handle UI events and interactions."""
    
    @staticmethod
    def handle_analyze_button(
        file_obj: Any, 
        wod_type: str, 
        current_button_value: str
    ) -> Tuple[str, pd.DataFrame, Dict[str, Any], gr.update, gr.update]:
        """Handle analyze button click."""
        # Check if this is a reset action
        if current_button_value == UIConstants.BUTTON_RESET:
            return (
                "",
                pd.DataFrame(),
                {},
                gr.update(visible=False),
                gr.update(value=UIConstants.BUTTON_ANALYZE, interactive=True)
            )
        
        # Perform analysis
        result = WODAnalyzer.analyze_document(file_obj, wod_type)
        
        if not result.success:
            if result.error_message:
                if "Please select" in result.error_message or "Please upload" in result.error_message:
                    gr.Warning(result.error_message)
                else:
                    gr.Error(result.error_message)
            
            return (
                "",
                pd.DataFrame(),
                {},
                gr.update(visible=False),
                gr.update(value=UIConstants.BUTTON_ANALYZE, interactive=True)
            )
        
        # Success case
        gr.Info(f"Analysis completed! Overall prediction: {result.prediction}")
        prediction_display = UIBuilder.format_prediction_display(result.prediction)
        
        return (
            prediction_display,
            result.dataframe,
            result.json_data,
            gr.update(visible=result.accordion_visible),
            gr.update(value=UIConstants.BUTTON_RESET, interactive=True)
        )
    
    @staticmethod
    def set_processing_state(current_button_value: str) -> gr.update:
        """Set button to processing state."""
        if current_button_value == UIConstants.BUTTON_RESET:
            return gr.update()  # Don't change if it's reset
        return gr.update(value=UIConstants.BUTTON_PROCESSING, interactive=False)


# === JAVASCRIPT FUNCTIONS ===
class JavaScriptFunctions:
    """JavaScript functions for the UI."""
    
    THEME_SETUP = """
    function refresh() {
        const url = new URL(window.location);
        if (url.searchParams.get('__theme') !== 'light') {
            url.searchParams.set('__theme', 'light');
            window.location.href = url.href;
        }
    }
    """
    
    REFRESH_ON_RESET = """
    function(button_value) {
        if (button_value === "Reset") {
            window.location.reload();
            return false;
        }
        return true;
    }
    """


# === MAIN APPLICATION ===
class WODAnalyzerApp:
    """Main application class."""
    
    def __init__(self):
        self.ui_builder = UIBuilder()
        self.event_handlers = EventHandlers()
    
    def create_interface(self) -> gr.Blocks:
        """Create the Gradio interface."""
        with gr.Blocks(
            theme=self.ui_builder.create_theme(),
            js=JavaScriptFunctions.THEME_SETUP,
            css=self.ui_builder.create_custom_css()
        ) as demo:
            
            # Header
            gr.Markdown(f"{UIConstants.TITLE}\n{UIConstants.DESCRIPTION}")
            
            # Input Section
            with gr.Row():
                file_input = gr.File(label="Upload WOD PDF")
                type_input = gr.Dropdown(
                    choices=self.ui_builder.get_wod_type_options(),
                    label="Type",
                    value=UIConstants.DEFAULT_WOD_TYPE,
                    info="Select the type of work order."
                )
            
            # Action Button
            analyze_btn = gr.Button(UIConstants.BUTTON_ANALYZE, variant="primary")
            
            # Results Section
            gr.Markdown("---\n## Results")
            
            # Prediction display
            prediction_output = gr.Markdown(value="", visible=True)
            
            # JSON display for extracted data
            with gr.Accordion("Extraction Result from Page 1", open=False, visible=False) as json_accordion:
                json_output = gr.JSON(label="Extracted Data", show_label=False, open=True)
            
            # Results table
            results_output = gr.DataFrame(
                headers=["Requirement", "Reason", "Status"],
                datatype=["str", "str", "str"],
                interactive=False,
                max_height=UIConstants.TABLE_MAX_HEIGHT,
                column_widths=UIConstants.COLUMN_WIDTHS,
                wrap=True
            )
            
            # Event handling
            analyze_btn.click(
                fn=self.event_handlers.set_processing_state,
                inputs=[analyze_btn],
                outputs=[analyze_btn],
                show_progress=False,
                js=JavaScriptFunctions.REFRESH_ON_RESET
            ).then(
                fn=self.event_handlers.handle_analyze_button,
                inputs=[file_input, type_input, analyze_btn],
                outputs=[prediction_output, results_output, json_output, json_accordion, analyze_btn],
                show_progress=True
            )
        
        return demo
    
    def launch(self, enable_auth: bool = False) -> None:
        """Launch the application."""
        demo = self.create_interface()
        
        if enable_auth:
            demo.launch(
                auth=AuthManager.authenticate_user,
                auth_message="Please enter your credentials to access the WOD Analyzer",
                debug=Config.DEBUG,
                ssr_mode=False
            )
        else:
            demo.launch(debug=Config.DEBUG)


# === MAIN ENTRY POINT ===
def main():
    """Main entry point."""
    app = WODAnalyzerApp()
    
    app.launch(enable_auth=PRODUCTION)


if __name__ == "__main__":
    main()