app.py

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

import gradio as gr
try:
    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(1)

try:
    from langchain_openai import ChatOpenAI
except ImportError:
    print("Error: langchain-openai not found. Please install it: pip install langchain-openai")
    exit(1)

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

# Suppress warnings
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
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s'
)

# --- RAGSystem Class ---
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 CONTEXT ---
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 == "Select a state..." or "Error" in state:
            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'>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'>{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}")
            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 ---
    def gradio_interface(self):
        # Wrapper function for the Gradio interface logic
        def query_interface_wrapper(api_key: str, query: str, state: str) -> str:
            logging.info(f"Gradio interface received query: '{query[:50]}...', state: '{state}'")

            # Re-validate inputs robustly
            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'>Get one here</a>.</div>"
            if not state or state == "Select a state..." or "Error" in state:
                return "<div class='error-message'><span class='error-icon'>⚠️</span>Please select a valid state from the dropdown.</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)

            # Format the response for display
            answer = result.get("answer", "<div class='error-message'><span class='error-icon'>⚠️</span>An unexpected error occurred, and no answer was generated. Please check the logs or try again.</div>")

            # Add a header *only* if the answer is not an error message itself
            if not "<div class='error-message'>" in answer:
                 formatted_response = f"<div class='response-header'><span class='response-icon'>📜</span>Response for {state}</div><hr class='divider'>{answer}"
            else:
                 formatted_response = answer # Pass through error messages directly

            # Log context length for debugging (optional)
            context_used = result.get("context_used", "N/A")
            if isinstance(context_used, str) and "N/A" not in context_used:
                logging.debug(f"Context length used for query: {len(context_used)} characters.")
            else:
                logging.debug(f"No context was used or available for this query ({context_used}).")

            return formatted_response

        # --- Get Available States for Dropdown ---
        try:
            available_states_list = self.get_states()
            if not available_states_list or "Error" in available_states_list[0]:
                dropdown_choices = ["Error: Could not load states"]
                initial_value = dropdown_choices[0]
                logging.error("Could not load states for dropdown. UI will show error.")
            else:
                dropdown_choices = ["Select a state..."] + available_states_list
                initial_value = dropdown_choices[0]
        except Exception as e:
            logging.error(f"Unexpected critical error getting states: {e}", exc_info=True)
            dropdown_choices = ["Error: Critical failure loading states"]
            initial_value = dropdown_choices[0]

        # --- Prepare Example Queries ---
        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"],
            ["What are the limits on rent increases in my state?", "Florida"],
            ["Is my lease automatically renewed if I don't move out?", "Illinois"],
            ["What happens if I break my lease early?", "Washington"]
        ]
        example_queries = []
        if available_states_list and "Error" not in available_states_list[0]:
            loaded_states_set = set(available_states_list)
            example_queries = [ex for ex in example_queries_base if ex[1] in loaded_states_set]
        if not example_queries:
             fallback_state = available_states_list[0] if available_states_list and "Error" not in available_states_list[0] else "California"
             example_queries.append(["What basic rights do tenants have?", fallback_state])

        # --- Refined Custom CSS ---
        custom_css = """
        @import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap');

        /* --- CSS Variables for Theme Consistency --- */
        :root {
            --primary-color: #2563EB; /* Tailwind Blue 600 */
            --primary-hover: #1D4ED8; /* Tailwind Blue 700 */
            --secondary-color: #4B5563; /* Gray 600 */
            --secondary-hover: #374151; /* Gray 700 */
            --text-primary: #111827; /* Gray 900 */
            --text-secondary: #6B7280; /* Gray 500 */
            --background: #F3F4F6; /* Gray 100 */
            --card-background: #FFFFFF; /* White */
            --border-color: #D1D5DB; /* Gray 300 */
            --shadow: 0 2px 4px rgba(0, 0, 0, 0.04), 0 4px 8px rgba(0, 0, 0, 0.06); /* Softer, layered shadow */
            --error-bg: #FEF2F2; /* Red 50 */
            --error-border: #FECACA; /* Red 300 */
            --error-accent: #EF4444; /* Red 500 */
            --error-text: #B91C1C; /* Red 700 */
            --success-bg: #F0FDF4; /* Green 50 */
            --success-border: #A7F3D0; /* Green 300 */
            --success-text: #15803D; /* Green 700 */
            --divider: #E5E7EB; /* Gray 200 */
            --focus-ring: rgba(37, 99, 235, 0.3); /* Based on new primary */
        }

        /* Dark Mode Variables */
        @media (prefers-color-scheme: dark) {
            :root {
                --primary-color: #3B82F6; /* Tailwind Blue 500 (lighter for dark mode primary) */
                --primary-hover: #60A5FA; /* Tailwind Blue 400 */
                --text-primary: #F3F4F6; /* Gray 100 */
                --text-secondary: #9CA3AF; /* Gray 400 */
                --background: #111827; /* Gray 900 (very dark) */
                --card-background: #1F2937; /* Gray 800 (main content card bg) */
                --border-color: #4B5563; /* Gray 600 */
                --shadow: 0 2px 4px rgba(0, 0, 0, 0.1), 0 4px 8px rgba(0, 0, 0, 0.2);
                --error-bg: #450A0A; /* Darker Red */
                --error-border: #7F1D1D; /* Darker Red */
                --error-accent: #F87171; /* Lighter Red for accent */
                --error-text: #FECACA; /* Lighter Red for text */
                --success-bg: #064E3B; /* Darker Green */
                --success-border: #15803D; /* Darker Green */
                --success-text: #A7F3D0; /* Lighter Green */
                --divider: #374151; /* Gray 700 */
                --focus-ring: rgba(59, 130, 246, 0.4);
            }
        }

        /* --- Base & Body --- */
        body, .gradio-container {
            font-family: 'Poppins', -apple-system, BlinkMacSystemFont, sans-serif !important;
            background: var(--background) !important;
            color: var(--text-primary) !important;
            margin: 0;
            padding: 0;
            min-height: 100vh;
            font-size: 15px; /* Slightly smaller base font */
            line-height: 1.6; /* Increased line height for readability */
            -webkit-font-smoothing: antialiased;
            -moz-osx-font-smoothing: grayscale;
        }
        * {
            box-sizing: border-box;
        }

        /* --- Main Content Container --- */
        .gradio-container > .flex.flex-col {
            max-width: 960px; /* Slightly narrower for focus */
            margin: 0 auto !important;
            padding: 2.5rem 1.5rem !important; /* Adjusted padding */
            gap: 2rem !important; /* Adjusted gap */
            background: transparent !important;
        }

        /* --- Card Styling --- */
        .card-style {
            background: var(--card-background) !important;
            border: 1px solid var(--border-color) !important;
            border-radius: 12px !important; /* More modern radius */
            padding: 1.75rem !important; /* Adjusted padding */
            box-shadow: var(--shadow) !important;
            transition: transform 0.2s ease, background 0.2s ease, border 0.2s ease;
        }
        .card-style:hover {
            transform: translateY(-2px);
        }

        /* --- Header Section --- */
        .header-section {
            background: var(--primary-color) !important; /* Solid primary color */
            border-radius: 12px !important;
            padding: 2.5rem 2rem !important;
            text-align: center !important;
            color: #FFFFFF !important;
            box-shadow: var(--shadow) !important;
            position: relative;
            overflow: hidden;
        }
        .header-section::before { /* Subtle background pattern */
            content: '';
            position: absolute;
            top: 0; left: 0; width: 100%; height: 100%;
            background-image: radial-gradient(rgba(255, 255, 255, 0.07) 1px, transparent 1.2px);
            background-size: 8px 8px;
            opacity: 0.5;
            pointer-events: none;
        }
        .header-logo {
            font-size: 2.5rem; /* Adjusted size */
            margin-bottom: 0.75rem;
        }
        .header-title {
            font-size: 2rem; /* Adjusted size */
            font-weight: 600; /* Adjusted weight */
            margin: 0 0 0.5rem 0;
        }
        .header-tagline {
            font-size: 1.1rem; /* Adjusted size */
            font-weight: 400; /* Adjusted weight */
            opacity: 0.85;
        }

        /* --- Introduction Section --- */
        .intro-card h3 { /* Title like "Know Your Rights" */
            font-size: 1.5rem; /* Adjusted relative to new base */
            font-weight: 600;
            color: var(--primary-color);
            margin: 0 0 1rem 0;
            padding-bottom: 0.5rem;
            border-bottom: 2px solid var(--primary-color);
            display: inline-block;
        }
        .intro-card p {
            font-size: 0.95rem; /* Adjusted relative to new base */
            line-height: 1.7;
            color: var(--text-secondary);
            margin: 0 0 0.75rem 0;
        }
        .intro-card a {
            color: var(--primary-color);
            text-decoration: none;
            font-weight: 500;
            transition: color 0.2s ease;
        }
        .intro-card a:hover {
            color: var(--primary-hover);
            text-decoration: underline;
        }
        .intro-card strong {
            font-weight: 600;
            color: var(--text-primary);
        }

        /* --- Input Form Section --- */
        .input-form-card h3 { /* Title like "Ask Your Question" */
            font-size: 1.4rem; /* Adjusted */
            font-weight: 600;
            color: var(--text-primary);
            margin: 0 0 1.25rem 0;
            padding-bottom: 0.5rem;
            border-bottom: 1px solid var(--divider); /* Thinner divider */
        }
        .input-field-group {
            margin-bottom: 1.25rem;
        }
        .input-row {
            display: flex;
            gap: 1.25rem;
            flex-wrap: wrap;
            margin-bottom: 1.25rem;
        }
        .input-field {
            flex: 1;
            min-width: 200px; /* Adjusted min-width */
        }
        .gradio-textbox textarea,
        .gradio-dropdown select,
        .gradio-textbox input[type=password] {
            border: 1px solid var(--border-color) !important;
            border-radius: 8px !important; /* Sharper radius */
            padding: 0.8rem 1rem !important; /* Adjusted padding */
            font-size: 0.95rem !important; /* Adjusted font size */
            background: var(--card-background) !important; /* Use card for consistency, can be var(--background) for contrast */
            color: var(--text-primary) !important;
            transition: border-color 0.2s ease, box-shadow 0.2s ease, background 0.2s ease;
            width: 100% !important;
            box-shadow: inset 0 1px 2px rgba(0, 0, 0, 0.03);
        }
        .gradio-textbox textarea {
            min-height: 110px; /* Adjusted height */
            resize: vertical;
        }
        .gradio-textbox textarea:focus,
        .gradio-dropdown select:focus,
        .gradio-textbox input[type=password]:focus {
            border-color: var(--primary-color) !important;
            box-shadow: 0 0 0 3px var(--focus-ring) !important; /* Slightly smaller focus ring */
            outline: none !important;
            /* background: var(--background) !important; */ /* Keep card background on focus or change to main for contrast */
        }
        .gradio-input-label,
        .gradio-output-label {
            font-size: 0.9rem !important; /* Adjusted */
            font-weight: 500 !important;
            color: var(--text-primary) !important;
            margin-bottom: 0.4rem !important;
            display: block !important;
        }
        .gradio-input-info {
            font-size: 0.8rem !important; /* Adjusted */
            color: var(--text-secondary) !important;
            margin-top: 0.3rem;
            font-style: italic;
        }
        /* Buttons */
        .button-row {
            display: flex;
            gap: 0.75rem; /* Tighter gap */
            margin-top: 1.25rem;
            flex-wrap: wrap;
            justify-content: flex-end;
        }
        .gradio-button {
            border-radius: 8px !important; /* Sharper radius */
            padding: 0.75rem 1.5rem !important; /* Adjusted padding */
            font-size: 0.95rem !important; /* Adjusted font size */
            font-weight: 500 !important;
            border: none !important;
            cursor: pointer;
            transition: background-color 0.2s ease, transform 0.1s ease, box-shadow 0.2s ease;
            box-shadow: 0 1px 2px rgba(0,0,0,0.05) !important; /* More subtle shadow */
        }
        .gradio-button:hover:not(:disabled) {
            transform: translateY(-1px);
            box-shadow: 0 2px 4px rgba(0,0,0,0.07) !important;
        }
        .gradio-button:active:not(:disabled) {
            transform: scale(0.98) translateY(0);
        }
        .gradio-button:disabled {
            background: var(--border-color) !important;
            color: var(--text-secondary) !important;
            cursor: not-allowed;
            box-shadow: none !important;
        }
        .gr-button-primary {
            background: var(--primary-color) !important;
            color: #FFFFFF !important;
        }
        .gr-button-primary:hover:not(:disabled) {
            background: var(--primary-hover) !important;
        }
        .gr-button-secondary {
            background: var(--card-background) !important; /* Changed from transparent */
            color: var(--text-primary) !important;
            border: 1px solid var(--border-color) !important;
            box-shadow: none !important;
        }
        .gr-button-secondary:hover:not(:disabled) {
            background: var(--background) !important; /* Use main bg for hover */
            border-color: var(--secondary-hover) !important;
        }

        /* --- Output Section --- */
        .output-card .response-header {
            font-size: 1.3rem; /* Adjusted */
            font-weight: 600;
            color: var(--text-primary);
            margin: 0 0 0.75rem 0;
            display: flex;
            align-items: center;
            gap: 0.5rem;
        }
        .output-card .response-icon {
            font-size: 1.5rem; /* Adjusted */
        }
        .output-card .divider {
            border: none;
            border-top: 1px solid var(--divider);
            margin: 0.75rem 0 1.25rem 0; /* Adjusted margins */
        }
        .output-card .output-content-wrapper {
            font-size: 0.95rem; /* Adjusted */
            line-height: 1.7; /* Adjusted */
            color: var(--text-primary);
        }
        .output-card .output-content-wrapper p {
            margin-bottom: 0.85rem;
        }
        .output-card .output-content-wrapper ul,
        .output-card .output-content-wrapper ol {
            margin-left: 1.25rem;
            margin-bottom: 0.85rem;
            padding-left: 0.85rem;
        }
        .output-card .output-content-wrapper li {
            margin-bottom: 0.4rem;
        }
        .output-card .output-content-wrapper strong,
        .output-card .output-content-wrapper b {
            font-weight: 600;
            color: var(--text-primary);
        }
        .output-card .output-content-wrapper a {
            color: var(--primary-color);
            text-decoration: none;
            font-weight: 500;
        }
        .output-card .output-content-wrapper a:hover {
            color: var(--primary-hover);
            text-decoration: underline;
        }
        /* Error and Success Messages */
        .output-card .error-message,
        .output-card .success-message {
            display: flex;
            align-items: flex-start;
            gap: 0.6rem;
            border-radius: 8px; /* Match other radii */
            padding: 0.85rem 1.25rem; /* Adjusted padding */
            margin-top: 0.5rem;
            font-weight: 500;
            line-height: 1.5;
        }
        .output-card .error-message {
            background: var(--error-bg);
            border: 1px solid var(--error-border);
            border-left: 3px solid var(--error-accent); /* Thinner accent line */
            color: var(--error-text);
        }
        .output-card .success-message {
            background: var(--success-bg);
            border: 1px solid var(--success-border);
            color: var(--success-text);
            border-left: 3px solid var(--success-text);
        }
        .output-card .error-icon,
        .output-card .success-icon {
            font-size: 1.1rem; /* Adjusted */
            line-height: 1.5;
            margin-top: 2px; /* Align icon better with text */
        }
        .output-card .error-details {
            font-size: 0.85rem; /* Adjusted */
            color: var(--error-text); /* Ensure correct text color for dark mode if needed */
            margin-top: 0.4rem;
            font-style: italic;
        }
        /* Placeholder text */
        .output-card .placeholder {
            color: var(--text-secondary);
            font-style: italic;
            text-align: center;
            padding: 1.5rem 1rem; /* Adjusted padding */
            display: block;
            font-size: 1rem; /* Adjusted */
        }

        /* --- Examples Section --- */
        .examples-card .gr-examples-header {
            font-size: 1.3rem !important; /* Adjusted */
            font-weight: 600 !important;
            color: var(--text-primary) !important;
            margin: 0 0 1.25rem 0 !important;
            padding-bottom: 0.5rem !important;
            border-bottom: 1px solid var(--divider) !important; /* Thinner divider */
        }
        .examples-card .gr-examples-table {
            border-collapse: collapse !important;
            width: 100% !important;
            background: var(--card-background) !important;
            border-radius: 8px !important; /* Match other radii */
            overflow: hidden;
            border: 1px solid var(--border-color) !important; /* Add outer border to table */
        }
        .examples-card .gr-examples-table th,
        .examples-card .gr-examples-table td {
            text-align: left !important;
            padding: 0.75rem 1rem !important; /* Adjusted padding */
            border: 1px solid var(--border-color) !important;
            font-size: 0.9rem !important; /* Adjusted font size */
            color: var(--text-primary) !important;
            background: transparent !important;
        }
        .examples-card .gr-examples-table th {
            font-weight: 500 !important;
            background: var(--background) !important; /* Use main bg for header */
        }
        .examples-card .gr-examples-table tr {
            cursor: pointer;
            transition: background 0.2s ease;
        }
        .examples-card .gr-examples-table tr:hover td {
            background: var(--background) !important; /* Use main bg for hover */
        }

        /* --- Footer Section --- */
        .footer-section {
            background: transparent !important;
            border-top: 1px solid var(--divider) !important;
            padding: 1.5rem 1rem !important; /* Adjusted padding */
            margin-top: 1.5rem !important; /* Adjusted margin */
            text-align: center !important;
            color: var(--text-secondary) !important;
            font-size: 0.85rem !important; /* Adjusted font size */
            line-height: 1.5 !important;
        }
        .footer-section strong {
            color: var(--text-primary);
            font-weight: 500;
        }
        .footer-section a {
            color: var(--primary-color);
            text-decoration: none;
            font-weight: 500;
        }
        .footer-section a:hover {
            color: var(--primary-hover);
            text-decoration: underline;
        }

        /* --- Accessibility & Focus --- */
        :focus-visible { /* Standard focus visibility */
            outline: 2px solid var(--primary-color) !important;
            outline-offset: 2px;
        }
        /* Remove custom box-shadow focus for inputs/selects if :focus-visible is preferred */
        .gradio-textbox textarea:focus,
        .gradio-dropdown select:focus,
        .gradio-textbox input[type=password]:focus {
            border-color: var(--primary-color) !important;
            box-shadow: 0 0 0 3px var(--focus-ring) !important; /* Keep this for consistent focus */
            outline: none !important;
        }
        .gradio-button span:focus { /* Remove Gradio's default focus on button text */
            outline: none !important;
        }


        /* --- Responsive Adjustments --- */
        @media (max-width: 768px) {
            body { font-size: 14px; }
            .gradio-container > .flex.flex-col {
                padding: 2rem 1rem !important;
                gap: 1.5rem !important;
            }
            .card-style {
                padding: 1.5rem !important;
                border-radius: 10px !important;
            }
            .header-section {
                padding: 2rem 1.5rem !important;
                border-radius: 10px !important;
            }
            .header-title { font-size: 1.8rem; }
            .header-tagline { font-size: 1rem; }
            .input-row {
                flex-direction: column;
                gap: 1rem;
            }
            .button-row { justify-content: center; }
            .intro-card h3, .input-form-card h3, .output-card .response-header, .examples-card .gr-examples-header {
                font-size: 1.2rem !important;
            }
        }
        @media (max-width: 480px) {
            body { font-size: 14px; } /* Keep 14px or adjust if too small */
            .gradio-container > .flex.flex-col {
                padding: 1.25rem 0.75rem !important;
                gap: 1.25rem !important;
            }
            .card-style {
                padding: 1rem !important;
                border-radius: 8px !important;
            }
            .header-section {
                padding: 1.5rem 1rem !important;
                border-radius: 8px !important;
            }
            .header-logo { font-size: 2rem; }
            .header-title { font-size: 1.5rem; }
            .header-tagline { font-size: 0.9rem; }

            .intro-card h3, .input-form-card h3, .output-card .response-header, .examples-card .gr-examples-header {
                font-size: 1.1rem !important;
            }
            .gradio-textbox textarea,
            .gradio-dropdown select,
            .gradio-textbox input[type=password] {
                font-size: 0.9rem !important;
                padding: 0.7rem 0.9rem !important;
            }
            .gradio-button {
                width: 100%;
                padding: 0.7rem 1.25rem !important;
                font-size: 0.9rem !important;
            }
            .button-row {
                flex-direction: column;
                gap: 0.5rem;
            }
            .examples-card .gr-examples-table th,
            .examples-card .gr-examples-table td {
                padding: 0.5rem 0.7rem !important;
                font-size: 0.85rem !important;
            }
        }

        /* --- Gradio Overrides --- */
        .gradio-container > .flex {
            gap: 2rem !important; /* Match main gap */
        }
        .gradio-markdown > *:first-child { margin-top: 0; }
        .gradio-markdown > *:last-child { margin-bottom: 0; }
        .gradio-dropdown,
        .gradio-textbox { /* Remove Gradio default borders/padding around components */
            border: none !important;
            padding: 0 !important;
            background: transparent !important;
        }
        """

        # --- Gradio Blocks Layout ---
        with gr.Blocks(theme=None, css=custom_css, title="Landlord-Tenant Rights Assistant") as demo:
            # Header Section
            with gr.Group(elem_classes="header-section"):
                gr.Markdown(
                    """
                    <span class="header-logo">⚖️</span>
                    <h1 class="header-title">Landlord-Tenant Rights Assistant</h1>
                    <p class="header-tagline">Empowering You with State-Specific Legal Insights</p>
                    """, elem_id="app-title"
                )

            # Introduction Section
            with gr.Group(elem_classes="card-style intro-card"):
                gr.Markdown(
                    """
                    <h3>Know Your Rights</h3>
                    <p>Navigate landlord-tenant laws with ease. Enter your <strong>OpenAI API key</strong>, select your state, and ask your question to get detailed, state-specific answers.</p>
                    <p>Don't have an API key? <a href='https://platform.openai.com/api-keys' target='_blank'>Get one free from OpenAI</a>.</p>
                    <p><strong>Disclaimer:</strong> This tool provides information only, not legal advice. For legal guidance, consult a licensed attorney.</p>
                    """, elem_id="app-description"
                )

            # Input Form Section
            with gr.Group(elem_classes="card-style input-form-card"):
                gr.Markdown("<h3>Ask Your Question</h3>", elem_id="form-heading")

                with gr.Column(elem_classes="input-field-group"):
                    api_key_input = gr.Textbox(
                        label="OpenAI API Key",
                        type="password",
                        placeholder="Enter your API key (e.g., sk-...)",
                        info="Required to process your query. Securely used per request, not stored.",
                        elem_id="api-key-input",
                        lines=1
                    )

                with gr.Row(elem_classes="input-row"):
                    with gr.Column(elem_classes="input-field"):
                        query_input = gr.Textbox(
                            label="Your Question",
                            placeholder="E.g., What are the rules for security deposit returns in my state?",
                            lines=4,
                            max_lines=8,
                            elem_id="query-input"
                        )
                    with gr.Column(elem_classes="input-field"):
                        state_input = gr.Dropdown(
                            label="Select State",
                            choices=dropdown_choices,
                            value=initial_value,
                            allow_custom_value=False,
                            elem_id="state-dropdown"
                        )

                with gr.Row(elem_classes="button-row"):
                    clear_button = gr.Button(
                        "Clear",
                        variant="secondary",
                        elem_id="clear-button",
                        elem_classes=["gr-button-secondary"]
                    )
                    submit_button = gr.Button(
                        "Submit Query",
                        variant="primary",
                        elem_id="submit-button",
                        elem_classes=["gr-button-primary"]
                    )

            # Output Section
            with gr.Group(elem_classes="card-style output-card"):
                with gr.Column():
                    output = gr.Markdown(
                        value="<div class='placeholder'>Your answer will appear here after submitting your query.</div>",
                        elem_id="output-content",
                        elem_classes="output-content-wrapper"
                    )

            # Example Questions Section
            if example_queries:
                with gr.Group(elem_classes="card-style examples-card"):
                    gr.Examples(
                        examples=example_queries,
                        inputs=[query_input, state_input],
                        label="Explore Sample Questions",
                        examples_per_page=6
                    )
            else:
                with gr.Group(elem_classes="card-style examples-card"):
                    gr.Markdown(
                        "<div class='placeholder'>Sample questions could not be loaded. Please ensure states are available.</div>"
                    )

            # Footer Section
            with gr.Group(elem_classes="footer-section"):
                gr.Markdown(
                    """
                    **Disclaimer**: This tool is for informational purposes only and does not constitute legal advice.
                    <br><br>
                    Developed by **Nischal Subedi**. Connect on <a href="https://www.linkedin.com/in/nischal1/" target="_blank">LinkedIn</a> or explore insights at <a href="https://datascientistinsights.substack.com/" target="_blank">Substack</a>.
                    """, elem_id="app-footer"
                )

            # --- Event Listeners ---
            submit_button.click(
                fn=query_interface_wrapper,
                inputs=[api_key_input, query_input, state_input],
                outputs=output,
                api_name="submit_query"
            )

            clear_button.click(
                fn=lambda: (
                    "", "", initial_value,
                    "<div class='placeholder'>Inputs cleared. Ready for your next question.</div>"
                ),
                inputs=[],
                outputs=[api_key_input, query_input, state_input, output]
            )

        logging.info("Refined Gradio interface created successfully.")
        return demo

# --- Main Execution Block ---
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, "data/tenant-landlord.pdf")
        DEFAULT_DB_PATH = os.path.join(SCRIPT_DIR, "data/chroma_db")

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

        os.makedirs(os.path.dirname(VECTOR_DB_PATH), exist_ok=True)
        os.makedirs(os.path.dirname(PDF_PATH), exist_ok=True)

        logging.info(f"Using PDF path: {PDF_PATH}")
        logging.info(f"Using Vector DB path: {VECTOR_DB_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 ---")
            print(f"The required PDF file ('{os.path.basename(PDF_PATH)}') was not found at:")
            print(f"  {PDF_PATH}")
            print(f"Please ensure the file exists or set 'PDF_PATH' environment variable.")
            print(f"---------------------------\n")
            exit(1)

        logging.info("Initializing Vector Database...")
        vector_db_instance = VectorDatabase(persist_directory=VECTOR_DB_PATH)
        logging.info("Initializing RAG System...")
        rag = RAGSystem(vector_db=vector_db_instance)

        logging.info(f"Loading/Verifying data from PDF: {PDF_PATH}")
        states_loaded_count = rag.load_pdf(PDF_PATH)
        doc_count = vector_db_instance.document_collection.count() if vector_db_instance.document_collection else 0
        state_count = vector_db_instance.state_collection.count() if vector_db_instance.state_collection else 0
        total_items = doc_count + state_count

        if total_items > 0:
            logging.info(f"Data loading/verification complete. Vector DB contains {total_items} items. Found {states_loaded_count} distinct states.")
        else:
            logging.warning("Potential issue: PDF processed but Vector DB appears empty. Check PDF content/format and logs.")
            print("\nWarning: No data loaded from PDF or found in DB. Application might not function correctly.\n")

        logging.info("Setting up Gradio interface...")
        app_interface = rag.gradio_interface()

        SERVER_PORT = 7860
        logging.info(f"Launching Gradio app on http://0.0.0.0:{SERVER_PORT}")
        print("\n--- Gradio App Running ---")
        print(f"Access the interface in your browser at: http://localhost:{SERVER_PORT} or http://<your-ip-address>:{SERVER_PORT}")
        print("--------------------------\n")
        app_interface.launch(
            server_name="0.0.0.0",
            server_port=SERVER_PORT,
            share=True # Set to False if you don't want public sharing links
        )

    except FileNotFoundError as fnf_error:
        logging.error(f"Initialization failed due to a missing file: {str(fnf_error)}", exc_info=True)
        print(f"\n--- STARTUP ERROR: File Not Found ---")
        print(f"{str(fnf_error)}")
        print(f"---------------------------------------\n")
        exit(1)
    except ImportError as import_error:
        logging.error(f"Import error: {str(import_error)}. Check dependencies.", exc_info=True)
        print(f"\n--- STARTUP ERROR: Missing Dependency ---")
        print(f"Import Error: {str(import_error)}")
        print(f"Please ensure required libraries are installed (e.g., pip install -r requirements.txt).")
        print(f"-----------------------------------------\n")
        exit(1)
    except RuntimeError as runtime_error:
        logging.error(f"A runtime error occurred during setup: {str(runtime_error)}", exc_info=True)
        print(f"\n--- STARTUP ERROR: Runtime Problem ---")
        print(f"Runtime Error: {str(runtime_error)}")
        print(f"Check logs for details, often related to data loading or DB setup.")
        print(f"--------------------------------------\n")
        exit(1)
    except Exception as e:
        logging.error(f"An unexpected error occurred during application startup: {str(e)}", exc_info=True)
        print(f"\n--- FATAL STARTUP ERROR ---")
        print(f"An unexpected error stopped the application: {str(e)}")
        print(f"Check logs for detailed traceback.")
        print(f"---------------------------\n")
        exit(1)