app.py

import os
import logging
from typing import Dict, List, Optional
from functools import lru_cache
import re

import gradio as gr

try:
    # Assuming vector_db.py exists in the same directory or is installed
    from vector_db import VectorDatabase
except ImportError:
    print("Error: Could not import VectorDatabase from vector_db.py.")
    print("Please ensure vector_db.py exists in the same directory and is correctly defined.")
    # Exit if critical dependency is missing at import time
    exit(1)

try:
    from langchain_openai import ChatOpenAI
except ImportError:
    print("Error: langchain-openai not found. Please install it: pip install langchain-openai")
    # Exit if critical dependency is missing at import time
    exit(1)

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Suppress warnings for cleaner console output
import warnings
warnings.filterwarnings("ignore", category=SyntaxWarning)
warnings.filterwarnings("ignore", category=UserWarning, message=".*You are using gradio version.*")
warnings.filterwarnings("ignore", category=DeprecationWarning)

# Enhanced logging configuration
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s'
)

# --- RAGSystem Class (Processing Logic - KEPT INTACT AS REQUESTED) ---
class RAGSystem:
    def __init__(self, vector_db: Optional[VectorDatabase] = None):
        logging.info("Initializing RAGSystem")
        self.vector_db = vector_db if vector_db else VectorDatabase()
        self.llm = None
        self.chain = None
        self.prompt_template_str = """You are a legal assistant specializing in tenant rights and landlord-tenant laws. Your goal is to provide accurate, detailed, and helpful answers grounded in legal authority. Use the provided statutes as the primary source when available. If no relevant statutes are found in the context, rely on your general knowledge to provide a pertinent and practical response, clearly indicating when you are doing so and prioritizing state-specific information over federal laws for state-specific queries.
Instructions:
* Use the context and statutes as the primary basis for your answer when available.
* For state-specific queries, prioritize statutes or legal principles from the specified state over federal laws.
* Cite relevant statutes (e.g., (AS § 34.03.220(a)(2))) explicitly in your answer when applicable.
* If multiple statutes apply, list all relevant ones.
* If no specific statute is found in the context, state this clearly (e.g., 'No specific statute was found in the provided context'), then provide a general answer based on common legal principles or practices, marked as such.
* Include practical examples or scenarios to enhance clarity and usefulness.
* Use bullet points or numbered lists for readability when appropriate.
* Maintain a professional and neutral tone.
Question: {query}
State: {state}
Statutes from context:
{statutes}
Context information:
--- START CONTEXT ---
{context}
--- END CONCONTEXT ---
Answer:"""
        self.prompt_template = PromptTemplate(
            input_variables=["query", "context", "state", "statutes"],
            template=self.prompt_template_str
        )
        logging.info("RAGSystem initialized.")

    def extract_statutes(self, text: str) -> str:
        statute_pattern = r'\b(?:[A-Z]{2,}\.?\s+(?:Rev\.\s+)?Stat\.?|Code(?:\s+Ann\.?)?|Ann\.?\s+Laws|Statutes|CCP|USC|ILCS|Civ\.\s+Code|Penal\s+Code|Gen\.\s+Oblig\.\s+Law|R\.?S\.?|P\.?L\.?)\s+§\s*[\d\-]+(?:\.\d+)?(?:[\(\w\.\)]+)?|Title\s+\d+\s+USC\s+§\s*\d+(?:-\d+)?\b'
        statutes = re.findall(statute_pattern, text, re.IGNORECASE)
        valid_statutes = []
        for statute in statutes:
            statute = statute.strip()
            if '§' in statute and any(char.isdigit() for char in statute):
                if not re.match(r'^\([\w\.]+\)$', statute) and 'http' not in statute:
                    if len(statute) > 5:
                        valid_statutes.append(statute)

        if valid_statutes:
            seen = set()
            unique_statutes = [s for s in valid_statutes if not (s.rstrip('.,;') in seen or seen.add(s.rstrip('.,;')))]
            logging.info(f"Extracted {len(unique_statutes)} unique statutes.")
            return "\n".join(f"- {s}" for s in unique_statutes)

        logging.info("No statutes found matching the pattern in the context.")
        return "No specific statutes found in the provided context."

    @lru_cache(maxsize=50)
    def process_query_cached(self, query: str, state: str, openai_api_key: str, n_results: int = 5) -> Dict[str, any]:
        logging.info(f"Processing query (cache key: '{query}'|'{state}'|key_hidden) with n_results={n_results}")

        if not state or state is None: # Removed "Select a state..." from error check as it's not a choice
            logging.warning("No valid state provided for query.")
            return {"answer": "<div class='error-message'>Error: Please select a valid state.</div>", "context_used": "N/A - Invalid Input"}
        if not query or not query.strip():
            logging.warning("No query provided.")
            return {"answer": "<div class='error-message'>Error: Please enter your question.</div>", "context_used": "N/A - Invalid Input"}
        if not openai_api_key or not openai_api_key.strip() or not openai_api_key.startswith("sk-"):
            logging.warning("No valid OpenAI API key provided.")
            return {"answer": "<div class='error-message'>Error: Please provide a valid OpenAI API key (starting with 'sk-'). Get one from <a href='https://platform.openai.com/api-keys' target='_blank'>OpenAI</a>.</div>", "context_used": "N/A - Invalid Input"}

        try:
            logging.info("Initializing temporary LLM and Chain for this query...")
            temp_llm = ChatOpenAI(
                temperature=0.2, openai_api_key=openai_api_key, model_name="gpt-3.5-turbo",
                max_tokens=1500, request_timeout=45
            )
            temp_chain = LLMChain(llm=temp_llm, prompt=self.prompt_template)
            logging.info("Temporary LLM and Chain initialized successfully.")
        except Exception as e:
            logging.error(f"LLM Initialization failed: {str(e)}", exc_info=True)
            error_msg = "Error: Failed to initialize AI model. Please check your network connection and API key validity."
            if "authentication" in str(e).lower():
                error_msg = "Error: OpenAI API Key is invalid or expired. Please check your key."
            return {"answer": f"<div class='error-message'>{error_msg}</div><div class='error-details'>Details: {str(e)}</div>", "context_used": "N/A - LLM Init Failed"}

        context = "No relevant context found."
        statutes_from_context = "Statute retrieval skipped due to context issues."
        try:
            logging.info(f"Querying Vector DB for query: '{query[:50]}...' in state '{state}'...")
            results = self.vector_db.query(query, state=state, n_results=n_results)
            logging.info(f"Vector DB query successful for state '{state}'. Processing results...")

            context_parts = []
            doc_results = results.get("document_results", {})
            docs = doc_results.get("documents", [[]])[0]
            metadatas = doc_results.get("metadatas", [[]])[0]
            if docs and metadatas and len(docs) == len(metadatas):
                logging.info(f"Found {len(docs)} document chunks.")
                for i, doc_content in enumerate(docs):
                    metadata = metadatas[i]
                    state_label = metadata.get('state', 'Unknown State')
                    chunk_id = metadata.get('chunk_id', 'N/A')
                    context_parts.append(f"**Source: Document Chunk {chunk_id} (State: {state_label})**\n{doc_content}")

            state_results_data = results.get("state_results", {})
            state_docs = state_results_data.get("documents", [[]])[0]
            state_metadatas = state_results_data.get("metadatas", [[]])[0]
            if state_docs and state_metadatas and len(state_docs) == len(state_metadatas):
                logging.info(f"Found {len(state_docs)} state summary documents.")
                for i, state_doc_content in enumerate(state_docs):
                    metadata = state_metadatas[i]
                    state_label = metadata.get('state', state)
                    context_parts.append(f"**Source: State Summary (State: {state_label})**\n{state_doc_content}")

            if context_parts:
                context = "\n\n---\n\n".join(context_parts)
                logging.info(f"Constructed context with {len(context_parts)} parts. Length: {len(context)} chars.")
                try:
                    statutes_from_context = self.extract_statutes(context)
                except Exception as e:
                    logging.error(f"Error extracting statutes: {e}", exc_info=True)
                    statutes_from_context = "Error extracting statutes from context."
            else:
                logging.warning("No relevant context parts found from vector DB query.")
                context = "No relevant context could be retrieved from the knowledge base for this query and state. The AI will answer from its general knowledge."
                statutes_from_context = "No specific statutes found as no context was retrieved."

        except Exception as e:
            logging.error(f"Vector DB query/context processing failed: {str(e)}", exc_info=True)
            context = f"Warning: Error retrieving documents from the knowledge base ({str(e)}). The AI will attempt to answer from its general knowledge, which may be less specific or accurate."
            statutes_from_context = "Statute retrieval skipped due to error retrieving context."

        try:
            logging.info("Invoking LLMChain with constructed input...")
            llm_input = {"query": query, "context": context, "state": state, "statutes": statutes_from_context}
            answer_dict = temp_chain.invoke(llm_input)
            answer_text = answer_dict.get('text', '').strip()

            if not answer_text:
                logging.warning("LLM returned an empty answer.")
                answer_text = "<div class='error-message'><span class='error-icon'>⚠️</span>The AI model returned an empty response. This might be due to the query, context limitations, or temporary issues. Please try rephrasing your question or try again later.</div>"
            else:
                logging.info("LLM generated answer successfully.")

            return {"answer": answer_text, "context_used": context}

        except Exception as e:
            logging.error(f"LLM processing failed: {str(e)}", exc_info=True)
            error_message = "Error: AI answer generation failed."
            details = f"Details: {str(e)}"
            if "authentication" in str(e).lower():
                error_message = "Error: Authentication failed. Please double-check your OpenAI API key."
                details = ""
            elif "rate limit" in str(e).lower():
                error_message = "Error: You've exceeded your OpenAI API rate limit or quota. Please check your usage and plan limits, or wait and try again."
                details = ""
            elif "context length" in str(e).lower():
                error_message = "Error: The request was too long for the AI model. This can happen with very complex questions or extensive retrieved context."
                details = "Try simplifying your question or asking about a more specific aspect."
            elif "timeout" in str(e).lower():
                error_message = "Error: The request to the AI model timed out. The service might be busy."
                details = "Please try again in a few moments."

            formatted_error = f"<div class='error-message'><span class='error-icon'>❌</span>{error_message}</div>"
            if details:
                formatted_error += f"<div class='error-details'>{details}</div>"

            return {"answer": formatted_error, "context_used": context}

    def process_query(self, query: str, state: str, openai_api_key: str, n_results: int = 5) -> Dict[str, any]:
        return self.process_query_cached(query.strip(), state, openai_api_key.strip(), n_results)

    def get_states(self) -> List[str]:
        try:
            states = self.vector_db.get_states()
            if not states:
                logging.warning("No states retrieved from vector_db. Returning empty list.")
                return []
            valid_states = sorted(list(set(s for s in states if s and isinstance(s, str) and s != "Select a state...")))
            logging.info(f"Retrieved {len(valid_states)} unique, valid states from VectorDatabase.")
            return valid_states
        except Exception as e:
            logging.error(f"Failed to get states from VectorDatabase: {str(e)}", exc_info=True)
            return ["Error: Could not load states"]

    def load_pdf(self, pdf_path: str) -> int:
        if not os.path.exists(pdf_path):
            logging.error(f"PDF file not found at path: {pdf_path}")
            raise FileNotFoundError(f"PDF file not found: {pdf_path}")
        try:
            logging.info(f"Attempting to load/verify data from PDF: {pdf_path}")
            # Assuming process_and_load_pdf is part of VectorDatabase and correctly implemented
            num_states_processed = self.vector_db.process_and_load_pdf(pdf_path)
            doc_count = self.vector_db.document_collection.count()
            state_count = self.vector_db.state_collection.count()
            total_items = doc_count + state_count

            if total_items > 0:
                logging.info(f"Vector DB contains {total_items} items ({doc_count} docs, {state_count} states). PDF processed or data already existed.")
                current_states = self.get_states()
                return len(current_states) if current_states and "Error" not in current_states[0] else 0
            else:
                logging.warning(f"PDF processing completed, but the vector database appears empty. Check PDF content and processing logs.")
                return 0

        except Exception as e:
            logging.error(f"Failed to load or process PDF '{pdf_path}': {str(e)}", exc_info=True)
            raise RuntimeError(f"Failed to process PDF '{pdf_path}': {e}") from e

    # --- GRADIO INTERFACE (NEW UI DESIGN) ---
    def gradio_interface(self):
        def query_interface_wrapper(api_key: str, query: str, state: str) -> str:
            # Basic client-side validation for immediate feedback (redundant but good UX)
            if not api_key or not api_key.strip() or not api_key.startswith("sk-"):
                return "<div class='error-message'><span class='error-icon'>⚠️</span>Please provide a valid OpenAI API key (starting with 'sk-'). <a href='https://platform.openai.com/api-keys' target='_blank'>OpenAI</a>.</div>"
            if not state or state is None: # Removed "Select a state..." from error check as it's not a choice
                return "<div class='error-message'><span class='error-icon'>⚠️</span>Please select a valid state from the list.</div>"
            if not query or not query.strip():
                return "<div class='error-message'><span class='error-icon'>⚠️</span>Please enter your question in the text box.</div>"

            # Call the core processing logic
            result = self.process_query(query=query, state=state, openai_api_key=api_key)
            answer = result.get("answer", "<div class='error-message'><span class='error-icon'>⚠️</span>An unexpected error occurred.</div>")

            # Check if the answer already contains an error message (from deeper within process_query)
            if "<div class='error-message'>" in answer:
                # Error messages are returned directly as they contain their own styling
                return answer
            else:
                # Wrap successful response in a div with an animation class
                formatted_response_content = f"<div class='response-header'><span class='response-icon'>📜</span>Response for {state}</div><hr class='divider'>{answer}"
                return f"<div class='animated-output-content'>{formatted_response_content}</div>"

        try:
            available_states_list = self.get_states()
            # DEBUG: Print states to console to verify if data is loaded
            print(f"DEBUG: States loaded for selection: {available_states_list}")
            # For Radio, remove "Select a state..." from choices, initial value will be None
            radio_choices = available_states_list if available_states_list and "Error" not in available_states_list[0] else ["Error: States unavailable"]
            initial_value_radio = None # Set initial value to None, forcing user to pick
        except Exception as e: # Catch-all for safety
            print(f"DEBUG: Error loading states for selection: {e}")
            radio_choices = ["Error: Critical failure loading states"]
            initial_value_radio = None


        # Define example queries, filtering based on available states
        example_queries_base = [
            ["What are the rules for security deposit returns?", "California"],
            ["Can a landlord enter my apartment without notice?", "New York"],
            ["My landlord hasn't made necessary repairs. What can I do?", "Texas"],
            ["How much notice must a landlord give to raise rent?", "Florida"],
            ["What is an implied warranty of habitability?", "Illinois"]
        ]
        example_queries = []
        if available_states_list and "Error" not in available_states_list[0] and len(available_states_list) > 0:
            loaded_states_set = set(available_states_list)
            # Filter for examples whose state is in the loaded states
            example_queries = [ex for ex in example_queries_base if ex[1] in loaded_states_set]
            # Add a generic example if no specific state examples match or if list is empty
            if not example_queries:
                # Add one example using the first available state, or a common one if no states
                example_queries.append(["What basic rights do tenants have?", available_states_list[0] if available_states_list else "California"])
        else: # Fallback if states list is problematic (e.g., empty or error)
            example_queries.append(["What basic rights do tenants have?", "California"])


        # Custom CSS for better UI design, clear boundaries, and text alignment for HuggingFace
        custom_css = """
        /* Import legible fonts from Google Fonts */
        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=Poppins:wght@600;700;800&display=swap');
        
        /* Root variables for consistent theming - adjusted for a very warm, clean palette */
        :root {
            --primary-color: #FF8C00; /* Darker Orange for buttons/accents */
            --primary-hover: #E07B00; /* Slightly darker orange for hover */
            --background-primary: hsl(30, 100%, 99.8%); /* Almost pure white, with a very subtle warm tint */
            --background-secondary: hsl(30, 100%, 96%); /* Clear, light, warm peach/cream */
            --text-primary: #4A3C32; /* Dark warm brown/charcoal for main text */
            --text-secondary: #8C7B6F; /* Muted warm gray/brown for secondary text */
            --border-color: hsl(30, 70%, 85%); /* Light, warm orange-brown for borders */
            --border-focus: #FF8C00; /* Focus color matches primary */
            --shadow-sm: 0 1px 3px rgba(0,0,0,0.08);
            --shadow-md: 0 4px 10px rgba(0,0,0,0.1);
            --shadow-lg: 0 10px 20px rgba(0,0,0,0.15);
            --error-bg: #FFF0E0; /* Light orange-pink for errors */
            --error-border: #FFD2B2; /* Medium orange-pink for error borders */
            --error-text: #E05C00; /* Darker, strong orange-red for error text */
        }
        
        /* Dark mode variables - for consistency if a dark mode toggle were present */
        body.dark {
            --primary-color: #FFA500; /* Bright orange for dark mode */
            --primary-hover: #CC8400;
            --background-primary: #2C2C2C; /* Dark charcoal */
            --background-secondary: #1F1F1F; /* Even darker charcoal */
            --text-primary: #F0F0F0;
            --text-secondary: #B0B0B0;
            --border-color: #555555;
            --border-focus: #FFA500;
            --error-bg: #400000;
            --error-border: #800000;
            --error-text: #FF6666;
        }

        /* Ensure the very outer body background is also set, overriding any Gradio defaults */
        body, html {
            background-color: var(--background-secondary) !important;
        }

        /* Base container improvements */
        .gradio-container {
            max-width: 900px !important; /* Slightly smaller for focused content */
            margin: 0 auto !important; /* Center the whole app */
            padding: 1.5rem !important;
            font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
            background-color: var(--background-secondary) !important; /* Overall background of the container */
            box-shadow: none !important; /* Remove default gradio container shadow */
        }
        /* All main content sections (cards) will have the primary background, not a gradient */
        .main-dashboard-container > * {
            background-color: var(--background-primary) !important;
        }

        /* Header styling - centered and prominent */
        .app-header-wrapper {
            background: linear-gradient(145deg, var(--background-primary) 0%, var(--background-secondary) 100%) !important; /* Gradient background */
            border: 2px solid var(--border-color) !important;
            border-radius: 16px !important;
            padding: 2.5rem 1.5rem !important; /* More vertical padding */
            margin-bottom: 1.5rem !important;
            box-shadow: var(--shadow-md) !important;
            position: relative; /* For potential pseudo-element effects */
            overflow: hidden; /* For any overflow animations */
            text-align: center !important; /* Centers block content within the wrapper */
        }

        .app-header-wrapper::before { /* Subtle background pattern for dynamism */
            content: '';
            position: absolute;
            top: 0;
            left: 0;
            width: 100%;
            height: 100%;
            background: radial-gradient(circle at top left, rgba(255,140,0,0.3) 0%, transparent 60%), /* More vibrant orange tint */
                        radial-gradient(circle at bottom right, rgba(255,140,0,0.3) 0%, transparent 60%);
            z-index: 0;
            opacity: 0.8;
            pointer-events: none;
        }

        .app-header-logo {
            font-size: 8.5rem !important; /* Significantly larger icon */
            margin-bottom: 0.75rem !important;
            display: block !important; 
            color: var(--primary-color) !important; /* Theme color */
            position: relative;
            z-index: 1; /* Bring icon to front of pseudo-element */
            animation: float-icon 3s ease-in-out infinite alternate;
        }
        /* Keyframes for floating icon */
        @keyframes float-icon {
            0% { transform: translateY(0px); }
            50% { transform: translateY(-5px); }
            100% { transform: translateY(0px); }
        }

        .app-header-title {
            font-family: 'Poppins', sans-serif !important;
            font-size: 3rem !important; /* Even larger title */
            font-weight: 800 !important; /* Bolder */
            color: var(--text-primary) !important;
            margin: 0 0 0.75rem 0 !important;
            line-height: 1.1 !important;
            letter-spacing: -0.03em !important; /* Tighter spacing */
            position: relative;
            z-index: 1;
            display: inline-block; /* Essential for text-align: center on parent to work for this block */
            max-width: 100%; /* Prevent overflow on smaller screens */
        }
        .app-header-tagline {
            font-size: 1.25rem !important; /* Slightly larger tagline */
            color: var(--text-secondary) !important;
            font-weight: 400 !important;
            margin: 0 !important;
            max-width: 700px; /* Constrain tagline width */
            display: inline-block; /* Essential for text-align: center on parent to work for this block */
            position: relative;
            z-index: 1;
        }

        /* Main container with consistent spacing */
        .main-dashboard-container {
            display: flex !important;
            flex-direction: column !important;
            gap: 1.25rem !important; /* Consistent spacing between cards */
        }
        /* Card sections with clear boundaries and subtle dynamic effects */
        .dashboard-card-section {
            background-color: var(--background-primary) !important; /* Solid primary background for card body */
            border: 2px solid var(--border-color) !important; /* Distinct border */
            border-radius: 12px !important;
            padding: 0 !important; /* Removed padding to allow gradient bar and content area to control it */
            box-shadow: var(--shadow-sm) !important; /* Subtle shadow */
            transition: all 0.3s ease-out !important; /* Smoother transition */
            cursor: default; /* Indicate not directly clickable (unless examples) */
        }
        .dashboard-card-section:hover {
            box-shadow: var(--shadow-md) !important;
            transform: translateY(-3px) !important; /* More pronounced lift */
        }

        /* Class for Markdown blocks to center their content */
        .full-width-center {
            display: flex !important;
            justify-content: center !important;
            align-items: center !important;
            width: 100% !important;
            flex-direction: column !important; /* Ensure content stacks vertically if needed */
        }

        /* NEW: Class for the solid color title bar within each card */
        .section-title-gradient-bar { 
            background-color: var(--background-secondary) !important; /* Solid warm peach/cream */
            padding: 1.25rem 1.75rem !important; /* Inner padding for the title bar */
            border-top-left-radius: 10px !important; /* Match parent's border radius */
            border-top-right-radius: 10px !important;
            margin-bottom: 1rem !important; /* Space below the title bar */
            text-align: center !important; /* Ensure content inside this bar is centered */
            box-sizing: border-box; /* Include padding in width */
            width: 100%; /* Ensure it spans full width */
        }

        /* Section titles (h3 inside markdown) */
        .section-title {
            font-family: 'Poppins', sans-serif !important;
            font-size: 1.7rem !important; /* Slightly larger */
            font-weight: 700 !important; /* Bolder */
            color: var(--text-primary) !important;
            margin: 0 !important; /* No margin on h3 itself as parent handles spacing */
            padding-bottom: 0.1rem !important; /* REDUCED further to bring text closer to border */
            border-bottom: 2px solid var(--border-color) !important; /* Underline effect */
            line-height: 1.1 !important; /* Ensure line height does not add extra space */
            display: inline-block !important; /* Allow centering within text-align: center of parent */
            text-align: center !important; /* Fallback centering */
            letter-spacing: -0.01em !important;
        }
        
        /* General content area padding within dashboard sections (below the title bar) */
        .dashboard-card-content-area {
            padding: 0 1.75rem 1.75rem 1.75rem !important; /* Match overall padding, 0 top because title bar handles it */
            background-color: var(--background-primary) !important; /* Pure white background */
            box-sizing: border-box; /* Include padding in width */
            width: 100%; /* Ensure it spans full width */
        }
        .dashboard-card-section p {
            line-height: 1.7 !important;
            color: var(--text-primary) !important;
            font-size: 1rem !important;
            text-align: left !important; /* Ensure content text is left-aligned */
            background-color: transparent !important; /* Ensure p tags don't have unexpected backgrounds */
            margin: 0 !important; /* Remove default paragraph margins */
            padding: 0 !important; /* Remove default paragraph padding */
            white-space: normal !important; /* Ensure text wraps */
        }
        .dashboard-card-section strong, .dashboard-card-section b {
            font-weight: 700 !important; /* Ensure bold is actually bold */
            color: var(--primary-color) !important; /* Highlight strong text with primary color */
        }

        /* Overrides for common Gradio internal elements that might have default backgrounds */
        /* These ensure transparency or explicit primary background for common Gradio containers */
        .gr-block, .gr-box, .gr-prose, .gr-form, .gr-panel, .gr-columns, .gr-column,
        .gradio-html, .gradio-markdown, .gradio-textbox, .gradio-radio, .gradio-button {
            background-color: transparent !important; 
            color: var(--text-primary) !important; /* Ensure text color is consistent */
            /* Ensure text wrapping for markdown as well */
            white-space: normal !important;
            overflow-wrap: break-word;
            word-break: break-word;
        }


        /* Improved input styling with clear boundaries and focus */
        .gradio-textbox {
            margin-bottom: 0.75rem !important;
        }
        /* Target the actual input elements for background color */
        .gradio-textbox textarea,
        .gradio-textbox input {
            background-color: var(--background-primary) !important; /* Pure white for content area */
            border: 2px solid var(--border-color) !important; /* Clear border */
            border-radius: 8px !important;
            padding: 0.85rem 1rem !important; /* Slightly more padding */
            font-size: 0.98rem !important;
            font-family: 'Inter', sans-serif !important;
            color: var(--text-primary) !important; /* Text color inside inputs */
            transition: border-color 0.2s ease, box-shadow 0.2s ease !important; /* Smooth transitions */
            box-shadow: var(--shadow-sm) !important;
        }
        /* Target the internal scrollable div within gradio-textbox */
        .gradio-textbox .scroll-hide {
            background-color: var(--background-primary) !important; /* Ensure scrollable area is white */
        }
        /* Focus styles for textboxes */
        .gradio-textbox textarea:focus,
        .gradio-textbox input:focus {
            outline: none !important;
            border-color: var(--border-focus) !important; /* Distinct border on focus */
            box-shadow: 0 0 0 4px rgba(255, 140, 0, 0.2) !important; /* Broader, softer glow on focus */
        }

        /* Styling for the radio button group (state selection) */
        .gradio-radio {
            padding: 0 !important; /* Remove any default padding */
            margin-top: 1rem !important; /* Add a little space above */
        }
        /* Hide default radio circle/dot */
        .gradio-radio input[type="radio"] {
            display: none !important;
        }

        .gradio-radio label {
            /* Style the clickable area for each radio option */
            display: flex !important; /* Use flexbox for internal alignment */
            justify-content: center !important; /* Center content horizontally */
            align-items: center !important;
            padding: 0.75rem 1rem !important;
            border: 2px solid var(--border-color) !important;
            border-radius: 8px !important;
            background-color: var(--background-primary) !important; /* Matches card background (white) */
            color: var(--text-primary) !important;
            font-weight: 500 !important;
            cursor: pointer !important;
            transition: all 0.2s ease-out !important;
            box-shadow: var(--shadow-sm) !important;
            margin: 0.4rem 0 !important; /* Increased vertical margin between options */
            width: 100% !important; /* Ensure options take full width of their column */
            box-sizing: border-box !important; /* Include padding/border in width */
        }

        /* Style the text/content within the radio label */
        .gradio-radio label span.text-lg { /* Gradio uses text-lg for the label text by default */
            font-weight: 600 !important; /* Make text bold */
            color: var(--text-primary) !important;
            font-size: 0.98rem !important; /* Match input text size */
        }

        /* Hover effect for radio options */
        .gradio-radio label:hover {
            background-color: var(--background-secondary) !important; /* Slightly darker cream on hover */
            border-color: var(--primary-color) !important; /* Highlight border with primary color */
            box-shadow: var(--shadow-md) !important;
            transform: translateY(-2px) !important;
        }

        /* Selected state for radio options */
        .gradio-radio input[type="radio"]:checked + label { /* Target label when its radio input is checked */
            background-color: var(--primary-color) !important; /* Primary color for selected item */
            color: white !important; /* White text on selected */
            border-color: var(--primary-hover) !important;
            box-shadow: var(--shadow-md) !important;
            transform: translateY(-1px) !important;
        }
        .gradio-radio input[type="radio"]:checked + label span.text-lg {
            color: white !important; /* Ensure text is white when selected */
        }
        /* Gradio's internal wrapper for radio buttons, ensure it doesn't add unwanted padding */
        .gradio-radio .gr-form {
            padding: 0 !important;
        }


        /* Label styling for better readability (for Query, State labels) */
        .gradio-textbox label,
        .gradio-radio > label { /* Target the main label for the radio group */
            font-weight: 600 !important; /* Bolder labels */
            color: var(--text-primary) !important;
            font-size: 1rem !important;
            margin-bottom: 0.6rem !important;
            display: block !important;
            text-align: left !important; /* Ensure these labels are left-aligned */
        }
        /* Info text styling below inputs (e.g., for API Key) */
        /* Specifically target gr-prose if it is the direct child of a gr.Block or gr.Group */
        .gr-prose { 
            font-size: 0.9rem !important;
            color: var(--text-secondary) !important;
            margin-top: 0.4rem !important; /* More space for info text */
            text-align: left !important; /* Ensure info text is left aligned */
            background-color: transparent !important; /* Ensure no unwanted background */
        }
        /* Input column layout improvements */
        .input-column { /* Renamed from .input-row */
            display: flex !important;
            flex-direction: column !important; /* Stack items vertically */
            gap: 1.25rem !important; /* Consistent gap between query and state input */
            margin-bottom: 0.5rem !important;
        }
        .input-field {
            flex: none !important; /* Remove flex sizing as items are stacked */
            width: 100% !important; /* Ensure each field takes full width */
        }

        /* Button styling improvements with active state for dynamism */
        .button-row {
            display: flex !important;
            gap: 1rem !important;
            justify-content: flex-end !important; /* Align buttons to the right */
            margin-top: 1.5rem !important; /* More space above buttons */
        }
        .gradio-button {
            padding: 0.85rem 1.8rem !important; /* More padding for bigger buttons */
            border-radius: 9px !important; /* Slightly more rounded */
            font-weight: 600 !important; /* Bolder text */
            font-size: 1rem !important;
            transition: all 0.2s ease-out !important; /* Smooth transition for hover/active */
            cursor: pointer !important;
            border: 2px solid transparent !important;
            text-align: center !important; /* Ensure button text is centered */
        }
        .gr-button-primary {
            background-color: var(--primary-color) !important; /* Explicitly set background */
            color: white !important;
            box-shadow: var(--shadow-sm) !important;
        }
        .gr-button-primary:hover {
            background-color: var(--primary-hover) !important; /* Explicitly set background */
            box-shadow: var(--shadow-md) !important;
            transform: translateY(-2px) !important; /* Subtle lift effect on hover */
        }
        .gr-button-primary:active { /* Press down effect on click */
            transform: translateY(1px) !important;
            box-shadow: none !important;
        }
        .gr-button-secondary {
            background-color: transparent !important; /* Explicitly set background */
            color: var(--text-primary) !important;
            border-color: var(--border-color) !important;
        }
        .gr-button-secondary:hover {
            background-color: var(--background-secondary) !important; /* Explicitly set background */
            border-color: var(--primary-color) !important;
            transform: translateY(-2px) !important;
        }
        .gr-button-secondary:active { /* Press down effect on click */
            transform: translateY(1px) !important;
            box-shadow: none !important;
        }

        /* Output styling with clear boundaries and dynamic fade-in */
        .output-content-wrapper {
            background-color: var(--background-primary) !important; /* Explicitly set background */
            border: 2px solid var(--border-color) !important; /* Clear border */
            border-radius: 8px !important;
            padding: 1.5rem !important;
            min-height: 150px !important; /* More space for output */
            color: var(--text-primary) !important;
            /* Ensure the inner animated content fits well */
            display: flex;
            flex-direction: column;
            justify-content: center; /* Center content vertically if small */
            align-items: center; /* Center content horizontally if small */
        }
        /* The div holding the actual response content, enabling fade-in animation */
        .animated-output-content {
            opacity: 0;
            animation: fadeInAndSlideUp 0.7s ease-out forwards; /* More pronounced animation */
            width: 100%; /* Take full width of parent */
            /* Preserve formatting within the animated content */
            white-space: pre-wrap;
            overflow-wrap: break-word;
            word-break: break-word;
            text-align: left !important; /* Ensure text is left-aligned within this div */
        }
        @keyframes fadeInAndSlideUp {
            from { opacity: 0; transform: translateY(15px); }
            to { opacity: 1; transform: translateY(0); }
        }

        .response-header {
            font-size: 1.3rem !important;
            font-weight: 700 !important;
            color: var(--primary-color) !important; /* Matches primary color */
            margin-bottom: 0.75rem !important;
            display: flex !important;
            align-items: center !important;
            gap: 0.6rem !important;
            text-align: left !important; /* Ensure header itself is not affected by parent centering */
            width: 100%; /* Take full width */
            justify-content: flex-start; /* Align content to the start */
        }
        .response-icon {
            font-size: 1.5rem !important;
            color: var(--primary-color) !important;
        }
        .divider {
            border: none !important;
            border-top: 1px dashed var(--border-color) !important; /* Dashed divider for visual separation */
            margin: 1rem 0 !important;
        }
        /* Error message styling */
        .error-message {
            background-color: var(--error-bg) !important; /* Explicitly set background */
            border: 2px solid var(--error-border) !important;
            color: var(--error-text) !important;
            padding: 1.25rem !important;
            border-radius: 8px !important;
            display: flex !important;
            align-items: flex-start !important;
            gap: 0.8rem !important;
            font-size: 0.95rem !important;
            font-weight: 500 !important;
            line-height: 1.6 !important;
            text-align: left !important; /* Ensure error message text is left aligned */
            width: 100%; /* Take full width of parent */
            box-sizing: border-box; /* Include padding/border in width */
        }
        .error-message a {
            color: var(--error-text) !important;
            text-decoration: underline !important;
        }
        .error-icon {
            font-size: 1.4rem !important;
            line-height: 1 !important;
            margin-top: 0.1rem !important;
        }
        .error-details {
            font-size: 0.85rem !important;
            color: var(--error-text) !important;
            margin-top: 0.5rem !important;
            opacity: 0.8;
        }
        /* Placeholder styling for empty output */
        .placeholder {
            background-color: var(--background-primary) !important; /* Explicitly set background to primary white */
            border: 2px dashed var(--border-color) !important;
            border-radius: 8px !important;
            padding: 2.5rem 1.5rem !important;
            text-align: center !important; /* Ensure placeholder text is centered */
            color: var(--text-secondary) !important;
            font-style: italic !important;
            font-size: 1.1rem !important;
            width: 100%; /* Ensure it takes full width of parent */
            box-sizing: border-box; /* Include padding/border in width */
        }

        /* Examples table styling with dynamic hover */
        .examples-section .gr-samples-table {
            border: 2px solid var(--border-color) !important;
            border-radius: 8px !important;
            overflow: hidden !important;
            margin-top: 1rem !important;
        }
        .examples-section .gr-samples-table th,
        .examples-section .gr-samples-table td {
            padding: 0.9rem !important;
            border: none !important;
            font-size: 0.95rem !important;
            text-align: left !important; /* Ensure example text is left-aligned */
        }
        .examples-section .gr-samples-table th {
            background-color: var(--background-secondary) !important; /* Explicitly set background */
            font-weight: 700 !important;
            color: var(--text-primary) !important;
        }
        .examples-section .gr-samples-table td {
            background-color: var(--background-primary) !important; /* Explicitly set background */
            color: var(--text-primary) !important;
            border-top: 1px solid var(--border-color) !important;
            cursor: pointer !important;
            transition: background 0.2s ease, transform 0.1s ease !important; /* Smooth transitions */
        }
        .examples-section .gr-samples-table tr:hover td {
            background-color: var(--background-secondary) !important; /* Explicitly set background */
            transform: translateX(5px); /* Subtle slide on hover */
        }
        /* Hide Gradio default elements for examples for cleaner look */
        .gr-examples .gr-label,
        .gr-examples .label-wrap,
        .gr-examples .gr-accordion-header {
            display: none !important;
        }
        
        /* Footer styling - centered text */
        .app-footer-wrapper {
            background: linear-gradient(145deg, var(--background-primary) 0%, var(--background-secondary) 100%) !important; /* Gradient background */
            border: 2px solid var(--border-color) !important;
            border-radius: 12px !important;
            padding: 1.75rem !important;
            margin-top: 1.5rem !important;
            margin-bottom: 1.5rem !important; /* Ensures space at the very bottom before the container padding */
            text-align: center !important; /* Centered footer text */
        }
        .app-footer p {
            margin: 0.6rem auto !important; /* Auto margins to center block */
            max-width: 90% !important; /* Constrain width for wrapping */
            font-size: 0.95rem !important;
            color: var(--text-secondary) !important;
            line-height: 1.6 !important;
            background-color: transparent !important; /* Ensure paragraph in footer does not get unexpected background */
            text-align: center !important; /* Ensure footer text is centered */
            white-space: normal !important; /* Allow text to wrap */
        }
        .app-footer strong, .app-footer b {
            font-weight: 700 !important; /* Ensure bold is actually bold */
            color: var(--primary-color) !important; /* Highlight strong text with primary color */
        }
        .app-footer a {
            color: var(--primary-color) !important;
            text-decoration: underline !important; /* Ensure links are underlined */
            font-weight: 600 !important;
        }
        .app-footer a:hover {
            text-decoration: none !important; /* Remove underline on hover for subtle effect */
        }

        /* Responsive design for smaller screens */
        @media (max-width: 768px) {
            .gradio-container {
                padding: 1rem !important;
            }
            .app-header-title {
                font-size: 2.2rem !important;
            }
            .app-header-tagline {
                font-size: 1rem !important;
            }
            .section-title {
                font-size: 1.4rem !important;
            }
            .input-column { /* Apply vertical stacking for columns */
                flex-direction: column !important; 
            }
            .button-row {
                flex-direction: column !important; /* Stack buttons vertically */
            }
            .gradio-button {
                width: 100% !important; /* Full width buttons */
            }
            .dashboard-card-section {
                /* padding handled by section-title-gradient-bar and dashboard-card-content-area */
            }
            .section-title-gradient-bar {
                padding: 0.8rem 1rem !important;
            }
            .dashboard-card-content-area {
                padding: 0 1rem 1rem 1rem !important;
            }
            .output-content-wrapper {
                min-height: 120px !important;
            }
            .placeholder {
                padding: 1.5rem 1rem !important;
                font-size: 1rem !important;
            }
        }
        """

        # Using gr.Blocks with custom CSS only to ensure no undesired colors are inherited from a theme.
        with gr.Blocks(css=custom_css, title="Landlord-Tenant Rights Assistant") as demo:
            # Header Section - uses gr.Group for distinct card-like styling
            with gr.Group(elem_classes="app-header-wrapper"):
                gr.Markdown(
                    """
                    <span class='app-header-logo'>⚖️</span>
                    <h1 class='app-header-title'>Landlord-Tenant Rights Assistant</h1>
                    <p class='app-header-tagline'>Empowering You with State-Specific Legal Insights</p>
                    """,
                    elem_classes="full-width-center"
                )

            # Main Dashboard Container - acts as a column to stack various sections
            with gr.Column(elem_classes="main-dashboard-container"):

                # How This Assistant Works Box (formerly Welcome & Disclaimer)
                with gr.Group(elem_classes="dashboard-card-section"):
                    # Title bar with solid warm peach/cream background
                    gr.Markdown("<h3 class='section-title'>How This Assistant Works</h3>", elem_classes="full-width-center section-title-gradient-bar")
                    # Content area with pure primary background
                    with gr.Column(elem_classes="dashboard-card-content-area"): # New wrapper for content area
                        gr.Markdown(
                            """
                            This AI-powered assistant helps navigate complex landlord-tenant laws. Simply ask a question about your state's regulations, and it will provide detailed, legally-grounded insights.
                            """
                        )

                # OpenAI API Key Input Card
                with gr.Group(elem_classes="dashboard-card-section"):
                    # Title bar with solid warm peach/cream background
                    gr.Markdown("<h3 class='section-title'>OpenAI API Key</h3>", elem_classes="full-width-center section-title-gradient-bar")
                    # Content area with pure primary background
                    with gr.Column(elem_classes="dashboard-card-content-area"): # New wrapper for content area
                        api_key_input = gr.Textbox(
                            label="API Key",
                            type="password", # Hides the input for security
                            placeholder="Enter your OpenAI API key (e.g., sk-...)",
                            lines=1,
                            elem_classes=["input-field-group"] # Custom class for input styling
                        )
                        # Separate Markdown for the clickable info link, using standard Markdown link syntax
                        gr.Markdown(
                            "Required to process your query. Get one from OpenAI: [platform.openai.com/api-keys](https://platform.openai.com/api-keys)",
                            elem_classes="gr-prose" # Reuse Gradio's info text class for styling
                        )

                # Query Input and State Selection Card
                with gr.Group(elem_classes="dashboard-card-section"):
                    # Title bar with solid warm peach/cream background
                    gr.Markdown("<h3 class='section-title'>Ask Your Question</h3>", elem_classes="full-width-center section-title-gradient-bar")
                    # Content area with pure primary background
                    with gr.Column(elem_classes="dashboard-card-content-area"): # New wrapper for content area
                        with gr.Column(elem_classes="input-column"): # Column to stack inputs
                            with gr.Column(elem_classes="input-field", scale=1): # Query box takes full width
                                query_input = gr.Textbox(
                                    label="Your Question",
                                    placeholder="E.g., What are the rules for security deposit returns in my state?",
                                    lines=8, # Increased default height
                                    max_lines=15, # Increased max height
                                    elem_classes=["input-field-group"]
                                )
                            with gr.Column(elem_classes="input-field", scale=1): # State radio buttons take full width below
                                state_input = gr.Radio(
                                    label="Select State",
                                    choices=radio_choices,
                                    value=initial_value_radio, # Set initial value to None
                                    elem_classes=["input-field-group", "gradio-radio-custom"], # Added custom class for specific styling
                                    interactive=True # Ensure it's interactive
                                )
                        
                        with gr.Row(elem_classes="button-row"): # Row for action buttons
                            clear_button = gr.Button("Clear", variant="secondary", elem_classes=["gr-button-secondary"])
                            submit_button = gr.Button("Submit Query", variant="primary", elem_classes=["gr-button-primary"])

                # Output Display Card - Using gr.HTML for better animation control
                with gr.Group(elem_classes="dashboard-card-section"):
                    # Title bar with solid warm peach/cream background
                    gr.Markdown("<h3 class='section-title'>Legal Assistant's Response</h3>", elem_classes="full-width-center section-title-gradient-bar")
                    # Content area with pure primary background
                    with gr.Column(elem_classes="dashboard-card-content-area"): # New wrapper for content area
                        output = gr.HTML( # Changed to gr.HTML to wrap content with animation class
                            value="<div class='placeholder'>The answer will appear here after submitting your query.</div>",
                            elem_classes="output-content-wrapper" # Custom class for output styling
                        )

                # Example Questions Section
                with gr.Group(elem_classes="dashboard-card-section examples-section"):
                    # Title bar with solid warm peach/cream background
                    gr.Markdown("<h3 class='section-title'>Example Questions</h3>", elem_classes="full-width-center section-title-gradient-bar")
                    # Content area with pure primary background
                    with gr.Column(elem_classes="dashboard-card-content-area"): # New wrapper for content area
                        if example_queries:
                            gr.Examples(
                                examples=example_queries,
                                inputs=[query_input, state_input],
                                examples_per_page=5,
                                label="" # Hide default Gradio label for examples to use our custom title
                            )
                        else:
                            gr.Markdown("<div class='placeholder'>Sample questions could not be loaded. Please ensure the vector database is populated.</div>")

            # Footer Section - contains disclaimer and developer info (now including the full disclaimer)
            with gr.Group(elem_classes="app-footer-wrapper"):
                # Using paragraph tags for proper text wrapping and bolding in Markdown
                gr.Markdown(
                    """
                    <p>**Disclaimer:** This tool is for informational purposes only and does not constitute legal advice. For specific legal guidance, always consult with a licensed attorney in your jurisdiction.</p>
                    <p>Developed by **Nischal Subedi**. Connect on [LinkedIn](https://www.linkedin.com/in/nischal1/) or explore insights at [Substack](https://datascientistinsights.substack.com/).</p>
                    """
                )

            # Event Listeners for buttons
            submit_button.click(
                fn=query_interface_wrapper,
                inputs=[api_key_input, query_input, state_input],
                outputs=output,
                api_name="submit_query" # Useful for debugging / external calls
            )
            
            clear_button.click(
                fn=lambda: (
                    "", # Clear API key input
                    "", # Clear query input
                    initial_value_radio, # Reset state radio to default (None)
                    "<div class='placeholder'>Inputs cleared. Ready for your next question.</div>" # Reset output message
                ),
                inputs=[],
                outputs=[api_key_input, query_input, state_input, output]
            )

            return demo

# --- Main Execution Block (UNCHANGED from original logic) ---
if __name__ == "__main__":
    logging.info("Starting Landlord-Tenant Rights Bot application...")
    try:
        SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
        DEFAULT_PDF_PATH = os.path.join(SCRIPT_DIR, "tenant-landlord.pdf")
        DEFAULT_DB_PATH = os.path.join(SCRIPT_DIR, "chroma_db")

        PDF_PATH = os.getenv("PDF_PATH", DEFAULT_PDF_PATH)
        VECTOR_DB_PATH = os.getenv("VECTOR_DB_PATH", DEFAULT_DB_PATH)

        # Ensure vector DB directory exists before initialization
        os.makedirs(os.path.dirname(VECTOR_DB_PATH), exist_ok=True)

        logging.info(f"Attempting to load PDF from: {PDF_PATH}")
        if not os.path.exists(PDF_PATH):
            logging.error(f"FATAL: PDF file not found at the specified path: {PDF_PATH}")
            print(f"\n--- CONFIGURATION ERROR ---\nPDF file ('{os.path.basename(PDF_PATH)}') not found at: {PDF_PATH}.\nPlease ensure it exists or set 'PDF_PATH' environment variable.\n---------------------------\n")
            exit(1) # Correctly exits if PDF is not found

        if not os.access(PDF_PATH, os.R_OK):
            logging.error(f"FATAL: PDF file at '{PDF_PATH}' exists but is not readable. Check file permissions.")
            print(f"\n--- PERMISSION ERROR ---\nPDF file ('{os.path.basename(PDF_PATH)}') found but not readable at: {PDF_PATH}\nPlease check file permissions (e.g., using 'chmod +r' in terminal).\n---------------------------\n")
            exit(1) # Correctly exits if PDF is unreadable

        logging.info(f"PDF file '{os.path.basename(PDF_PATH)}' found and is readable.")

        # Initialize VectorDatabase and RAGSystem
        vector_db_instance = VectorDatabase(persist_directory=VECTOR_DB_PATH)
        rag = RAGSystem(vector_db=vector_db_instance)

        # Load PDF data into the vector DB (or verify it's already loaded)
        rag.load_pdf(PDF_PATH)

        # Get the Gradio interface object
        app_interface = rag.gradio_interface()
        SERVER_PORT = int(os.getenv("PORT", 7860)) # Use PORT env var for Hugging Face Spaces

        logging.info(f"Launching Gradio app on http://0.0.0.0:{SERVER_PORT}")
        print(f"\n--- Gradio App Running ---\nAccess at: http://localhost:{SERVER_PORT} or your public Spaces URL\n--------------------------\n")
        app_interface.launch(server_name="0.0.0.0", server_port=SERVER_PORT, share=False) # share=False is typical for Spaces

    except ModuleNotFoundError as e:
        if "vector_db" in str(e):
             logging.error(f"FATAL: Could not import VectorDatabase. Ensure 'vector_db.py' is in the same directory and 'chromadb', 'langchain', 'pypdf', 'sentence-transformers' are installed.", exc_info=True)
             print(f"\n--- MISSING DEPENDENCY OR FILE ---\nCould not find/import 'vector_db.py' or one of its dependencies.\nError: {e}\nPlease ensure 'vector_db.py' is present and all required packages (chromadb, langchain, pypdf, sentence-transformers, etc.) are in your requirements.txt and installed.\n---------------------------\n")
        else:
            logging.error(f"Application startup failed due to a missing module: {str(e)}", exc_info=True)
            print(f"\n--- FATAL STARTUP ERROR - MISSING MODULE ---\n{str(e)}\nPlease ensure all dependencies are installed.\nCheck logs for more details.\n---------------------------\n")
        exit(1)
    except FileNotFoundError as e:
        logging.error(f"Application startup failed due to a missing file: {str(e)}", exc_info=True)
        print(f"\n--- FATAL STARTUP ERROR - FILE NOT FOUND ---\n{str(e)}\nPlease ensure the file exists at the specified path.\nCheck logs for more details.\n---------------------------\n")
        exit(1)
    except Exception as e:
        logging.error(f"Application startup failed: {str(e)}", exc_info=True)
        print(f"\n--- FATAL STARTUP ERROR ---\n{str(e)}\nCheck logs for more details.\n---------------------------\n")
        exit(1)