# app.py
import streamlit as st
import os
import time
import plotly.express as px
from groq import Groq
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize Groq client
try:
    client = Groq(api_key=os.getenv("GROQ_API_KEY"))
except:
    st.error("Failed to initialize Groq client. Please check your API key.")
    st.stop()

# Set up Streamlit page
st.set_page_config(
    page_title="Legal Rights Explorer Pro",
    page_icon="⚖️",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# Custom CSS for professional styling
st.markdown("""
<style>
    :root {
        --primary: #2c3e50;
        --secondary: #3498db;
        --accent: #1abc9c;
        --light: #ecf0f1;
        --dark: #2c3e50;
        --success: #27ae60;
        --card-shadow: 0 6px 20px rgba(0,0,0,0.08);
        --transition: all 0.3s ease;
    }
    
    .stApp {
        background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%);
        color: #333;
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
        max-width: 1200px;
        margin: 0 auto;
        padding: 0 20px;
    }
    
    .header {
        background: linear-gradient(90deg, var(--primary) 0%, var(--secondary) 100%);
        color: white;
        padding: 2.5rem 2rem;
        border-radius: 0 0 25px 25px;
        box-shadow: var(--card-shadow);
        margin-bottom: 2.5rem;
        text-align: center;
    }
    
    .card {
        background: white;
        border-radius: 18px;
        padding: 2rem;
        box-shadow: var(--card-shadow);
        margin-bottom: 2rem;
        transition: var(--transition);
        border-left: 5px solid var(--accent);
    }
    
    .card:hover {
        box-shadow: 0 12px 30px rgba(0,0,0,0.15);
    }
    
    .scenario-btn {
        width: 100%;
        padding: 1.2rem;
        background: linear-gradient(135deg, var(--secondary) 0%, #2980b9 100%);
        color: white;
        border: none;
        border-radius: 15px;
        font-weight: 600;
        margin: 0.8rem 0;
        cursor: pointer;
        transition: var(--transition);
        text-align: left;
        padding-left: 25px;
        font-size: 1.1rem;
        box-shadow: 0 4px 10px rgba(0,0,0,0.1);
    }
    
    .scenario-btn:hover {
        background: linear-gradient(135deg, var(--primary) 0%, #2c3e50 100%);
        transform: translateY(-3px);
        box-shadow: 0 8px 20px rgba(0,0,0,0.15);
    }
    
    .response-card {
        background: white;
        border-left: 5px solid var(--success);
        border-radius: 15px;
        padding: 2.5rem;
        margin-top: 2.5rem;
        box-shadow: var(--card-shadow);
        animation: fadeIn 0.8s ease;
    }
    
    .country-btn {
        display: flex;
        align-items: center;
        justify-content: center;
        gap: 10px;
        width: 100%;
        padding: 1.2rem;
        border-radius: 15px;
        background: white;
        box-shadow: var(--card-shadow);
        margin-bottom: 15px;
        cursor: pointer;
        transition: var(--transition);
        border: 2px solid transparent;
        font-weight: 600;
        font-size: 1.1rem;
    }
    
    .country-btn:hover {
        transform: translateY(-5px);
        box-shadow: 0 10px 25px rgba(0,0,0,0.15);
    }
    
    .country-btn.selected {
        background: linear-gradient(135deg, var(--light) 0%, #d6eaf8 100%);
        border: 2px solid var(--accent);
        box-shadow: 0 8px 25px rgba(0,0,0,0.12);
    }
    
    .footer {
        text-align: center;
        padding: 2.5rem;
        margin-top: 3.5rem;
        color: var(--dark);
        font-size: 1rem;
        background: rgba(236, 240, 241, 0.7);
        border-radius: 20px;
        box-shadow: var(--card-shadow);
    }
    
    .spinner {
        display: flex;
        justify-content: center;
        margin: 3rem 0;
    }
    
    @keyframes fadeIn {
        from { opacity: 0; transform: translateY(30px); }
        to { opacity: 1; transform: translateY(0); }
    }
    
    @keyframes pulse {
        0% { transform: scale(1); }
        50% { transform: scale(1.05); }
        100% { transform: scale(1); }
    }
    
    .pulse {
        animation: pulse 2s infinite;
    }
    
    @media (max-width: 992px) {
        .stApp {
            padding: 0 15px;
        }
        
        .header {
            padding: 2rem 1.5rem;
        }
        
        .card {
            padding: 1.8rem;
        }
        
        .scenario-btn, .country-btn {
            padding: 1.1rem;
            font-size: 1.05rem;
        }
    }
    
    @media (max-width: 768px) {
        .header {
            padding: 1.8rem 1.2rem;
            border-radius: 0 0 20px 20px;
        }
        
        .card {
            padding: 1.5rem;
        }
        
        .scenario-btn, .country-btn {
            padding: 1rem;
            font-size: 1rem;
        }
        
        .response-card {
            padding: 1.8rem;
        }
    }
    
    @media (max-width: 576px) {
        .stApp {
            padding: 0 10px;
        }
        
        .header h1 {
            font-size: 2.2rem !important;
        }
        
        .header p {
            font-size: 1.1rem !important;
        }
        
        .card {
            padding: 1.2rem;
        }
        
        .scenario-btn, .country-btn {
            padding: 0.9rem;
            font-size: 0.95rem;
        }
        
        .response-card {
            padding: 1.5rem;
        }
    }
</style>
""", unsafe_allow_html=True)

# Country data with flags
COUNTRIES = {
    "🇺🇸 United States": "US",
    "🇬🇧 United Kingdom": "UK",
    "🇨🇦 Canada": "CA",
    "🇦🇺 Australia": "AU",
    "🇮🇳 India": "IN",
    "🇩🇪 Germany": "DE",
    "🇫🇷 France": "FR",
    "🇯🇵 Japan": "JP",
    "🇧🇷 Brazil": "BR",
    "🇿🇦 South Africa": "ZA",
    "🇪🇸 Spain": "ES",
    "🇸🇬 Singapore": "SG",
    "🇲🇽 Mexico": "MX",
    "🇳🇱 Netherlands": "NL",
    "🇸🇪 Sweden": "SE",
    "🇳🇴 Norway": "NO"
}

# Common legal scenarios
SCENARIOS = [
    "Traffic Stop by Police",
    "Workplace Discrimination",
    "Rental Housing Issues",
    "Medical Emergency Rights",
    "Consumer Protection",
    "Arrest or Detention",
    "Border Crossing",
    "Online Privacy",
    "Employment Termination",
    "Debt Collection",
    "Freedom of Speech",
    "Immigration Rights"
]

# LLM models available on Groq
MODELS = {
    "Llama3-70b (Highest Accuracy)": "llama3-70b-8192",
    "Llama3-8b (Fast Response)": "llama3-8b-8192",
    "Mixtral-8x7b (Balanced)": "mixtral-8x7b-32768"
}

# Function to get rights information from Groq API
def get_legal_rights(country, scenario, model_name):
    """Get legal rights information using Groq API"""
    system_prompt = f"""
    You are an expert legal assistant specializing in {country} law. 
    Provide clear, accurate information about an individual's rights in the given scenario.
    
    Guidelines:
    - List 5-7 key rights as bullet points
    - Start with a brief context about the legal situation
    - Use plain language understandable to non-lawyers
    - Include relevant legal references when appropriate
    - Highlight any critical actions the person should take
    - Mention any country-specific variations
    - Keep response under 350 words
    - Format with emojis for readability
    - End with important disclaimers
    """
    
    user_prompt = f"""
    Scenario: {scenario}
    Country: {country}
    
    Please provide a comprehensive list of rights for this situation in {country}.
    """
    
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": user_prompt
                }
            ],
            model=model_name,
            temperature=0.3,
            max_tokens=1024,
            top_p=1,
            stream=False,
            stop=None,
        )
        
        return chat_completion.choices[0].message.content
    
    except Exception as e:
        st.error(f"Error fetching data: {str(e)}")
        return None

# Function to display response with animation
def display_response(response):
    """Display the response with typing animation effect"""
    message_placeholder = st.empty()
    full_response = ""
    
    # Simulate stream of response with milliseconds delay
    for chunk in response.split():
        full_response += chunk + " "
        time.sleep(0.03)
        # Add a blinking cursor to simulate typing
        message_placeholder.markdown(f'<div class="response-card">{full_response}▌</div>', unsafe_allow_html=True)
    
    # Display final message without the cursor
    message_placeholder.markdown(f'<div class="response-card">{response}</div>', unsafe_allow_html=True)

# Main app
def main():
    # Header section
    st.markdown("""
    <div class="header">
        <h1 style="margin:0;font-size:2.8rem;">Legal Rights Explorer Pro</h1>
        <p style="margin:0;font-size:1.3rem;opacity:0.9;margin-top:12px;">Know your rights • Protect your freedom</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Initialize session state
    if 'selected_country' not in st.session_state:
        st.session_state.selected_country = "🇺🇸 United States"
    
    if 'selected_scenario' not in st.session_state:
        st.session_state.selected_scenario = None
    
    # Model selection at top
    st.markdown("### ⚙️ AI Model Selection")
    model_cols = st.columns([1, 1, 2])
    with model_cols[0]:
        selected_model = st.selectbox("Choose AI Model", list(MODELS.keys()), index=0)
    
    with model_cols[1]:
        st.markdown("### 🌎 Selected Country")
        st.markdown(f"<div style='font-size:1.2rem;font-weight:bold;'>{st.session_state.selected_country}</div>", 
                    unsafe_allow_html=True)
    
    # Main content columns
    col1, col2 = st.columns([1, 1.2], gap="large")
    
    with col1:
        st.markdown("### 🌍 Select Your Country")
        st.markdown("Choose your country to see jurisdiction-specific rights information")
        
        # Create 2 columns for country buttons
        country_col1, country_col2 = st.columns(2)
        
        # Split countries into two groups
        countries_list = list(COUNTRIES.keys())
        mid_index = len(countries_list) // 2
        countries_col1 = countries_list[:mid_index]
        countries_col2 = countries_list[mid_index:]
        
        # First column of countries
        with country_col1:
            for country_display in countries_col1:
                is_selected = st.session_state.selected_country == country_display
                btn_class = "country-btn selected" if is_selected else "country-btn"
                
                if st.button(
                    country_display,
                    key=f"btn_{country_display}",
                    use_container_width=True
                ):
                    st.session_state.selected_country = country_display
                    st.session_state.selected_scenario = None
                    st.rerun()
        
        # Second column of countries
        with country_col2:
            for country_display in countries_col2:
                is_selected = st.session_state.selected_country == country_display
                btn_class = "country-btn selected" if is_selected else "country-btn"
                
                if st.button(
                    country_display,
                    key=f"btn_{country_display}",
                    use_container_width=True
                ):
                    st.session_state.selected_country = country_display
                    st.session_state.selected_scenario = None
                    st.rerun()
        
        # Country visualization
        st.markdown("### 📊 Global Rights Index")
        countries = list(COUNTRIES.values())
        awareness = [92, 89, 87, 85, 82, 88, 86, 84, 79, 81, 85, 83, 80, 78, 90, 85]  # Simulated data
        
        fig = px.bar(
            x=countries,
            y=awareness,
            labels={'x': 'Country', 'y': 'Rights Awareness'},
            color=awareness,
            color_continuous_scale='Teal',
            height=350
        )
        fig.update_layout(
            plot_bgcolor='rgba(0,0,0,0)',
            paper_bgcolor='rgba(0,0,0,0)',
            margin=dict(l=0, r=0, t=30, b=0),
            xaxis_tickangle=-45
        )
        st.plotly_chart(fig, use_container_width=True)
        
        # Legal resources
        with st.expander("📚 Legal Resources & References", expanded=True):
            st.markdown("""
            - [United Nations Human Rights](https://www.ohchr.org/)
            - [International Justice Resource Center](https://ijrcenter.org/)
            - [Global Legal Information Network](https://www.loc.gov/law/help/legal-info.php)
            - [World Legal Information Institute](https://www.worldlii.org/)
            - [Amnesty International](https://www.amnesty.org/)
            - [Human Rights Watch](https://www.hrw.org/)
            """)
    
    with col2:
        st.markdown("### 🚨 Select a Legal Scenario")
        st.markdown("Choose a situation to understand your rights and protections")
        
        # Create buttons for each scenario in 2 columns
        scenario_cols = st.columns(2)
        for i, scenario in enumerate(SCENARIOS):
            with scenario_cols[i % 2]:
                if st.button(
                    f"**{scenario}**",
                    key=f"scen_{scenario}",
                    use_container_width=True,
                    help=f"Click to see your rights for {scenario}"
                ):
                    st.session_state.selected_scenario = scenario
        
        # Custom scenario input
        st.markdown("#### ✏️ Or describe your specific situation:")
        custom_scenario = st.text_input(
            "Describe your legal concern:",
            placeholder="e.g., 'Rights during home search', 'Employee privacy rights'",
            label_visibility="collapsed"
        )
        if custom_scenario:
            st.session_state.selected_scenario = custom_scenario
        
        # Response area
        st.markdown("### ⚖️ Your Legal Rights Analysis")
        
        if st.session_state.selected_scenario:
            country_code = COUNTRIES[st.session_state.selected_country]
            
            with st.spinner(f"🔍 Analyzing your rights for '{st.session_state.selected_scenario}' in {country_code}..."):
                response = get_legal_rights(
                    country_code, 
                    st.session_state.selected_scenario,
                    MODELS[selected_model]
                )
                
                if response:
                    display_response(response)
                else:
                    st.error("Failed to get response. Please try again.")
        else:
            st.info("👆 Please select a legal scenario to see your rights information")
            st.markdown("""
            <div style="text-align:center; margin:20px 0;">
                <img src="https://images.unsplash.com/photo-1589391886645-d51941baf7fb?auto=format&fit=crop&w=600&h=400" 
                     style="width:100%; border-radius:15px; box-shadow:0 6px 20px rgba(0,0,0,0.1);">
                <p style="margin-top:10px;font-style:italic;color:#666;">Know Your Rights, Exercise Your Freedom</p>
            </div>
            """, unsafe_allow_html=True)
            
            st.markdown("""
            <div class="card">
                <h3>Why Knowing Your Rights Matters</h3>
                <p>Understanding your legal rights empowers you to:</p>
                <ul>
                    <li>Protect yourself in challenging legal situations</li>
                    <li>Make informed decisions when facing authorities</li>
                    <li>Prevent exploitation and rights violations</li>
                    <li>Access justice when your rights are infringed</li>
                    <li>Confidently navigate complex legal systems</li>
                </ul>
                <p style="margin-top:15px;font-style:italic;border-left:3px solid #3498db;padding-left:15px;">
                "The first defense against injustice is knowledge of your rights."
                </p>
            </div>
            """, unsafe_allow_html=True)
    
    # Footer
    st.markdown("""
    <div class="footer">
        <h4>Legal Rights Explorer Pro</h4>
        <p>© 2023 • For Educational Purposes Only • Not Legal Advice</p>
        <p style="font-size:0.9rem;margin-top:10px;">Always consult a qualified attorney in your jurisdiction for your specific legal situation</p>
        <p style="font-size:0.8rem;opacity:0.7;margin-top:15px;">Powered by Groq Cloud & Open-Source LLMs</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()