"""
NurseLex — Legal Literacy Agent for All Nurses and Nursing Students
Architecture:
  1. Local legislation.parquet — 219K health/social care Acts & SIs for browsing
  2. cached_legislation.py — 1,128 sections loaded from nursing_sections.json
  3. Gemini Flash REST API — Plain English explanations (with retry logic)
"""
import os
import asyncio
import httpx
import logging
import pandas as pd
import gradio as gr

from cached_legislation import search_cached
from local_search import search_scenarios_locally

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Load local legislation index ---
PARQUET_PATH = os.path.join(os.path.dirname(__file__), "legislation.parquet")
try:
    LEG_DF = pd.read_parquet(PARQUET_PATH)
    logger.info(f"Loaded {len(LEG_DF)} legislation entries from parquet")
except Exception as e:
    logger.warning(f"Could not load parquet: {e}")
    LEG_DF = pd.DataFrame()

# --- Key nursing legislation IDs ---
NURSING_ACTS = {
    "Mental Health Act 1983": "ukpga/1983/20",
    "Mental Capacity Act 2005": "ukpga/2005/9",
    "Care Act 2014": "ukpga/2014/23",
    "Human Rights Act 1998": "ukpga/1998/42",
    "Equality Act 2010": "ukpga/2010/15",
    "Health and Social Care Act 2012": "ukpga/2012/7",
    "Mental Health Units (Use of Force) Act 2018": "ukpga/2018/27",
    "Autism Act 2009": "ukpga/2009/15",
    "Children Act 1989": "ukpga/1989/41",
    "Children Act 2004": "ukpga/2004/31",
    "Safeguarding Vulnerable Groups Act 2006": "ukpga/2006/47",
    "Health and Care Act 2022": "ukpga/2022/31",
}
REVERSE_ACTS = {v: k for k, v in NURSING_ACTS.items()}

# --- Gemini REST API ---
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
GEMINI_BASE = "https://generativelanguage.googleapis.com/v1beta/models"
GEMINI_MODELS = ["gemini-2.0-flash-lite", "gemini-2.0-flash"]

SYSTEM_PROMPT = """You are NurseLex, a legal literacy assistant for all UK nurses and nursing students.

Your role:
1. Answer legal questions using ONLY the legislation text provided in the context.
2. Explain the law in clear, plain English suitable for all nurses and nursing students.
3. Always cite the specific Act, section number, and year.
4. If the context doesn't contain enough information, say so clearly.
5. Add practical nursing implications (e.g., "In practice, this means...").
6. Include professional reminders (e.g., NMC Code, duty of care).

Disclaimers to include:
- "This is for educational purposes only — always consult your trust's legal team for specific cases."
- "This reflects the legislation as written — local trust policies may add additional requirements."

Format with clear headings, bullet points, and bold key terms."""

QUICK_QUESTIONS = [
    "What is Section 5(4) of the Mental Health Act and when can a nurse use it?",
    "What does the Mental Capacity Act say about best interests decisions?",
    "When can a patient be detained under Section 2 vs Section 3?",
    "What are the legal requirements for using restraint?",
    "What does Section 117 aftercare mean and who is entitled?",
    "What are a nurse's legal duties under the Care Act 2014 for safeguarding?",
    "What is Deprivation of Liberty and when do DoLS apply?",
    "What rights does a patient have under Section 136?",
]

async def call_gemini(prompt: str) -> str:
    """Call Gemini via REST API with retry logic and model fallback."""
    if not GEMINI_API_KEY:
        return ""

    payload = {
        "system_instruction": {"parts": [{"text": SYSTEM_PROMPT}]},
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"temperature": 0.3, "maxOutputTokens": 2048},
    }

    async with httpx.AsyncClient(timeout=60.0) as client:
        for model in GEMINI_MODELS:
            url = f"{GEMINI_BASE}/{model}:generateContent?key={GEMINI_API_KEY}"
            for attempt in range(3):
                try:
                    resp = await client.post(url, json=payload)
                    if resp.status_code == 429:
                        wait = 2 ** (attempt + 1)
                        logger.warning(f"Rate limited ({model}), retrying in {wait}s")
                        await asyncio.sleep(wait)
                        continue
                    resp.raise_for_status()
                    data = resp.json()
                    return data["candidates"][0]["content"]["parts"][0]["text"]
                except httpx.HTTPStatusError as e:
                    if e.response.status_code == 429:
                        wait = 2 ** (attempt + 1)
                        logger.warning(f"Rate limited ({model}), retrying in {wait}s")
                        await asyncio.sleep(wait)
                        continue
                    logger.error(f"Gemini API error ({model}): {e.response.status_code}")
                    break
                except Exception as e:
                    logger.error(f"Gemini error ({model}): {type(e).__name__}")
                    break
            logger.info(f"Model {model} exhausted, trying next...")

    logger.error("All Gemini models failed")
    return ""

def search_legislation_index(query: str, max_results: int = 10) -> pd.DataFrame:
    """Search the full legislation index parquet by title."""
    if LEG_DF.empty:
        return pd.DataFrame()

    mask = LEG_DF["title"].str.contains(query, case=False, na=False)
    results = LEG_DF[mask].sort_values("year", ascending=False).head(max_results)
    return results

async def query_and_respond(user_question: str, history: list) -> str:
    """Main RAG pipeline: local cached sections (1,128) + Gemini explanation."""
    if not user_question.strip():
        return "Please enter a question about UK healthcare legislation."

    # Step 1: Search local legislation sections
    sections = search_cached(user_question, max_results=5)
    logger.info(f"Local search returned {len(sections)} sections for: {user_question[:60]}")

    # Step 2: Search parquet index for related Acts
    related_acts = search_legislation_index(user_question, max_results=5)

    # Step 3: Build context
    context_parts = []
    for section in sections:
        title = section.get("title", "Untitled")
        text = section.get("text", "")
        leg_id = section.get("legislation_id", "")
        num = section.get("number", "")
        context_parts.append(f"### {title}\n**Source:** {leg_id}, Section {num}\n\n{text}\n")

    context = "\n---\n".join(context_parts) if context_parts else "No matching legislation sections found in cache."

    # Step 4: Generate Gemini response
    prompt = f"## Nurse's Question\n{user_question}\n\n## Relevant UK Legislation\n{context}\n\nPlease answer the nurse's question using the legislation above."
    
    answer = await call_gemini(prompt)

    if not answer:
        # Fallback: show raw legislation if Gemini fails or is missing key
        answer = _build_fallback(user_question, sections)
        if not GEMINI_API_KEY:
            answer += "\n\n⚠️ *Set `GEMINI_API_KEY` in Space secrets for AI-powered plain English explanations.*"
        elif "rate limit" in answer.lower():
             answer += "\n\n⚠️ *Gemini is currently rate limited, falling back to raw legislation.*"

    # Add source citations
    source_acts = set()
    for s in sections:
        leg_id = s.get("legislation_id", "")
        if leg_id:
            source_acts.add(leg_id)

    if source_acts:
        answer += "\n\n---\n📚 **Sources:** "
        answer += " | ".join(f"[{sid}](https://www.legislation.gov.uk/id/{sid})" for sid in sorted(source_acts))

    # Add related Acts from parquet
    if not related_acts.empty:
        answer += "\n\n📖 **Related legislation:** "
        act_links = []
        for _, row in related_acts.head(3).iterrows():
            uri = row.get("uri", "")
            title = row.get("title", "")
            if uri and title:
                act_links.append(f"[{title}]({uri})")
        if act_links:
            answer += " | ".join(act_links)

    answer += "\n\n🏛️ *Data from [legislation.gov.uk](https://www.legislation.gov.uk/) — Crown Copyright, OGL v3.0*"
    return answer

def _build_fallback(question: str, sections: list) -> str:
    """Show raw legislation without LLM."""
    response = f"## Legislation relevant to: *{question}*\n\n"

    if not sections:
        response += (
            "No matching sections found in cache. Try searching the full **Browse Legislation** tab for the Act title, or try specific terms like:\n"
            "- **\"Section 5(4)\"** or **\"nurse holding power\"**\n"
            "- **\"best interests\"** or **\"capacity\"**\n"
            "- **\"safeguarding\"** or **\"Section 42\"**\n"
            "- **\"Section 136\"** or **\"place of safety\"**\n"
        )
        return response

    for i, section in enumerate(sections[:5], 1):
        title = section.get("title", "Untitled")
        text = section.get("text", "No text available")
        leg_id = section.get("legislation_id", "")
        num = section.get("number", "")
        uri = section.get("uri", "")

        response += f"### {i}. {title}\n"
        response += f"**Act:** `{leg_id}` | **Section:** {num}\n\n"
        response += f"{text}\n\n"
        if uri:
            response += f"🔗 [View on legislation.gov.uk]({uri})\n\n"
        response += "---\n\n"

    return response

async def section_lookup(act_name: str, section_input: str) -> str:
    """Look up sections from cached legislation."""
    legislation_id = NURSING_ACTS.get(act_name)
    if not legislation_id:
        return f"❌ Act not found in NurseLex."

    cache_query = f"{act_name} section {section_input}" if section_input.strip() else act_name
    sections = search_cached(cache_query, max_results=10)
    sections = [s for s in sections if s.get("legislation_id") == legislation_id]

    if section_input.strip() and sections:
        try:
            target_num = int(section_input.strip().replace("Section ", "").replace("s.", "").replace("S.", ""))
            matching = [s for s in sections if s.get("number") == target_num]
            if matching:
                sections = matching
        except ValueError:
            pass

    if not sections:
        return (
            f"⏳ Section not found in cache for **{act_name}**.\n\n"
            f"Try the **Chat tab** for a broader search, or visit "
            f"[legislation.gov.uk](https://www.legislation.gov.uk/id/{legislation_id}) directly."
        )

    result = f"## {act_name}\n\n"
    for section in sections[:5]:
        title = section.get("title", "Untitled")
        text = section.get("text", "No text")
        num = section.get("number", "")
        uri = section.get("uri", "")

        result += f"### Section {num}: {title}\n\n{text}\n\n"
        if uri:
            result += f"🔗 [View on legislation.gov.uk]({uri})\n\n"
        result += "---\n\n"

    result += "\n🏛️ *Crown Copyright, OGL v3.0*"
    return result

async def fetch_explanatory_note(act_name: str, section_input: str) -> str:
    """Dynamically fetch Explanatory Notes from the i.AI Lex API."""
    if not section_input.strip():
        return "Please specify a section number to view its Explanatory Note."
        
    try:
        # Extract the digits
        section_number = "".join([c for c in section_input if c.isdigit()])
        if not section_number:
            return "Please enter a valid section number."
            
        url = 'https://lex.lab.i.ai.gov.uk/explanatory_note/section/search'
        payload = {
            'query': f'"{act_name}" Section {section_number}',
            'limit': 5
        }
        
        async with httpx.AsyncClient() as client:
            r = await client.post(url, json=payload, timeout=10.0)
            if r.status_code == 200:
                data = r.json()
                if isinstance(data, list):
                    parent_id = NURSING_ACTS.get(act_name, "")
                    for note in data:
                        if parent_id and parent_id in note.get('legislation_id', ''):
                            text = note.get('text', '')
                            if text:
                                return f"### Official Explanatory Note\n\n{text}\n\n*Source: i.AI Lex API*"
                                
        return f"No official Explanatory Note found for {act_name} Section {section_number}.\n\n*(Note: Acts passed prior to 1999 generally do not have Explanatory Notes).*."
    except httpx.TimeoutException:
        return "⏳ API Timeout while fetching Explanatory Note."
    except Exception as e:
        return f"Error fetching note: {str(e)}"

async def scenario_search(scenario_text: str) -> str:
    """Use local i-dot-ai vector search to map a clinical scenario to legal sections."""
    if not scenario_text.strip():
        return "Please describe a clinical scenario."
        
    try:
        results = search_scenarios_locally(scenario_text, top_k=5)
        
        if not results:
            return "No matching legislation found for this scenario in the local cache."
            
        result = f"## ⚖️ Probable Legislation Matches for:\n*{scenario_text}*\n\n"
        
        for i, n in enumerate(results, 1):
            leg_id = n.get("legislation_id", "")
            
            # 1. Use the act_name from known mapping
            act_name = ""
            for known_id, known_name in REVERSE_ACTS.items():
                if known_id in leg_id:
                    act_name = known_name
                    break
            
            # 2. Final fallback: extract from the legislation_id URL
            if not act_name:
                act_name = leg_id.split("/id/")[-1] if "/id/" in leg_id else leg_id or "Legislation"
                    
            sec_num = n.get("number", "??")
            title = n.get("title", "Untitled Section")
            text = n.get("text", "")
            uri = n.get("uri", f"https://www.legislation.gov.uk/id/{leg_id}/section/{sec_num}")
            score = n.get("score", 0.0)
            
            result += f"### {i}. {act_name} — Section {sec_num}: {title} (Match Score: {score:.2f})\n"
            result += f"{text[:800]}...\n\n"
            result += f"🔗 [Read full text on legislation.gov.uk]({uri})\n\n---\n\n"
            
        return result
    except Exception as e:
        return f"Error during local scenario search: {str(e)}"

def browse_legislation(search_term: str, act_type: str) -> str:
    """Browse the legislation index from the parquet file."""
    if LEG_DF.empty:
        return "⚠️ Legislation index not loaded."

    filtered = LEG_DF.copy()

    if act_type != "All":
        type_map = {"Primary Acts": "ukpga", "Statutory Instruments": "uksi", "Scottish SIs": "ssi", "NI SRs": "nisr", "Welsh SIs": "wsi"}
        if act_type in type_map:
            filtered = filtered[filtered["type"] == type_map[act_type]]

    if search_term.strip():
        filtered = filtered[filtered["title"].str.contains(search_term, case=False, na=False)]

    filtered = filtered.sort_values("year", ascending=False).head(50)

    if filtered.empty:
        return f"No legislation found matching '{search_term}'."

    result = f"## 📖 Legislation Index ({len(filtered)} results)\n\n| Year | Title | Type |\n|---|---|---|\n"

    for _, row in filtered.iterrows():
        year = row.get("year", "—")
        title = row.get("title", "Untitled")
        uri = row.get("uri", "")
        leg_type = row.get("type", "")
        title_link = f"[{title}]({uri})" if uri else title
        result += f"| {year} | {title_link} | {leg_type} |\n"

    result += f"\n\n*Showing top 50 of {len(LEG_DF)} health & social care entries — {len(LEG_DF[LEG_DF['type']=='ukpga'])} Primary Acts*"
    result += "\n\n🏛️ *Data from i.AI Lex bulk downloads — Crown Copyright, OGL v3.0*"
    return result

# --- Gradio UI ---
THEME = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="violet",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
)

CSS = """
.gradio-container { max-width: 960px !important; }
.header-banner {
    background: linear-gradient(135deg, #312e81 0%, #4338ca 50%, #6366f1 100%);
    border-radius: 16px;
    padding: 28px 32px;
    margin-bottom: 16px;
    color: white;
}
.header-banner h1 { color: white; font-size: 2em; margin: 0 0 8px 0; }
.header-banner p { color: #c7d2fe; margin: 0; font-size: 1.05em; }
.disclaimer-box {
    background: #fef3c7;
    border-left: 4px solid #f59e0b;
    border-radius: 8px;
    padding: 12px 16px;
    margin-bottom: 12px;
    font-size: 0.9em;
    color: #92400e;
}
footer { display: none !important; }
"""

with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses") as app:
    gr.HTML("""
    <div class="header-banner">
        <h1>🏛️ NurseLex</h1>
        <p>Legal literacy for all nurses and nursing students — powered by UK Government legislation data</p>
    </div>
    """)

    gr.HTML("""
    <div class="disclaimer-box">
        ⚠️ <strong>Educational tool only.</strong> NurseLex provides legislation text and AI-generated explanations for learning purposes.
        It does not constitute legal advice. Always consult your trust's legal/governance team for specific cases.
    </div>
    """)

    with gr.Tabs():
        # --- Tab 1: Chat ---
        with gr.TabItem("💬 Ask a Legal Question", id="chat"):
            gr.Markdown("Ask about UK healthcare legislation — answers are grounded in **real statutory text**. (Cache: 1,128 Sections + 219K Acts)")

            chatbot = gr.Chatbot(
                label="NurseLex",
                height=480,
                type="messages",
                show_copy_button=True,
                avatar_images=(None, "https://em-content.zobj.net/source/twitter/376/classical-building_1f3db-fe0f.png"),
            )
            msg = gr.Textbox(
                label="Your question",
                placeholder="e.g., What is Section 5(4) of the Mental Health Act and when can a nurse use it?",
                lines=2,
            )
            with gr.Row():
                submit_btn = gr.Button("🔍 Search Legislation", variant="primary", scale=2)
                clear_btn = gr.ClearButton([msg, chatbot], value="🗑️ Clear", scale=1)

            gr.Markdown("### 💡 Quick Questions")
            with gr.Row(equal_height=True):
                for i in range(0, 4):
                    gr.Button(
                        QUICK_QUESTIONS[i][:55] + "..." if len(QUICK_QUESTIONS[i]) > 55 else QUICK_QUESTIONS[i],
                        size="sm",
                        variant="secondary",
                    ).click(fn=lambda q=QUICK_QUESTIONS[i]: q, outputs=msg)
            with gr.Row(equal_height=True):
                for i in range(4, 8):
                    gr.Button(
                        QUICK_QUESTIONS[i][:55] + "..." if len(QUICK_QUESTIONS[i]) > 55 else QUICK_QUESTIONS[i],
                        size="sm",
                        variant="secondary",
                    ).click(fn=lambda q=QUICK_QUESTIONS[i]: q, outputs=msg)

            async def respond(message, history):
                history = history or []
                history.append({"role": "user", "content": message})
                answer = await query_and_respond(message, history)
                history.append({"role": "assistant", "content": answer})
                return "", history

            submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
            msg.submit(respond, [msg, chatbot], [msg, chatbot])

        # --- Tab 2: Section Lookup ---
        with gr.TabItem("📖 Section Lookup", id="lookup"):
            gr.Markdown("Look up a **specific section** of key nursing Acts. Includes **Official Explanatory Notes** where available.")

            with gr.Row():
                act_dropdown = gr.Dropdown(
                    choices=list(NURSING_ACTS.keys()),
                    label="Select Act",
                    value="Mental Health Act 1983",
                )
                section_input_box = gr.Textbox(
                    label="Section number",
                    placeholder="e.g., 5 or 117 or 136",
                )

            lookup_btn = gr.Button("🔍 Look Up Law & Notes", variant="primary")
            
            with gr.Row():
                lookup_output = gr.Markdown(label="Statutory Text")
                note_output = gr.Markdown(label="Official Explanatory Note")

            lookup_btn.click(section_lookup, [act_dropdown, section_input_box], lookup_output)
            lookup_btn.click(fetch_explanatory_note, [act_dropdown, section_input_box], note_output)

        # --- Tab 3: Scenario Matcher ---
        with gr.TabItem("🧠 Scenario Matcher", id="scenario"):
            gr.Markdown("Describe a clinical scenario in plain English, and the **Lex Vector Search Engine** will map it to the most relevant UK laws.")
            
            with gr.Row():
                scenario_input = gr.Textbox(
                    label="Clinical Scenario",
                    placeholder="e.g. 'Patient wants to leave the ward but lacks capacity' or 'Doctor orders restraint without DoLS'",
                    lines=3
                )
            
            scenario_btn = gr.Button("🤖 Find Relevant Law", variant="primary")
            scenario_output = gr.Markdown(label="Semantic Search Results")
            
            scenario_btn.click(scenario_search, [scenario_input], scenario_output)

        # --- Tab 4: Browse Legislation ---
        with gr.TabItem("📚 Browse Legislation", id="browse"):
            gr.Markdown(f"Browse **219,678** health & social care Acts and Statutory Instruments from the i.AI Lex dataset.")

            with gr.Row():
                browse_search = gr.Textbox(
                    label="Search legislation titles",
                    placeholder="e.g., mental health, safeguarding, disability",
                )
                browse_type = gr.Dropdown(
                    choices=["All", "Primary Acts", "Statutory Instruments", "Scottish SIs", "NI SRs", "Welsh SIs"],
                    label="Type",
                    value="All",
                )

            browse_btn = gr.Button("🔍 Search", variant="primary")
            browse_output = gr.Markdown(label="Results")

            browse_btn.click(browse_legislation, [browse_search, browse_type], browse_output)

        # --- Tab 4: About ---
        with gr.TabItem("ℹ️ About", id="about"):
            gr.Markdown(f"""
## About NurseLex

**NurseLex** is a universal legal literacy tool for **all nurses and nursing students**.

### How It Works

1. **You ask a question** about UK healthcare law
2. **Cached legislation** provides the actual statutory text instantly
3. **Gemini Flash** explains it in plain English with practical nursing implications
4. **Every answer cites** the specific Act, section, and year

### Data

- **219,678 legislation entries** from the [i.AI Lex](https://lex.lab.i.ai.gov.uk/) bulk dataset
- **1,128 key sections** pre-cached with full text (MHA 1983, MCA 2005, Care Act 2014)
- **Crown Copyright** — Open Government Licence v3.0

### Key Acts Covered

| Act | Key Sections | Nursing Relevance |
|---|---|---|
| Mental Health Act 1983 | S.2, S.3, S.4, S.5(2), S.5(4), S.17, S.117, S.135, S.136 | Detention, holding powers, leave, aftercare |
| Mental Capacity Act 2005 | S.1 (Principles), S.2-3 (Capacity), S.4 (Best Interests), S.5 | Capacity assessments, best interests, DoLS |
| Care Act 2014 | S.42 (Safeguarding), S.67 (Advocacy) | Safeguarding adults, independent advocates |

### Built By

**NurseCitizenDeveloper** — NHS Registered Nurse building AI tools for nursing education.

🤗 [Hugging Face](https://huggingface.co/NurseCitizenDeveloper) · 🐙 [GitHub](https://github.com/Clinical-Quality-Intelligence)
            """)

app.queue()

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)