app.py

"""
SafeGuard AI - Women's Legal Rights Assistant (Pakistan)
"""

import os
import hashlib
import sqlite3
from io import BytesIO
from datetime import datetime
from pathlib import Path

import streamlit as st
import streamlit.components.v1 as components
from dotenv import load_dotenv

load_dotenv()

try:
    from google import genai
    from google.genai import types
    from PIL import Image

    from reportlab.lib.pagesizes import A4
    from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
    from reportlab.lib.units import inch
    from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
    from reportlab.lib import colors
    from reportlab.lib.enums import TA_CENTER
except ImportError:
    st.error("Install: pip install google-genai streamlit pillow python-dotenv reportlab")
    st.stop()


# -----------------------------------------------------------------------------
# DATABASE SETUP
# -----------------------------------------------------------------------------
DB_PATH = Path("safeguard_evidence.db")


def init_database():
    """Initialize SQLite database and create tables if they don't exist"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS evidence_timeline (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            logged_at TEXT NOT NULL,
            incident_date TEXT NOT NULL,
            incident_time TEXT NOT NULL,
            incident_type TEXT NOT NULL,
            description TEXT NOT NULL,
            witnesses TEXT,
            location TEXT,
            evidence_files TEXT
        )
    """)
    
    conn.commit()
    conn.close()


def save_evidence_to_db(entry):
    """Save evidence entry to database"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # Serialize evidence_files to JSON string
    import json
    evidence_files_json = json.dumps(entry.get("evidence_files", []))
    
    cursor.execute("""
        INSERT INTO evidence_timeline 
        (logged_at, incident_date, incident_time, incident_type, description, witnesses, location, evidence_files)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        entry["logged_at"],
        entry["incident_date"],
        entry["incident_time"],
        entry["incident_type"],
        entry["description"],
        entry.get("witnesses", ""),
        entry.get("location", ""),
        evidence_files_json
    ))
    
    entry_id = cursor.lastrowid
    conn.commit()
    conn.close()
    
    return entry_id


def load_evidence_from_db():
    """Load all evidence entries from database"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("""
        SELECT id, logged_at, incident_date, incident_time, incident_type, 
               description, witnesses, location, evidence_files
        FROM evidence_timeline
        ORDER BY incident_date DESC, incident_time DESC
    """)
    
    rows = cursor.fetchall()
    conn.close()
    
    # Convert to list of dictionaries
    import json
    evidence_list = []
    for row in rows:
        evidence_list.append({
            "id": row[0],
            "logged_at": row[1],
            "incident_date": row[2],
            "incident_time": row[3],
            "incident_type": row[4],
            "description": row[5],
            "witnesses": row[6] or "",
            "location": row[7] or "",
            "evidence_files": json.loads(row[8]) if row[8] else []
        })
    
    return evidence_list


def delete_evidence_from_db(entry_id):
    """Delete an evidence entry from database"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("DELETE FROM evidence_timeline WHERE id = ?", (entry_id,))
    
    conn.commit()
    conn.close()


# -----------------------------------------------------------------------------
# OPTIONALS
# -----------------------------------------------------------------------------
SHOW_EVIDENCE_UPLOAD_CARD = True 


# -----------------------------------------------------------------------------
# MODELS
# -----------------------------------------------------------------------------
MODEL_CANDIDATES = ["gemini-3-flash-preview", "gemini-3-pro-preview"]


def _model_variants(name: str):
    return [name, f"models/{name}"]


# -----------------------------------------------------------------------------
# CONFIG
# -----------------------------------------------------------------------------
st.set_page_config(
    page_title="SafeGuard AI - Women's Legal Rights Assistant",
    page_icon="🛡️",
    layout="wide",
    initial_sidebar_state="expanded",
)

# -----------------------------------------------------------------------------
# IMPROVED CSS with better contrast and browser compatibility
# -----------------------------------------------------------------------------
st.markdown(
    """
<style>
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=Poppins:wght@500;600;700;800&display=swap');
* { font-family: Inter, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }

#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
header {visibility: hidden;}
.stDeployButton {display: none;}
div[data-testid="stNotification"] {display: none !important;}
.stAlert {display:none !important;}

/* Force color scheme to light */
:root {
  color-scheme: light !important;
}

/* Page - Force white background */
[data-testid="stAppViewContainer"] { 
  background: #f7f7fb !important; 
  color-scheme: light !important;
}

[data-testid="stAppViewContainer"] > div {
  background: #f7f7fb !important;
}

.main {
  background: #f7f7fb !important;
}

.block-container{ 
  max-width: 1500px; 
  padding-top: 1.0rem; 
  padding-bottom: 1.0rem;
  background: #f7f7fb !important;
}

/* Sidebar */
[data-testid="stSidebar"] { 
  min-width: 310px !important; 
  max-width: 310px !important; 
}

[data-testid="stSidebar"]{
  background: linear-gradient(180deg, #2b0f55 0%, #35126a 35%, #4c1d95 75%, #5b21b6 100%) !important;
  color-scheme: dark !important;
}

[data-testid="stSidebar"] > div:first-child { 
  padding: 1.35rem 1.05rem 1.0rem 1.05rem; 
}

/* Sidebar header */
.sg-sidebar-header{
  display:flex; gap: 12px; align-items:flex-start;
  padding-bottom: 1.05rem;
  border-bottom: 1px solid rgba(255,255,255,0.12);
  margin-bottom: 1rem;
}

.sg-sidebar-logo{
  width: 46px; height: 46px; border-radius: 14px;
  background: rgba(255,255,255,0.12);
  display:flex; align-items:center; justify-content:center;
  font-size: 20px; color: #fff !important;
}

.sg-sidebar-title{
  font-family: Poppins, Inter, sans-serif;
  color: #fff !important; 
  font-size: 18px; 
  font-weight: 800; 
  line-height: 1.15; 
  margin:0;
}

.sg-sidebar-subtitle{
  color: rgba(255,255,255,0.70) !important; 
  font-size: 12px; 
  margin: 2px 0 0 0; 
  line-height: 1.25;
}

/* Sidebar - Force all text to white */
[data-testid="stSidebar"] * {
  color: #ffffff !important;
}

[data-testid="stSidebar"] p,
[data-testid="stSidebar"] span,
[data-testid="stSidebar"] div,
[data-testid="stSidebar"] label,
[data-testid="stSidebar"] h1,
[data-testid="stSidebar"] h2,
[data-testid="stSidebar"] h3 {
  color: #ffffff !important;
}

/* Sidebar navigation (radio -> cards) */
[data-testid="stSidebar"] .stRadio > label { display:none !important; }

[data-testid="stSidebar"] div[role="radiogroup"]{
  display:flex !important;
  flex-direction:column !important;
  gap: 12px !important;
}

[data-testid="stSidebar"] div[role="radiogroup"] label{
  background: rgba(255,255,255,0.10) !important;
  border: 1px solid rgba(255,255,255,0.12) !important;
  border-radius: 14px !important;
  padding: 14px 14px !important;
  cursor: pointer !important;
  transition: all .18s ease !important;
}

[data-testid="stSidebar"] div[role="radiogroup"] label:hover{
  background: rgba(255,255,255,0.14) !important;
  transform: translateY(-1px);
}

[data-testid="stSidebar"] div[role="radiogroup"] label[aria-checked="true"]{
  background: rgba(255,255,255,0.20) !important;
  border-color: rgba(255,255,255,0.22) !important;
}

/* Hide the radio circle */
[data-testid="stSidebar"] div[role="radiogroup"] label > div:first-child{
  display:none !important;
}

/* Make label text white and allow newline */
[data-testid="stSidebar"] div[role="radiogroup"] label span{
  white-space: pre-line !important;
  color: rgba(255,255,255,0.94) !important;
  font-size: 14px !important;
  font-weight: 700 !important;
  line-height: 1.25 !important;
}

/* Sidebar chat uploader hidden offscreen (keep functional) */
[data-testid="stSidebar"] div[data-testid="stFileUploader"]{
  position: fixed !important;
  left: -9999px !important;
  top: -9999px !important;
  width: 1px !important;
  height: 1px !important;
  opacity: 0 !important;
  overflow: hidden !important;
}

/* Privacy notice */
.sg-privacy{
  margin-top: 18px;
  background: rgba(255,255,255,0.10);
  border: 1px solid rgba(255,255,255,0.12);
  border-radius: 14px;
  padding: 14px;
}

.sg-privacy strong{ 
  display:block; 
  color:#fff !important; 
  font-size: 12px; 
  font-weight: 800; 
  margin-bottom: 4px; 
}

.sg-privacy p{ 
  margin:0; 
  color: rgba(255,255,255,0.82) !important; 
  font-size: 11px; 
  line-height: 1.45; 
}

/* Top header */
.sg-topline { 
  height: 1px; 
  background: rgba(91,33,182,0.12); 
  margin: .55rem 0 .95rem 0; 
}

.secure-badge{
  display:inline-flex; 
  align-items:center; 
  gap: 12px;
  background: #ffffff !important;
  border: 1px solid rgba(47,21,94,0.12);
  padding: 8px 14px; 
  border-radius: 999px;
  font-size: 13px; 
  color: #2f155e !important; 
  font-weight: 700;
}

.secure-badge .muted{ 
  opacity:.65; 
  font-weight: 600; 
  color: #2f155e !important;
}

/* Quick Exit */
.stButton button[kind="primary"]{
  background: #ef4444 !important;
  border: 1px solid #ef4444 !important;
  color: #fff !important;
  border-radius: 999px !important;
  padding: 10px 18px !important;
  font-weight: 800 !important;
  font-size: 13px !important;
}

/* Export PDF button */
.stDownloadButton button{
  background: #2f155e !important;
  color: #fff !important;
  border: 1px solid #2f155e !important;
  border-radius: 999px !important;
  padding: 10px 18px !important;
  font-weight: 800 !important;
  font-size: 13px !important;
}

.stDownloadButton button:hover{
  background: #3a1a73 !important;
  border-color: #3a1a73 !important;
}

/* Welcome */
.sg-welcome{
  text-align:center;
  padding: 3.3rem 1.5rem 1.6rem 1.5rem;
  max-width: 980px;
  margin: 0 auto;
  background: #f7f7fb !important;
}

.sg-welcome-icon{
  width: 66px; 
  height: 66px; 
  border-radius: 18px;
  background: #2f155e !important;
  display:flex; 
  align-items:center; 
  justify-content:center;
  color:#fff !important; 
  font-size: 28px;
  margin: 0 auto 1.2rem auto;
}

.sg-welcome-title{
  font-family: Poppins, Inter, sans-serif;
  font-size: 30px; 
  font-weight: 800; 
  color:#1f1147 !important;
  margin: 0 0 10px 0;
}

.sg-welcome-sub{
  color: rgba(47,21,94,.62) !important;
  font-size: 15px;
  line-height: 1.6;
  margin: 0 auto 18px auto;
  max-width: 720px;
}

/* Quick action pills */
div[data-testid="column"] .stButton button{
  background: #ffffff !important;
  color: #2f155e !important;
  border: 1px solid rgba(47,21,94,0.12) !important;
  border-radius: 999px !important;
  padding: 10px 16px !important;
  font-size: 13px !important;
  font-weight: 700 !important;
}

div[data-testid="column"] .stButton button:hover{
  background: #f9f9fb !important;
  border-color: rgba(47,21,94,0.20) !important;
}

/* Chat messages */
div[data-testid="stChatMessage"]{
  border-radius: 16px !important;
  padding: 1rem 1.1rem !important;
  margin: .7rem 0 !important;
  border: none !important;
  box-shadow: 0 8px 18px rgba(17,24,39,.06);
}

/* AI Assistant messages - Purple box with white text */
div[data-testid="stChatMessage"]:has(div[data-testid="chatAvatarIcon-assistant"]){
  background: linear-gradient(135deg, #5b21b6 0%, #7c3aed 100%) !important;
}

div[data-testid="stChatMessage"]:has(div[data-testid="chatAvatarIcon-assistant"]) [data-testid="stChatMessageContent"]{
  color: #ffffff !important;
  font-size: 14px !important;
  line-height: 1.65 !important;
}

div[data-testid="stChatMessage"]:has(div[data-testid="chatAvatarIcon-assistant"]) [data-testid="stChatMessageContent"] *{
  color: #ffffff !important;
}

/* User messages - Light purple-pink gradient with black text */
div[data-testid="stChatMessage"]:has(div[data-testid="chatAvatarIcon-user"]){
  background: linear-gradient(135deg, #fae8ff 0%, #fce7f3 50%, #ffe4e6 100%) !important;
  border: 1px solid rgba(147,51,234,0.15) !important;
}

div[data-testid="stChatMessage"]:has(div[data-testid="chatAvatarIcon-user"]) [data-testid="stChatMessageContent"]{
  color: #1f1147 !important;
  font-size: 14px !important;
  line-height: 1.65 !important;
}

div[data-testid="stChatMessage"]:has(div[data-testid="chatAvatarIcon-user"]) [data-testid="stChatMessageContent"] *{
  color: #1f1147 !important;
}

/* Avatars */
[data-testid="chatAvatarIcon-assistant"]{
  background: #2b0f55 !important;
  color: #fff !important;
  border-radius: 12px !important;
  width: 34px !important;
  height: 34px !important;
  font-weight: 800 !important;
}

[data-testid="chatAvatarIcon-user"]{
  background: #ef4444 !important;
  color: #fff !important;
  border-radius: 999px !important;
  width: 34px !important;
  height: 34px !important;
  font-weight: 800 !important;
}

.message-time{ 
  font-size: 11px; 
  opacity: .65; 
  margin-top: 8px; 
}

/* Chat input */
div[data-testid="stChatInput"]{
  background: #ffffff !important;
  border: 1px solid rgba(47,21,94,0.18) !important;
  border-radius: 16px !important;
  box-shadow: 0 10px 20px rgba(17,24,39,.06) !important;
  position: relative;
}

div[data-testid="stChatInput"] textarea{
  padding-left: 52px !important;
  font-size: 14px !important;
  color: #1f1147 !important;
  background: #ffffff !important;
}

div[data-testid="stChatInput"] textarea::placeholder{
  color: rgba(47,21,94,0.45) !important;
}

/* Paperclip label injected into chat input */
.sg-clip-label{
  position: absolute;
  left: 14px;
  bottom: 12px;
  width: 34px;
  height: 34px;
  border-radius: 10px;
  border: 1px solid rgba(47,21,94,0.14);
  background: rgba(255,255,255,0.96) !important;
  color: rgba(47,21,94,.75) !important;
  display:flex;
  align-items:center;
  justify-content:center;
  cursor:pointer;
  user-select:none;
  z-index: 6000;
}

.sg-clip-label:hover{ 
  border-color: rgba(47,21,94,0.22); 
  background: #ffffff !important;
}

/* Evidence upload card */
.sg-evidence-card{
  width: min(820px, 100%);
  margin: 18px auto 18px auto;
  border-radius: 18px;
  border: 2px dashed rgba(47,21,94,0.20);
  background: #fbfbfe !important;
  padding: 34px 26px;
  position: relative;
  text-align:center;
}

.sg-evidence-icon{
  width: 52px; 
  height: 52px; 
  border-radius: 16px;
  margin: 0 auto 14px auto;
  background: rgba(91,33,182,0.14) !important;
  display:flex; 
  align-items:center; 
  justify-content:center;
  color: #2f155e !important; 
  font-size: 22px;
}

.sg-evidence-title{ 
  font-family: Poppins, Inter, sans-serif; 
  color:#1f1147 !important; 
  font-size:18px; 
  font-weight:800; 
}

.sg-evidence-sub{ 
  color: rgba(47,21,94,.62) !important; 
  font-size:13px; 
  margin-top:6px; 
}

.sg-evidence-tags{ 
  display:flex; 
  gap:16px; 
  justify-content:center; 
  flex-wrap:wrap; 
  margin-top:10px; 
  color: rgba(47,21,94,.70) !important; 
  font-weight:600; 
}

/* Evidence overlay uploader */
.sg-evidence-uploader{
  position:absolute !important;
  inset: 0 !important;
  opacity: 0 !important;
  z-index: 5 !important;
}

.sg-evidence-uploader section{ height:100% !important; }
.sg-evidence-uploader [data-testid="stFileUploaderDropzone"]{ height: 100% !important; }

/* Resource cards */
.resource-card{
  background: #ffffff !important;
  border: 1px solid rgba(47,21,94,0.10);
  border-radius: 16px;
  padding: 16px 18px;
  margin: 10px 0;
  position: relative;
}

.resource-card::after{
  content: '↗';
  position:absolute;
  right: 16px;
  top: 16px;
  color: rgba(47,21,94,0.35) !important;
  font-size: 16px;
}

.resource-title{ 
  font-weight: 800; 
  color:#1f1147 !important; 
  font-size: 15px; 
  margin-bottom: 4px; 
}

.resource-desc{ 
  color: rgba(47,21,94,.62) !important; 
  font-size: 13px; 
  line-height: 1.55; 
  padding-right: 24px; 
}

.section-h{ 
  margin: 22px 0 10px 0; 
  font-weight: 900; 
  color: #1f1147 !important; 
  display:flex; 
  align-items:center; 
  gap: 8px; 
}

/* Disclaimer */
.disclaimer{
  margin-top: 14px;
  background: rgba(255, 243, 199, 0.85) !important;
  border: 1px solid rgba(245,158,11,.25);
  border-left: 4px solid #f59e0b;
  padding: 12px 14px;
  border-radius: 14px;
  color: #92400e !important;
  font-size: 12px;
  line-height: 1.5;
}

/* Forms and inputs */
.stTextInput input,
.stTextArea textarea,
.stSelectbox select,
.stDateInput input,
.stTimeInput input {
  background: #ffffff !important;
  color: #1f1147 !important;
  border-color: rgba(47,21,94,0.18) !important;
}

.stTextInput label,
.stTextArea label,
.stSelectbox label,
.stDateInput label,
.stTimeInput label {
  color: #1f1147 !important;
}

/* Success/Info messages */
.stSuccess {
  background: #d1fae5 !important;
  color: #065f46 !important;
}

.stInfo {
  background: #dbeafe !important;
  color: #1e40af !important;
}

.stError {
  background: #fee2e2 !important;
  color: #991b1b !important;
}

/* Expander */
.streamlit-expanderHeader {
  background: #ffffff !important;
  color: #1f1147 !important;
  border-color: rgba(47,21,94,0.12) !important;
}

.streamlit-expanderContent {
  background: #ffffff !important;
}

/* Delete button */
.stButton button[kind="secondary"]{
  background: transparent !important;
  border: 1px solid #ef4444 !important;
  color: #ef4444 !important;
  border-radius: 8px !important;
  padding: 6px 12px !important;
  font-weight: 600 !important;
  font-size: 12px !important;
}

.stButton button[kind="secondary"]:hover{
  background: #fee2e2 !important;
}
</style>
""",
    unsafe_allow_html=True,
)


# -----------------------------------------------------------------------------
# SYSTEM PROMPT
# -----------------------------------------------------------------------------
SYSTEM_INSTRUCTION = """You are SafeGuard AI, a legal assistant for women's rights in Pakistan.

RULES:
1. Start with: "⚠️ SafeGuard AI is not a lawyer. It is a first step toward safety, awareness, and documentation."
2. If danger: "🚨 IMMEDIATE DANGER: Call 1099 or Police 15"
3. Be supportive and trauma-informed.

PAKISTANI LAWS:
- Protection Against Harassment of Women at Workplace Act 2010
- Domestic Violence Acts (Sindh 2012, Punjab 2016)
- Prevention of Electronic Crimes Act (PECA) 2016
"""


# -----------------------------------------------------------------------------
# STATE
# -----------------------------------------------------------------------------
def init_state():
    if "nav_page" not in st.session_state:
        st.session_state.nav_page = "chat"

    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "message_translations" not in st.session_state:
        st.session_state.message_translations = {}

    if "uploaded_image" not in st.session_state:
        st.session_state.uploaded_image = None
    if "show_full_image" not in st.session_state:
        st.session_state.show_full_image = {}

    if "evidence_timeline" not in st.session_state:
        # Load from database on startup
        st.session_state.evidence_timeline = load_evidence_from_db()

    if "_quick_exit" not in st.session_state:
        st.session_state._quick_exit = False

    if "_pdf_cache" not in st.session_state:
        st.session_state._pdf_cache = None
    if "_pdf_cache_hash" not in st.session_state:
        st.session_state._pdf_cache_hash = None

    if "_toast_incident_saved" not in st.session_state:
        st.session_state._toast_incident_saved = None

    if "_clear_chat_attach" not in st.session_state:
        st.session_state._clear_chat_attach = False


# -----------------------------------------------------------------------------
# GEMINI
# -----------------------------------------------------------------------------
def initialize_gemini():
    # Try Streamlit secrets first (for cloud deployment)
    try:
        api_key = st.secrets["GOOGLE_API_KEY"]
    except (KeyError, FileNotFoundError, AttributeError):
        # Fallback to .env (for local development)
        api_key = os.getenv("GOOGLE_API_KEY")
    
    if not api_key:
        st.error("🔑 Add GOOGLE_API_KEY to Streamlit secrets or .env file")
        st.stop()
    return genai.Client(api_key=api_key)

def get_gemini_response(client, user_message, image_bytes=None):
    """
    Uses only:
    - gemini-3-flash-preview
    - gemini-3-pro-preview
    Tries both bare and 'models/<name>' variants for compatibility.
    """
    try:
        parts = []
        if user_message:
            parts.append(types.Part(text=user_message))
        if image_bytes:
            parts.append(types.Part(inline_data=types.Blob(mime_type="image/jpeg", data=image_bytes)))

        contents = [types.Content(role="user", parts=parts)]
        config = types.GenerateContentConfig(
            system_instruction=SYSTEM_INSTRUCTION,
            temperature=0.7,
            max_output_tokens=2000,
        )

        last_error = None
        for name in MODEL_CANDIDATES:
            for model_name in _model_variants(name):
                try:
                    resp = client.models.generate_content(
                        model=model_name,
                        contents=contents,
                        config=config,
                    )
                    return resp.text, None
                except Exception as e:
                    last_error = e
                    continue

        return None, f"Model error: {last_error}"
    except Exception as e:
        return None, str(e)


def translate_to_urdu(client, text):
    try:
        prompt = (
            "Translate the following text to Urdu. "
            "Keep legal terms accurate and understandable:\n\n"
            f"{text}"
        )
        contents = [types.Content(role="user", parts=[types.Part(text=prompt)])]
        config = types.GenerateContentConfig(temperature=0.3, max_output_tokens=2048)

        last_error = None
        for name in MODEL_CANDIDATES:
            for model_name in _model_variants(name):
                try:
                    resp = client.models.generate_content(
                        model=model_name,
                        contents=contents,
                        config=config,
                    )
                    return resp.text
                except Exception as e:
                    last_error = e
                    continue

        return f"Translation unavailable. Model error: {last_error}"
    except Exception:
        return "Translation unavailable at the moment."


def process_image(uploaded_file):
    """Returns (jpeg_bytes, PIL_image_for_display)"""
    try:
        img = Image.open(uploaded_file)
        if img.mode != "RGB":
            img = img.convert("RGB")
        img.thumbnail((1920, 1920), Image.Resampling.LANCZOS)

        buf = BytesIO()
        img.save(buf, format="JPEG", quality=85)
        return buf.getvalue(), img
    except Exception as e:
        st.error(f"Image error: {e}")
        return None, None


# -----------------------------------------------------------------------------
# EVIDENCE / PDF
# -----------------------------------------------------------------------------
def serialize_uploaded_files(uploaded_files):
    if not uploaded_files:
        return []
    return [
        {"name": getattr(f, "name", "file"), "type": getattr(f, "type", ""), "size": getattr(f, "size", 0)}
        for f in uploaded_files
    ]


def add_evidence_entry(incident_date, incident_time, incident_type, description,
                       evidence_files=None, witnesses="", location=""):
    entry = {
        "logged_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "incident_date": incident_date.strftime("%Y-%m-%d"),
        "incident_time": incident_time.strftime("%H:%M"),
        "incident_type": incident_type,
        "description": description,
        "witnesses": witnesses,
        "location": location,
        "evidence_files": serialize_uploaded_files(evidence_files),
    }
    
    # Save to database and get the ID
    entry_id = save_evidence_to_db(entry)
    entry["id"] = entry_id
    
    # Add to session state (for current session)
    st.session_state.evidence_timeline.append(entry)
    
    return entry


def evidence_hash():
    parts = []
    for e in st.session_state.evidence_timeline:
        files = ",".join([f.get("name", "") for f in (e.get("evidence_files") or [])])
        parts.append(
            f"{e.get('incident_date')}|{e.get('incident_time')}|{e.get('incident_type')}|"
            f"{e.get('location')}|{e.get('witnesses')}|{e.get('description')}|{files}"
        )
    return hashlib.md5("\n".join(parts).encode("utf-8", errors="ignore")).hexdigest()


def generate_pdf_bytes():
    if not st.session_state.evidence_timeline:
        return None

    buf = BytesIO()
    doc = SimpleDocTemplate(buf, pagesize=A4)
    elements = []
    styles = getSampleStyleSheet()

    title_style = ParagraphStyle("Title", parent=styles["Heading1"], fontSize=22, alignment=TA_CENTER, spaceAfter=24)
    elements.append(Paragraph("EVIDENCE TIMELINE REPORT", title_style))
    elements.append(Spacer(1, 0.2 * inch))

    metadata = [
        ["Generated:", datetime.now().strftime("%Y-%m-%d %H:%M:%S")],
        ["Incidents:", str(len(st.session_state.evidence_timeline))],
    ]
    metadata_table = Table(metadata, colWidths=[2 * inch, 4 * inch])
    metadata_table.setStyle(TableStyle([("GRID", (0, 0), (-1, -1), 0.8, colors.grey)]))
    elements.append(metadata_table)
    elements.append(Spacer(1, 0.35 * inch))

    for idx, entry in enumerate(
        sorted(st.session_state.evidence_timeline, key=lambda x: (x["incident_date"], x["incident_time"])), 1
    ):
        file_names = ", ".join([f.get("name", "") for f in (entry.get("evidence_files") or [])]) or "N/A"
        entry_data = [
            ["#:", str(entry.get("id", idx))],
            ["Date:", f"{entry['incident_date']} at {entry['incident_time']}"],
            ["Type:", entry["incident_type"]],
            ["Location:", entry["location"] or "N/A"],
            ["Witnesses:", entry["witnesses"] or "N/A"],
            ["Evidence files:", file_names],
            ["Description:", entry["description"]],
        ]
        entry_table = Table(entry_data, colWidths=[1.5 * inch, 4.5 * inch])
        entry_table.setStyle(TableStyle([("GRID", (0, 0), (-1, -1), 0.5, colors.grey)]))
        elements.append(entry_table)
        elements.append(Spacer(1, 0.18 * inch))

    doc.build(elements)
    buf.seek(0)
    return buf.getvalue()


def get_cached_pdf():
    h = evidence_hash()
    if st.session_state._pdf_cache is None or st.session_state._pdf_cache_hash != h:
        st.session_state._pdf_cache = generate_pdf_bytes()
        st.session_state._pdf_cache_hash = h
    return st.session_state._pdf_cache


# -----------------------------------------------------------------------------
# JS helpers
# -----------------------------------------------------------------------------
def inject_chat_paperclip_label_js():
    """Inject a <label> into chat input linked to sidebar file input."""
    components.html(
        """
<script>
(function(){
  function run(){
    const doc = window.parent.document;

    const chatInput = doc.querySelector('div[data-testid="stChatInput"]');
    if(!chatInput) return false;

    const fileInput = doc.querySelector('[data-testid="stSidebar"] input[type="file"]:not([multiple])');
    if(!fileInput) return false;

    if(!fileInput.id) fileInput.id = "sg_chat_file_input";
    if(chatInput.querySelector('.sg-clip-label')) return true;

    const lbl = doc.createElement('label');
    lbl.className = 'sg-clip-label';
    lbl.setAttribute('for', fileInput.id);
    lbl.title = 'Attach image';
    lbl.textContent = '📎';

    chatInput.appendChild(lbl);
    return true;
  }

  let tries=0;
  const timer=setInterval(()=>{
    tries++;
    const ok = run();
    if(ok || tries>30) clearInterval(timer);
  }, 120);
})();
</script>
        """,
        height=0,
    )


def inject_evidence_uploader_overlay_js():
    """Overlay the multiple-uploader on the evidence card (optional UI)."""
    components.html(
        """
<script>
(function(){
  function run(){
    const doc = window.parent.document;
    const card = doc.getElementById('sg-evidence-card');
    if (!card) return false;

    const fileInput = doc.querySelector('input[type="file"][multiple]');
    if (!fileInput) return false;

    const uploader = fileInput.closest('div[data-testid="stFileUploader"]');
    if (!uploader) return false;

    const wrap = uploader.closest('.element-container') || uploader;
    if (wrap.classList.contains('sg-evidence-uploader')) return true;

    wrap.classList.add('sg-evidence-uploader');
    card.appendChild(wrap);
    return true;
  }

  let tries=0;
  const timer=setInterval(()=>{
    tries++;
    const ok = run();
    if(ok || tries>30) clearInterval(timer);
  }, 120);
})();
</script>
        """,
        height=0,
    )


# -----------------------------------------------------------------------------
# Topbar helpers
# -----------------------------------------------------------------------------
def render_quick_exit_redirect_if_needed():
    if st.session_state.get("_quick_exit"):
        components.html("""<script>window.top.location.href="https://www.google.com";</script>""", height=0)
        st.stop()


def render_top_bar():
    c1, c2 = st.columns([8, 1.8])
    with c1:
        st.markdown(
            """
            <div class="secure-badge">
              <span style="display:inline-flex; gap:8px; align-items:center;">
                <span style="font-size:14px;">🔒</span>
                <span>Secure Session</span>
              </span>
              <span class="muted">End-to-end encrypted</span>
            </div>
            """,
            unsafe_allow_html=True,
        )
    with c2:
        if st.button("✕  Quick Exit", type="primary", use_container_width=True):
            st.session_state._quick_exit = True
            st.rerun()
    st.markdown('<div class="sg-topline"></div>', unsafe_allow_html=True)


# -----------------------------------------------------------------------------
# Sidebar (stable navigation)
# -----------------------------------------------------------------------------
def render_sidebar() -> str:
    with st.sidebar:
        st.markdown(
            """
            <div class="sg-sidebar-header">
              <div class="sg-sidebar-logo">🛡️</div>
              <div>
                <div class="sg-sidebar-title">SafeGuard AI</div>
                <div class="sg-sidebar-subtitle">Legal Rights Assistant</div>
              </div>
            </div>
            """,
            unsafe_allow_html=True,
        )

        # Safe clear BEFORE instantiating uploader widget
        if st.session_state._clear_chat_attach:
            st.session_state.pop("chat_attach", None)
            st.session_state.uploaded_image = None
            st.session_state._clear_chat_attach = False

        def nav_label(key: str) -> str:
            if key == "chat":
                return "💬  Chat Assistant\nAI-powered legal guidance"
            if key == "evidence":
                return "🕐  Evidence Timeline\nTrack your documentation"
            return "📚  Legal Resources\nLaws and rights library"

        page_key = st.radio(
            "Navigation",
            options=["chat", "evidence", "resources"],
            format_func=nav_label,
            key="nav_page",
            label_visibility="collapsed",
        )

        # Chat file uploader in sidebar (hidden offscreen by CSS)
        if page_key == "chat":
            uploaded = st.file_uploader(
                "Attach image",
                type=["jpg", "jpeg", "png"],
                key="chat_attach",
                label_visibility="collapsed",
            )
            if uploaded:
                st.session_state.uploaded_image = uploaded

        st.markdown(
            """
            <div class="sg-privacy">
              <strong>🔒 Privacy First</strong>
              <p>Your conversations are confidential and encrypted. Your data is stored securely in a local database.</p>
            </div>
            """,
            unsafe_allow_html=True,
        )

    return page_key


# -----------------------------------------------------------------------------
# Pages
# -----------------------------------------------------------------------------
def page_chat(client):
    if len(st.session_state.messages) == 0:
        st.markdown(
            """
            <div class="sg-welcome">
              <div class="sg-welcome-icon">✨</div>
              <div class="sg-welcome-title">Welcome to SafeGuard AI</div>
              <div class="sg-welcome-sub">
                Your confidential AI assistant for understanding legal rights and protections for women in Pakistan.
              </div>
            </div>
            """,
            unsafe_allow_html=True,
        )

        b1, b2, b3 = st.columns(3)
        with b1:
            if st.button("What are my rights in a divorce?", use_container_width=True):
                st.session_state.messages.append({"role": "user", "content": "What are my rights in a divorce?"})
                with st.spinner("Thinking..."):
                    response, error = get_gemini_response(client, "What are my rights in a divorce?")
                st.session_state.messages.append({"role": "assistant", "content": response if not error else f"⚠️ {error}"})
                st.rerun()
        with b2:
            if st.button("How to report domestic violence?", use_container_width=True):
                st.session_state.messages.append({"role": "user", "content": "How to report domestic violence?"})
                with st.spinner("Thinking..."):
                    response, error = get_gemini_response(client, "How to report domestic violence?")
                st.session_state.messages.append({"role": "assistant", "content": response if not error else f"⚠️ {error}"})
                st.rerun()
        with b3:
            if st.button("Workplace harassment laws", use_container_width=True):
                st.session_state.messages.append({"role": "user", "content": "Workplace harassment laws"})
                with st.spinner("Thinking..."):
                    response, error = get_gemini_response(client, "Workplace harassment laws")
                st.session_state.messages.append({"role": "assistant", "content": response if not error else f"⚠️ {error}"})
                st.rerun()

        st.markdown("<br>", unsafe_allow_html=True)
        with st.chat_message("assistant"):
            st.markdown(
                "Assalam-o-Alaikum! I am your SafeGuard AI assistant. I am here to help you understand your legal rights "
                "under Pakistani law. Everything you share is completely confidential."
            )
            st.markdown('<div class="message-time">Just now</div>', unsafe_allow_html=True)

        with st.chat_message("assistant"):
            st.markdown(
                "You can ask me about domestic violence laws, workplace harassment protections, property rights, "
                "marriage and divorce laws, or any other legal matter. How can I assist you today?"
            )
            st.markdown('<div class="message-time">Just now</div>', unsafe_allow_html=True)

    # render messages
    for idx, msg in enumerate(st.session_state.messages):
        with st.chat_message(msg["role"]):
            msg_id = f"msg_{idx}"
            current_lang = st.session_state.message_translations.get(msg_id, "english")

            if msg["role"] == "assistant" and current_lang == "urdu":
                st.markdown(st.session_state.message_translations.get(f"{msg_id}_urdu", msg["content"]))
            else:
                st.markdown(msg["content"])

            if msg.get("image"):
                st.image(msg["image"], width=220)

            if msg["role"] == "assistant":
                c1, c2, _ = st.columns([1, 1, 7])
                with c1:
                    if st.button("English", key=f"en_{idx}"):
                        st.session_state.message_translations[msg_id] = "english"
                        st.rerun()
                with c2:
                    if st.button("اردو", key=f"ur_{idx}"):
                        if f"{msg_id}_urdu" not in st.session_state.message_translations:
                            with st.spinner("Translating..."):
                                st.session_state.message_translations[f"{msg_id}_urdu"] = translate_to_urdu(
                                    client, msg["content"]
                                )
                        st.session_state.message_translations[msg_id] = "urdu"
                        st.rerun()

            st.markdown('<div class="message-time">Just now</div>', unsafe_allow_html=True)

    # attachment preview + remove
    if st.session_state.uploaded_image:
        image_bytes, preview_img = process_image(st.session_state.uploaded_image)
        if preview_img:
            st.image(preview_img, caption="Attached image (will be sent with your next message)", width=170)
        if st.button("Remove attachment", key="remove_attach_btn"):
            st.session_state._clear_chat_attach = True
            st.rerun()

    prompt = st.chat_input("Type your question here...")
    inject_chat_paperclip_label_js()

    if prompt:
        if st.session_state.uploaded_image:
            image_bytes, display_img = process_image(st.session_state.uploaded_image)
            st.session_state.messages.append({"role": "user", "content": prompt, "image": display_img})

            with st.spinner("Analyzing..."):
                response, error = get_gemini_response(client, prompt, image_bytes=image_bytes)

            st.session_state.messages.append({"role": "assistant", "content": response if not error else f"⚠️ {error}"})
            st.session_state._clear_chat_attach = True
        else:
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.spinner("Thinking..."):
                response, error = get_gemini_response(client, prompt)
            st.session_state.messages.append({"role": "assistant", "content": response if not error else f"⚠️ {error}"})

        st.rerun()

    st.markdown(
        """
        <div class="disclaimer">
          SafeGuard AI provides legal information, not legal advice. Consult a qualified lawyer for case-specific guidance.
        </div>
        """,
        unsafe_allow_html=True,
    )


def page_evidence():
    # toast after save
    if st.session_state._toast_incident_saved:
        try:
            st.toast(st.session_state._toast_incident_saved, icon="✅")
        except Exception:
            st.success(st.session_state._toast_incident_saved)
        st.session_state._toast_incident_saved = None

    h1, h2 = st.columns([7, 2])
    with h1:
        st.markdown("<div class='section-h'>🕐 Evidence Timeline</div>", unsafe_allow_html=True)
        st.markdown(
            "<div style='color:rgba(47,21,94,.62);font-size:14px;margin-top:2px;'>Track and organize your documentation securely</div>",
            unsafe_allow_html=True
        )
    with h2:
        has_data = len(st.session_state.evidence_timeline) > 0
        pdf_bytes = get_cached_pdf() if has_data else b""
        st.download_button(
            "Export PDF",
            data=pdf_bytes,
            file_name=f"SafeGuard_Evidence_{datetime.now().strftime('%Y%m%d')}.pdf",
            mime="application/pdf",
            use_container_width=True,
            disabled=not has_data,
        )

    st.markdown("<br>", unsafe_allow_html=True)

    if SHOW_EVIDENCE_UPLOAD_CARD:
        st.markdown(
            """
            <div id="sg-evidence-card" class="sg-evidence-card">
              <div class="sg-evidence-icon">⤴</div>
              <div class="sg-evidence-title">Upload Evidence</div>
              <div class="sg-evidence-sub">Drag and drop files here, or click to browse</div>
              <div class="sg-evidence-tags">
                <span>🖼️ Images</span>
                <span>📄 Documents</span>
              </div>
            </div>
            """,
            unsafe_allow_html=True,
        )
        uploaded_files = st.file_uploader(
            "Upload Evidence",
            type=["jpg", "jpeg", "png", "pdf", "doc", "docx"],
            accept_multiple_files=True,
            label_visibility="collapsed",
            key="evidence_upload",
        )
        inject_evidence_uploader_overlay_js()
    else:
        uploaded_files = st.file_uploader(
            "Upload Evidence",
            type=["jpg", "jpeg", "png", "pdf", "doc", "docx"],
            accept_multiple_files=True,
            key="evidence_upload",
        )

    if uploaded_files:
        st.success(f"✅ {len(uploaded_files)} file(s) uploaded successfully!")

    with st.expander("➕ Log New Incident", expanded=False):
        with st.form("evidence_form"):
            c1, c2 = st.columns(2)
            with c1:
                incident_date = st.date_input("📅 Date", value=datetime.now())
                incident_time = st.time_input("🕐 Time")
                incident_type = st.selectbox(
                    "📝 Incident Type",
                    ["Workplace Harassment", "Domestic Violence", "Cyber Harassment", "Stalking", "Threat", "Assault", "Other"],
                )
            with c2:
                location = st.text_input("📍 Location")
                witnesses = st.text_input("👥 Witnesses")

            description = st.text_area("📄 Description", height=150, placeholder="Describe the incident in detail...")

            if st.form_submit_button("💾 Save Incident", use_container_width=True):
                if description.strip():
                    entry = add_evidence_entry(
                        incident_date, incident_time, incident_type, description,
                        uploaded_files, witnesses, location
                    )
                    st.session_state._pdf_cache = None
                    st.session_state._pdf_cache_hash = None
                    st.session_state._toast_incident_saved = f"Incident #{entry['id']} saved successfully!"
                    st.rerun()
                else:
                    st.error("⚠️ Description is required")

    st.markdown(
        f"<div class='section-h'>📊 Timeline <span style='opacity:.6;'>({len(st.session_state.evidence_timeline)} incidents)</span></div>",
        unsafe_allow_html=True,
    )

    if st.session_state.evidence_timeline:
        for entry in sorted(
            st.session_state.evidence_timeline,
            key=lambda x: (x["incident_date"], x["incident_time"]),
            reverse=True,
        ):
            files = entry.get("evidence_files") or []
            files_text = ", ".join([f.get("name", "") for f in files]) if files else ""
            
            with st.container():
                col1, col2 = st.columns([9, 1])
                with col1:
                    st.markdown(
                        f"""
                        <div class="resource-card">
                          <div class="resource-title">#{entry['id']} — {entry['incident_type']}</div>
                          <div class="resource-desc">📅 {entry['incident_date']} at {entry['incident_time']}</div>
                          <div class="resource-desc" style="margin-top:8px;">📝 {entry['description']}</div>
                          {"<div class='resource-desc' style='margin-top:6px;'>📍 " + entry['location'] + "</div>" if entry['location'] else ""}
                          {"<div class='resource-desc'>👥 " + entry['witnesses'] + "</div>" if entry['witnesses'] else ""}
                          {("<div class='resource-desc'>📎 " + files_text + "</div>") if files_text else ""}
                        </div>
                        """,
                        unsafe_allow_html=True,
                    )
                with col2:
                    if st.button("🗑️", key=f"del_{entry['id']}", type="secondary", help="Delete this incident"):
                        delete_evidence_from_db(entry['id'])
                        st.session_state.evidence_timeline = load_evidence_from_db()
                        st.session_state._pdf_cache = None
                        st.session_state._pdf_cache_hash = None
                        st.rerun()
    else:
        st.info("📝 No incidents logged yet. Use the form above to add your first entry.")


def page_resources():
    st.markdown(
        """
        <div style="display:flex; gap:14px; align-items:center; margin-top:.25rem;">
          <div class="sg-welcome-icon" style="width:54px; height:54px; font-size:22px;">📖</div>
          <div>
            <div style="font-family:Poppins;font-weight:800;font-size:28px;color:#1f1147;">Legal Resources</div>
            <div style="color:rgba(47,21,94,.62);font-size:14px;margin-top:2px;">Key laws, rights, and emergency contacts</div>
          </div>
        </div>
        """,
        unsafe_allow_html=True,
    )

    st.markdown('<div class="section-h">🤝 Support Organizations</div>', unsafe_allow_html=True)
    st.markdown(
        """
        <div class="resource-card">
          <div class="resource-title">Aurat Foundation</div>
          <div class="resource-desc"><b>Contact:</b> 051-2609596 (Islamabad)</div>
          <div class="resource-desc"><b>Services:</b> Women's rights advocacy and information.</div>
        </div>
        <div class="resource-card">
          <div class="resource-title">War Against Rape (WAR)</div>
          <div class="resource-desc"><b>Contact:</b> 021-35373008 (Sindh/Karachi)</div>
          <div class="resource-desc"><b>Services:</b> Legal aid and counseling for sexual assault survivors.</div>
        </div>
        <div class="resource-card">
          <div class="resource-title">Edhi Foundation</div>
          <div class="resource-desc"><b>Contact:</b> 115 (Emergency) or 051-2827844 (H-8 Islamabad Center)</div>
          <div class="resource-desc"><b>Services:</b> Emergency ambulance, shelter, and victim support.</div>
        </div>
        """,
        unsafe_allow_html=True,
    )

    st.markdown('<div class="section-h">📞 Emergency (General)</div>', unsafe_allow_html=True)
    st.markdown(
        """
        <div class="resource-card">
          <div class="resource-title">Police Emergency</div>
          <div class="resource-desc"><b>15</b></div>
        </div>
        <div class="resource-card">
          <div class="resource-title">Legal Aid / Human Rights (Govt)</div>
          <div class="resource-desc"><b>1099</b></div>
        </div>
        """,
        unsafe_allow_html=True,
    )


# -----------------------------------------------------------------------------
# MAIN
# -----------------------------------------------------------------------------
def main():
    # Initialize database first
    init_database()
    
    init_state()
    render_quick_exit_redirect_if_needed()

    page = render_sidebar()
    client = initialize_gemini()
    render_top_bar()

    if page == "chat":
        page_chat(client)
    elif page == "evidence":
        page_evidence()
    elif page == "resources":
        page_resources()


if __name__ == "__main__":
    main()