import streamlit as st
from utils import translate_text, fetch_pubmed_articles, fetch_arxiv_articles
import os

# Page configuration
st.set_page_config(
    page_title="Gov-LLM Transaction Flow",
    page_icon="🏛️",
    layout="wide"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: bold;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 2rem;
    }
    .sub-header {
        font-size: 1.5rem;
        font-weight: bold;
        color: #ff7f0e;
        margin-top: 1.5rem;
    }
    .built-with {
        font-size: 0.8rem;
        color: #666;
        text-align: right;
    }
    .phase-box {
        border: 2px solid #1f77b4;
        border-radius: 10px;
        padding: 15px;
        margin: 10px 0;
        background-color: #f0f8ff;
    }
    .reference-box {
        background-color: #e6f7ff;
        padding: 10px;
        border-left: 4px solid #1f77b4;
        margin: 10px 0;
    }
</style>
""", unsafe_allow_html=True)

# Header
st.markdown('<div class="main-header">Gov-LLM Transaction Flow</div>', unsafe_allow_html=True)
st.markdown('<div class="built-with">Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder">anycoder</a></div>', unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.title("Navigation")
    page = st.radio("Select Phase", ["Overview", "Phase 0", "Phase 1", "Phase 2", "Phase 3", "Phase 4", "Phase 5", "References"])

    st.divider()
    st.header("About")
    st.markdown("""
    This application demonstrates the Gov-LLM Transaction Flow, a comprehensive framework for ADA-compliant cognitive assistance in government services.
    """)

# Main content based on selected page
if page == "Overview":
    st.header("Gov-LLM Transaction Flow Overview")
    st.markdown("""
    This application presents a complete, executable framework for implementing cognitive accessibility in government services.
    The Gov-LLM Transaction Flow is designed to comply with ADA Title II and Due Process requirements while maintaining
    strict data privacy and security standards.
    """)

    st.subheader("Key Features")
    st.markdown("""
    - ✅ ADA Title II Compliance
    - ✅ Due Process Protection
    - ✅ Post-Quantum Cryptography
    - ✅ Cognitive Load Reduction
    - ✅ Audit Trail & Accountability
    """)

    st.subheader("Technical Implementation")
    st.markdown("""
    The system uses:
    - **Kyber-768** for post-quantum key exchange
    - **Dilithium** for digital signatures
    - **RAG (Retrieval-Augmented Generation)** for legal document processing
    - **NIST-compliant** security protocols
    """)

elif page == "Phase 0":
    st.markdown('<div class="sub-header">Phase 0: Event-Driven Trigger Mechanism</div>', unsafe_allow_html=True)

    st.markdown("""
    <div class="phase-box">
    <h3>Operational Trigger</h3>
    <p>A citizen interacts with a .gov interface (SSA, Federal Court, Benefits Portal).</p>

    <h3>System Action</h3>
    <p>Real-time detection of navigation latency or accessibility metadata triggers the "Cognitive Auxiliary Mode" (Gov-LLM Interface).</p>

    <h3>Legal Basis</h3>
    <p>Affirmative duty to provide "auxiliary aids and services" (28 C.F.R. § 35.160).</p>
    </div>
    """, unsafe_allow_html=True)

    st.markdown("""
    <div class="reference-box">
    <h4>Technical & Academic Substantiation</h4>
    <p><strong>Just-in-Time Adaptive Interventions (JITAI):</strong> The system utilizes JITAI principles to deliver support at the moment of cognitive need.</p>
    <p>Reference: Nahum-Shani, I., et al. (2018). "Just-in-Time Adaptive Interventions (JITAIs) in Mobile Health." Annals of Behavioral Medicine.</p>
    """, unsafe_allow_html=True)

elif page == "Phase 1":
    st.markdown('<div class="sub-header">Phase 1: Ephemeral Session Initialization</div>', unsafe_allow_html=True)

    st.markdown("""
    <div class="phase-box">
    <h3>Transaction Input</h3>
    <p>User explicitly consents to "Cognitive Assistance" (Rights-preserving opt-in).</p>

    <h3>Privacy Protocol</h3>
    <ul>
        <li>Generation of a Cryptographically Random Session ID (UUIDv4)</li>
        <li>No PII Retention: Identity is verified via token, not database storage</li>
        <li>Isolation: The session is "air-gapped" from third-party commercial model training pipelines</li>
    </ul>

    <h3>Time Point</h3>
    <p>T1 (Cognitive Assistance Begins)</p>
    </div>
    """, unsafe_allow_html=True)

elif page == "Phase 2":
    st.markdown('<div class="sub-header">Phase 2: Post-Quantum Cryptographic Enveloping</div>', unsafe_allow_html=True)

    st.markdown("""
    <div class="phase-box">
    <h3>Transaction Behavior</h3>
    <p>User inputs (medical evidence, pleadings, queries) are instantaneously encapsulated.</p>

    <h3>Security Architecture</h3>
    <ul>
        <li><strong>Key Exchange:</strong> Kyber-768/1024 (NIST FIPS 203)</li>
        <li><strong>Integrity Signing:</strong> Dilithium (NIST FIPS 204)</li>
    </ul>

    <h3>Legal Implications</h3>
    <ul>
        <li>✅ Complies with privacy protection standards</li>
        <li>✅ Does not constitute medical data forced disclosure</li>
    </ul>

    <h3>Time Point</h3>
    <p>T2 (Secure Transaction Established)</p>
    </div>
    """, unsafe_allow_html=True)

elif page == "Phase 3":
    st.markdown('<div class="sub-header">Phase 3: Retrieval-Augmented Inference</div>', unsafe_allow_html=True)

    st.markdown("""
    <div class="phase-box">
    <h3>Core Processing</h3>
    <p>The Gov-LLM processes the query using RAG (Retrieval-Augmented Generation).</p>

    <h3>Restricted Context Window</h3>
    <ul>
        <li><strong>Allowed Corpus:</strong> USC, CFR, POMS, HALLEX, Public Case Law</li>
        <li><strong>Prohibited Actions:</strong> No internet browsing, no hallucination, no training on user data</li>
    </ul>

    <h3>Cognitive Transformation</h3>
    <p>Text is processed for "Cognitive Load Reduction" (summarization, step-extraction).</p>

    <h3>Time Point</h3>
    <p>T3 (Cognitive Translation Complete)</p>
    </div>
    """, unsafe_allow_html=True)

elif page == "Phase 4":
    st.markdown('<div class="sub-header">Phase 4: Accessibility Rendering & Verification</div>', unsafe_allow_html=True)

    st.markdown("""
    <div class="phase-box">
    <h3>Output Format</h3>
    <ul>
        <li>Linearized Logic: Step-by-step instructions (List format)</li>
        <li>State Markers: "You are at Step X of Y"</li>
        <li>Rewind Capability: User can request "Repeat/Rephrase" without penalty</li>
    </ul>

    <h3>ADA Compliance Check</h3>
    <p>The system verifies the output meets WCAG 2.2 AA/AAA standards for cognitive accessibility.</p>

    <h3>Time Point</h3>
    <p>T4 (Understandability Achieved)</p>
    </div>
    """, unsafe_allow_html=True)

elif page == "Phase 5":
    st.markdown('<div class="sub-header">Phase 5: Secure Termination & Cryptographic Shredding</div>', unsafe_allow_html=True)

    st.markdown("""
    <div class="phase-box">
    <h3>Termination Condition</h3>
    <ul>
        <li>User closes browser</li>
        <li>Completes task</li>
        <li>Timeout</li>
    </ul>

    <h3>Data Destruction</h3>
    <ul>
        <li>Forward Secrecy: Ephemeral session keys are destroyed from memory</li>
        <li>Audit Trail: Only a hashed timestamp (SHA-3) remains</li>
    </ul>

    <h3>Guarantee</h3>
    <ul>
        <li>✅ No data persists for commercial reuse</li>
        <li>✅ No government surveillance</li>
    </ul>

    <h3>Time Point</h3>
    <p>T5 (Transaction Complete)</p>
    </div>
    """, unsafe_allow_html=True)

elif page == "References":
    st.header("Academic and Technical References")

    st.subheader("Phase 0 - JITAI")
    st.markdown("""
    <div class="reference-box">
    <p>Nahum-Shani, I., et al. (2018). "Just-in-Time Adaptive Interventions (JITAIs) in Mobile Health: Key Components and Design Principles for Engineering and Science." Annals of Behavioral Medicine.</p>
    </div>
    """, unsafe_allow_html=True)

    st.subheader("Phase 2 - Post-Quantum Cryptography")
    st.markdown("""
    <div class="reference-box">
    <p>Albrecht, M. R., et al. "Lattice-based cryptography and the NIST standardization process." ArXiv.</p>
    <p>NIST FIPS 203: Module-Lattice-Based Key-Encapsulation Mechanism Standard</p>
    <p>NIST FIPS 204: Module-Lattice-Based Digital Signature Standard</p>
    </div>
    """, unsafe_allow_html=True)

    st.subheader("Phase 3 - RAG")
    st.markdown("""
    <div class="reference-box">
    <p>Lewis, P., et al. (2020). "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks." NeurIPS.</p>
    </div>
    """, unsafe_allow_html=True)

    st.subheader("Phase 4 - Cognitive Load Theory")
    st.markdown("""
    <div class="reference-box">
    <p>Sweller, J. (2011). "Cognitive Load Theory." Psychology of Learning and Motivation.</p>
    <p>WCAG 2.2 Guidelines: Success Criterion 3.3.2 (Labels or Instructions) and 3.1.5 (Reading Level)</p>
    </div>
    """, unsafe_allow_html=True)

    st.subheader("Phase 5 - Secure Deletion")
    st.markdown("""
    <div class="reference-box">
    <p>NIST SP 800-88 Rev. 1: Guidelines for Media Sanitization</p>
    </div>
    """, unsafe_allow_html=True)

# Footer
st.divider()
st.markdown("""
<div style="text-align: center; color: #666; padding: 20px;">
    <p>Gov-LLM Transaction Flow Application | Built with Streamlit</p>
    <p>© 2023 | All rights reserved</p>
</div>
""", unsafe_allow_html=True)