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."

    def format_llm_response_for_html(self, text: str) -> str:
        """
        Formats the LLM's raw text response for display in a Gradio HTML component,
        converting Markdown-like elements to HTML. Handles bolding, paragraphs, and lists.
        """
        if not text:
            return ""

        # 1. Convert markdown bolding (**text**) to HTML <b>text</b>
        formatted_text = re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', text)

        # Split the text into blocks based on two or more consecutive newlines
        # This separates paragraphs and distinct list blocks.
        blocks = re.split(r'\n\s*\n+', formatted_text.strip())
        
        html_output_parts = []

        for block in blocks:
            block = block.strip()
            if not block:
                continue

            # Check if the block is a list (starts with a list marker on its first non-empty line)
            list_lines = block.split('\n')
            
            is_list = False
            list_type = None # 'ul' or 'ol'
            
            for line in list_lines:
                stripped_line = line.strip()
                if stripped_line: # Check first non-empty line
                    if re.match(r'^[*-]\s*(.*)', stripped_line):
                        is_list = True
                        list_type = 'ul'
                    elif re.match(r'^\d+\.\s*(.*)', stripped_line):
                        is_list = True
                        list_type = 'ol'
                    break # Only need to check the first actual line

            if is_list:
                # Process as a list
                list_html = f'<{list_type}>'
                for line in list_lines:
                    stripped_line = line.strip()
                    if not stripped_line:
                        continue # Skip empty lines within a list block
                    
                    # Extract content after the list marker
                    if list_type == 'ul':
                        item_content = re.sub(r'^[*-]\s*', '', stripped_line)
                    else: # ol
                        item_content = re.sub(r'^\d+\.\s*', '', stripped_line)
                    
                    list_html += f'<li>{item_content.strip()}</li>'
                list_html += f'</{list_type}>'
                html_output_parts.append(list_html)
            else:
                # Process as a paragraph
                # Replace single newlines within the paragraph block with <br>
                paragraph_content = block.replace('\n', '<br>')
                html_output_parts.append(f'<p>{paragraph_content}</p>')
        
        return ''.join(html_output_parts)


    @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)} 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.")
                # Apply the formatting function here to convert Markdown-like syntax to HTML
                answer_text = self.format_llm_response_for_html(answer_text)

            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



    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'>Get yours here</a>.</div>"
            if not state or state is None or "Error" in state:
                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 using 'self'
            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 process_query
            if "<div class='error-message'>" in answer:
                return answer
            else:
                formatted_response_content = f"<div class='response-header'><span class='response-icon'>📜</span>Response for {state}</div><div class='response-divider'></div><div class='response-content'>{answer}</div>"
                return f"<div class='animated-output-content'>{formatted_response_content}</div>"

        try:
            available_states_list = self.get_states()
            print(f"DEBUG: States loaded for selection: {available_states_list}")
            radio_choices = available_states_list if available_states_list and "Error" not in available_states_list[0] else ["Error: States unavailable"]
            initial_value_radio = radio_choices[0] if radio_choices and "Error" not in radio_choices[0] else None
        except Exception as e:
            print(f"DEBUG: Error loading states for selection: {e}")
            radio_choices = ["Error: Critical failure loading states"]
            initial_value_radio = None

        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 radio_choices and "Error" not in radio_choices[0] and len(radio_choices) > 0:
            loaded_states_set = set(radio_choices)
            example_queries = [ex for ex in example_queries_base if ex[1] in loaded_states_set]
            if not example_queries and radio_choices[0] != "Error: States unavailable":
                example_queries.append(["What basic rights do tenants have?", radio_choices[0]])
        else:
            example_queries.append(["What basic rights do tenants have?", "California"])

        custom_css = """
        /* Import modern fonts */
        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=Poppins:wght@600;700;800&display=swap');

        /* CSS Variables for theme consistency */
        :root {
            /* Light theme colors */
            --primary-color: #FF8C00;
            --primary-hover: #E07B00;
            --primary-light: rgba(255, 140, 0, 0.1);
            --primary-border: rgba(255, 140, 0, 0.3);
            
            /* Backgrounds */
            --bg-primary: #ffffff;
            --bg-secondary: #f8f9fa;
            --bg-tertiary: #f1f3f4;
            --bg-card: #ffffff;
            
            /* Text colors */
            --text-primary: #1a1a1a;
            --text-secondary: #6c757d;
            --text-muted: #adb5bd;
            
            /* Border colors */
            --border-light: #e9ecef;
            --border-medium: #dee2e6;
            --border-strong: #adb5bd;
            
            /* Status colors */
            --success-bg: #d1edff;
            --success-border: #0ea5e9;
            --success-text: #0369a1;
            --error-bg: #fef2f2;
            --error-border: #fca5a5;
            --error-text: #dc2626;
            --warning-bg: #fffbeb;
            --warning-border: #fbbf24;
            --warning-text: #d97706;
            
            /* Shadows */
            --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
            --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
            --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
            --shadow-xl: 0 20px 25px -5px rgba(0, 0, 0, 0.1), 0 10px 10px -5px rgba(0, 0, 0, 0.04);
        }

        /* Dark theme colors */
        @media (prefers-color-scheme: dark) {
            :root {
                --bg-primary: #1a1a1a;
                --bg-secondary: #2d2d2d;
                --bg-tertiary: #404040;
                --bg-card: #2d2d2d;
                
                --text-primary: #ffffff;
                --text-secondary: #a1a1aa;
                --text-muted: #71717a;
                
                --border-light: #404040;
                --border-medium: #525252;
                --border-strong: #737373;
                
                --success-bg: rgba(14, 165, 233, 0.1);
                --error-bg: rgba(220, 38, 38, 0.1);
                --warning-bg: rgba(217, 119, 6, 0.1);
            }
        }

        /* Force dark theme when .dark class is present */
        .dark {
            --bg-primary: #1a1a1a !important;
            --bg-secondary: #2d2d2d !important;
            --bg-tertiary: #404040 !important;
            --bg-card: #2d2d2d !important;
            
            --text-primary: #ffffff !important;
            --text-secondary: #a1a1aa !important;
            --text-muted: #71717a !important;
            
            --border-light: #404040 !important;
            --border-medium: #525252 !important;
            --border-strong: #737373 !important;
            
            --success-bg: rgba(14, 165, 233, 0.1) !important;
            --error-bg: rgba(220, 38, 38, 0.1) !important;
            --warning-bg: rgba(217, 119, 6, 0.1) !important;
        }

        /* Base styles */
        body, html {
            font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
            background-color: var(--bg-secondary) !important;
            color: var(--text-primary) !important;
            line-height: 1.6;
        }

        .gradio-container {
            background-color: var(--bg-secondary) !important;
            color: var(--text-primary) !important;
            max-width: 1200px !important;
            margin: 0 auto !important;
            padding: 1rem !important;
        }

        /* Header styles */
        .app-header {
            text-align: center;
            padding: 2rem 1rem;
            background: linear-gradient(135deg, var(--bg-card) 0%, var(--bg-tertiary) 100%);
            border: 1px solid var(--border-light);
            border-radius: 16px;
            margin-bottom: 2rem;
            box-shadow: var(--shadow-md);
        }

        .app-header-icon {
            font-size: 3rem;
            margin-bottom: 1rem;
            display: block;
        }

        .app-header-title {
            font-family: 'Poppins', sans-serif;
            font-size: 2.5rem;
            font-weight: 700;
            color: var(--primary-color) !important;
            margin: 0 0 0.5rem 0;
            text-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
            letter-spacing: -0.02em;
        }

        .app-header-tagline {
            font-size: 1.1rem;
            color: var(--text-secondary);
            margin: 0;
            font-weight: 400;
        }

        /* Card styles */
        .dashboard-card {
            background-color: var(--bg-card);
            border: 1px solid var(--border-light);
            border-radius: 12px;
            margin-bottom: 1.5rem;
            box-shadow: var(--shadow-sm);
            overflow: hidden;
            transition: all 0.2s ease;
        }

        .dashboard-card:hover {
            box-shadow: var(--shadow-md);
            transform: translateY(-1px);
        }

        .card-header {
            background: linear-gradient(135deg, var(--primary-color), var(--primary-hover));
            color: white;
            padding: 1rem 1.5rem;
            border-bottom: none;
        }

        .card-title {
            font-family: 'Poppins', sans-serif;
            font-size: 1.25rem;
            font-weight: 600;
            margin: 0;
            text-align: center;
        }

        .card-content {
            padding: 1.5rem;
            background-color: var(--bg-card);
        }

        /* Input styles */
        .input-group {
            margin-bottom: 1.5rem;
        }

        .input-group:last-child {
            margin-bottom: 0;
        }

        .gradio-textbox, .gradio-radio {
            margin-bottom: 1rem;
        }

        .gradio-textbox label, .gradio-radio legend {
            font-weight: 600;
            color: var(--text-primary) !important;
            margin-bottom: 0.5rem;
            display: block;
        }

        .gradio-textbox textarea, .gradio-textbox input {
            background-color: var(--bg-primary) !important;
            border: 2px solid var(--border-medium) !important;
            border-radius: 8px !important;
            color: var(--text-primary) !important;
            font-size: 0.95rem !important;
            padding: 0.75rem !important;
            transition: all 0.2s ease !important;
            box-shadow: var(--shadow-sm) !important;
        }

        .gradio-textbox textarea:focus, .gradio-textbox input:focus {
            border-color: var(--primary-color) !important;
            box-shadow: 0 0 0 3px var(--primary-light) !important;
            outline: none !important;
        }

        /* Radio button styles */
        .gradio-radio .wrap {
            display: grid !important;
            grid-template-columns: repeat(auto-fit, minmax(120px, 1fr)) !important;
            gap: 0.5rem !important;
            margin-top: 0.5rem !important;
        }

        .gradio-radio label {
            background-color: var(--bg-primary) !important;
            border: 2px solid var(--border-medium) !important;
            border-radius: 8px !important;
            padding: 0.75rem 1rem !important;
            cursor: pointer !important;
            transition: all 0.2s ease !important;
            text-align: center !important;
            font-weight: 500 !important;
            color: var(--text-primary) !important;
            box-shadow: var(--shadow-sm) !important;
        }

        .gradio-radio label:hover {
            border-color: var(--primary-color) !important;
            background-color: var(--primary-light) !important;
            transform: translateY(-1px) !important;
            box-shadow: var(--shadow-md) !important;
        }

        .gradio-radio input[type="radio"]:checked + label {
            background-color: var(--primary-color) !important;
            border-color: var(--primary-hover) !important;
            color: white !important;
            box-shadow: var(--shadow-md) !important;
        }

        /* Button styles */
        .button-row {
            display: flex;
            gap: 1rem;
            margin-top: 1.5rem;
            justify-content: center;
        }

        .gradio-button {
            padding: 0.75rem 2rem !important;
            border-radius: 8px !important;
            font-weight: 600 !important;
            font-size: 0.95rem !important;
            transition: all 0.2s ease !important;
            border: 2px solid transparent !important;
            cursor: pointer !important;
            min-width: 120px !important;
        }

        .gradio-button.primary {
            background: linear-gradient(135deg, var(--primary-color), var(--primary-hover)) !important;
            color: white !important;
            box-shadow: var(--shadow-md) !important;
        }

        .gradio-button.primary:hover {
            transform: translateY(-2px) !important;
            box-shadow: var(--shadow-lg) !important;
        }

        .gradio-button.secondary {
            background-color: transparent !important;
            color: var(--text-primary) !important;
            border-color: var(--border-medium) !important;
        }

        .gradio-button.secondary:hover {
            background-color: var(--bg-tertiary) !important;
            border-color: var(--primary-color) !important;
            transform: translateY(-1px) !important;
        }

        /* Output styles */
        .output-container {
            min-height: 200px;
            background-color: var(--bg-primary);
            border: 2px solid var(--border-light);
            border-radius: 8px;
            padding: 1.5rem;
            box-shadow: var(--shadow-sm);
        }

        .placeholder {
            text-align: center;
            color: var(--text-muted);
            font-style: italic;
            padding: 2rem;
            border: 2px dashed var(--border-medium);
            border-radius: 8px;
            background-color: var(--bg-tertiary);
        }

        .animated-output-content {
            animation: fadeInUp 0.3s ease-out;
        }

        @keyframes fadeInUp {
            from {
                opacity: 0;
                transform: translateY(10px);
            }
            to {
                opacity: 1;
                transform: translateY(0);
            }
        }

        .response-header {
            display: flex;
            align-items: center;
            gap: 0.5rem;
            font-weight: 600;
            color: var(--primary-color);
            margin-bottom: 1rem;
            font-size: 1.1rem;
        }

        .response-divider {
            height: 2px;
            background: linear-gradient(90deg, var(--primary-color), transparent);
            margin: 1rem 0;
            border-radius: 1px;
        }

        .response-content {
            color: var(--text-primary);
            line-height: 1.7;
        }
        /* Custom CSS for lists and paragraphs within response-content for better spacing */
        .response-content p {
            margin-bottom: 1em; /* Spacing between paragraphs */
        }
        .response-content ul, .response-content ol {
            margin-bottom: 1em; /* Spacing after lists */
            padding-left: 2em; /* Indent lists */
        }
        .response-content li {
            margin-bottom: 0.5em; /* Spacing between list items */
        }
        .response-content li:last-child {
            margin-bottom: 0;
        }


        /* Error message styles */
        .error-message {
            background-color: var(--error-bg);
            border: 2px solid var(--error-border);
            border-radius: 8px;
            padding: 1rem;
            color: var(--error-text);
            display: flex;
            align-items: center;
            gap: 0.5rem;
            font-weight: 500;
        }

        .error-icon {
            font-size: 1.2rem;
        }

        .error-message a {
            color: var(--error-text);
            text-decoration: underline;
            font-weight: 600;
        }

        /* Examples section */
        .examples-section {
            margin-top: 1rem;
        }

        .examples-section .gradio-examples {
            border: 1px solid var(--border-light);
            border-radius: 8px;
            overflow: hidden;
            box-shadow: var(--shadow-sm);
        }

        .examples-section table {
            width: 100%;
            border-collapse: collapse;
            background-color: var(--bg-card);
        }

        .examples-section th {
            background-color: var(--bg-tertiary);
            color: var(--text-primary);
            font-weight: 600;
            padding: 1rem;
            text-align: left;
            border-bottom: 2px solid var(--border-medium);
        }

        .examples-section td {
            padding: 1rem;
            border-bottom: 1px solid var(--border-light);
            color: var(--text-primary);
            cursor: pointer;
            transition: all 0.2s ease;
        }

        .examples-section tr:hover td {
            background-color: var(--primary-light);
            transform: translateX(4px);
        }

        /* Footer styles */
        .app-footer {
            margin-top: 3rem;
            padding: 2rem;
            background: linear-gradient(135deg, var(--bg-card) 0%, var(--bg-tertiary) 100%);
            border: 1px solid var(--border-light);
            border-radius: 12px;
            text-align: center;
            box-shadow: var(--shadow-sm);
        }

        .app-footer p {
            margin: 0.5rem 0;
            color: var(--text-secondary);
            line-height: 1.6;
        }

        .app-footer strong {
            color: var(--primary-color);
        }

        .app-footer a {
            color: var(--primary-color);
            text-decoration: none;
            font-weight: 600;
            transition: all 0.2s ease;
        }

        .app-footer a:hover {
            text-decoration: underline;
            color: var(--primary-hover);
        }

        /* Responsive design */
        @media (max-width: 768px) {
            .gradio-container {
                padding: 0.5rem !important;
            }
            
            .app-header {
                padding: 1.5rem 1rem;
            }
            
            .app-header-title {
                font-size: 2rem;
            }
            
            .card-content {
                padding: 1rem;
            }
            
            .button-row {
                flex-direction: column;
            }
            
            .gradio-button {
                width: 100% !important;
            }
            
            .gradio-radio .wrap {
                grid-template-columns: 1fr !important;
            }
        }

        @media (max-width: 480px) {
            .app-header-title {
                font-size: 1.8rem;
            }
            
            .app-header-tagline {
                font-size: 1rem;
            }
            
            .card-header {
                padding: 0.75rem 1rem;
            }
            
            .card-title {
                font-size: 1.1rem;
            }
        }

        /* Accessibility improvements */
        .gradio-button:focus,
        .gradio-textbox input:focus,
        .gradio-textbox textarea:focus,
        .gradio-radio label:focus {
            outline: 2px solid var(--primary-color);
            outline-offset: 2px;
        }

        /* Animation for loading states */
        .loading {
            opacity: 0.7;
            pointer-events: none;
        }

        .loading::after {
            content: '';
            position: absolute;
            top: 50%;
            left: 50%;
            width: 20px;
            height: 20px;
            margin: -10px 0 0 -10px;
            border: 2px solid var(--primary-color);
            border-radius: 50%;
            border-top-color: transparent;
            animation: spin 1s linear infinite;
        }

        @keyframes spin {
            to {
                transform: rotate(360deg);
            }
        }
        """

        with gr.Blocks(css=custom_css, title="Landlord-Tenant Rights Assistant", theme="default") as demo:
            # Header Section
            with gr.Group(elem_classes="app-header"):
                gr.HTML("""
                    <div class="app-header-icon">⚖️</div>
                    <h1 class="app-header-title">Landlord-Tenant Rights Assistant</h1>
                    <p class="app-header-tagline">Empowering You with State-Specific Legal Insights</p>
                """)

            # How It Works Card
            with gr.Group(elem_classes="dashboard-card"):
                gr.HTML('<div class="card-header"><h3 class="card-title">How This Assistant Works</h3></div>')
                with gr.Column(elem_classes="card-content"):
                    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 based on current legal statutes and precedents.
                    """)

            # API Key Input Card
            with gr.Group(elem_classes="dashboard-card"):
                gr.HTML('<div class="card-header"><h3 class="card-title">🔑 OpenAI API Key</h3></div>')
                with gr.Column(elem_classes="card-content"):
                    api_key_input = gr.Textbox(
                        label="API Key",
                        type="password",
                        placeholder="Enter your OpenAI API key (e.g., sk-...)",
                        lines=1,
                        elem_classes=["input-group"]
                    )
                    gr.Markdown("""
                    **Required to process your query.** Get your API key from [OpenAI Platform](https://platform.openai.com/api-keys).
                    Your key is processed securely and not stored.
                    """)

            # Query Input Card
            with gr.Group(elem_classes="dashboard-card"):
                gr.HTML('<div class="card-header"><h3 class="card-title">💬 Ask Your Question</h3></div>')
                with gr.Column(elem_classes="card-content"):
                    with gr.Row():
                        with gr.Column(scale=2):
                            query_input = gr.Textbox(
                                label="Your Question",
                                placeholder="E.g., What are the rules for security deposit returns in my state?",
                                lines=6,
                                max_lines=10,
                                elem_classes=["input-group"]
                            )
                        with gr.Column(scale=1):
                            state_input = gr.Radio(
                                label="Select State",
                                choices=radio_choices,
                                value=initial_value_radio,
                                elem_classes=["input-group"],
                                interactive=True
                            )
                    
                    with gr.Row(elem_classes="button-row"):
                        clear_button = gr.Button("Clear", variant="secondary", elem_classes=["secondary"])
                        submit_button = gr.Button("Submit Query", variant="primary", elem_classes=["primary"])

            # Output Card
            with gr.Group(elem_classes="dashboard-card"):
                gr.HTML('<div class="card-header"><h3 class="card-title">📋 Legal Assistant\'s Response</h3></div>')
                with gr.Column(elem_classes="card-content"):
                    output = gr.HTML(
                        value='<div class="placeholder">💡 The answer will appear here after submitting your query.</div>',
                        elem_classes=["output-container"]
                    )

            # Examples Card
            with gr.Group(elem_classes="dashboard-card examples-section"):
                gr.HTML('<div class="card-header"><h3 class="card-title">💡 Example Questions</h3></div>')
                with gr.Column(elem_classes="card-content"):
                    if example_queries:
                        gr.Examples(
                            examples=example_queries,
                            inputs=[query_input, state_input],
                            examples_per_page=5,
                            label=""
                        )
                    else:
                        gr.HTML('<div class="placeholder">⚠️ Sample questions could not be loaded. Please ensure the vector database is populated.</div>')

            # Footer
            with gr.Group(elem_classes="app-footer"):
                gr.HTML("""
                    <p><strong>⚠️ Disclaimer:</strong> 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 <strong>Nischal Subedi</strong>. 
                    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>.</p>
                """)

            # Event handlers
            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_radio, 
                    '<div class="placeholder">✨ Inputs cleared. Ready for your next question.</div>'
                ),
                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)