Gradio_UI.py

import gradio as gr
from fpdf import FPDF
from langdetect import detect
import time
import importlib

# === Map document choices to modules ===
DOCUMENT_TO_MODULE = {
    "AI Act": {
        "Compliance Register": "tools.ai_act_generator",
        "High-Risk System Record": "tools.high_risk_ai_register",
        "Technical Documentation Template": "tools.ai_technical_doc",
        "Post-Market Monitoring Plan": "tools.post_market_monitoring",
        "Risk-Management Plan": "tools.risk_management_plan",
        "Data Governance Record": "tools.data_governance_record",
        "Human Oversight Strategy (Art. 14)": "tools.human_oversight_strategy",
        "Corrective Action Log": "tools.corrective_action_log"
    },
    "GDPR": {
        "Data Processing Record": "tools.gdpr_data_record",
        "DPIA (Data Protection Impact Assessment)": "tools.gdpr_dpia",
        "Data Breach Notification Log": "tools.gdpr_breach_log",
        "Subject Access Request Record": "tools.gdpr_sar_log",
        "Consent Management Log": "tools.gdpr_consent_log",
        "Data Processor Agreements": "tools.gdpr_dpa_log"
    },
    "NIS2": {
        "Access Control Log": "tools.nis2_access_control_log",
        "Asset Inventory": "tools.nis2_asset_inventory",
        "Backup & Recovery Log": "tools.nis2_backup_recovery_log",
        "Business Continuity Plan": "tools.nis2_business_continuity",
        "Compliance Audit Log": "tools.nis2_compliance_audit_log",
        "Data Encryption Inventory": "tools.nis2_data_encryption_inventory",
        "Disaster Recovery Plan": "tools.nis2_disaster_recovery_plan",
        "Incident Communication Plan": "tools.nis2_incident_communication",
        "Incident Response Log": "tools.nis2_incident_log",
        "Logging & Monitoring Policy": "tools.nis2_logging_monitoring_policy",
        "Network Architecture Diagram Log": "tools.nis2_network_architecture_log",
        "Penetration Test Summary": "tools.nis2_penetration_test_summary",
        "Physical Security Checklist": "tools.nis2_physical_security_checklist",
        "Risk Assessment Matrix": "tools.nis2_risk_assessment_matrix",
        "Risk Register": "tools.nis2_risk_register",
        "Security Audit Record": "tools.nis2_security_audit",
        "Security Testing Strategy": "tools.nis2_security_testing_strategy",
        "Security Awareness Training Record": "tools.nis2_security_training_record",
        "Supply Chain Risk Record": "tools.nis2_supply_chain_risk",
        "System Testing Record": "tools.nis2_system_testing",
        "Third-Party Access Review": "tools.nis2_third_party_access_review",
        "Vulnerability Management Log": "tools.nis2_vulnerability_log"
    },
    "DORA": {
        "Third-Party ICT Risk Register": "tools.nis2_third_party_access_review",
        "ICT Asset Inventory": "tools.nis2_asset_inventory",
        "Incident Reporting Log": "tools.nis2_incident_log",
        "Risk Register": "tools.nis2_risk_register",
        "Backup & Recovery": "tools.nis2_backup_recovery_log",
        "Penetration Testing Summary": "tools.nis2_penetration_test_summary",
        "Disaster Recovery Plan": "tools.nis2_disaster_recovery_plan",
        "Security Testing Strategy": "tools.nis2_security_testing_strategy",
        "Backup Strategy Summary": "tools.dora_backup_strategy_summary",
        "Capacity Management Log": "tools.dora_capacity_management_log",
        "Cloud Resilience Assessment": "tools.dora_cloud_resilience_assessment",
        "Cloud Service Usage Log": "tools.dora_cloud_service_usage_log",
        "Communication Tree (Crisis)": "tools.dora_communication_tree",
        "Critical Function Inventory": "tools.dora_critical_function_inventory",
        "Data Integrity Strategy": "tools.dora_data_integrity_strategy",
        "Data Loss Prevention": "tools.dora_data_loss_prevention",
        "Encryption Strategy": "tools.dora_encryption_strategy",
        "ICT Asset Log": "tools.dora_ict_asset_log",
        "ICT Asset Register": "tools.dora_ict_asset_register",
        "Incident Notification Log": "tools.dora_incident_notification_log",
        "Network Resilience Assessment": "tools.dora_network_resilience_assessment",
        "Remote Access Policy": "tools.dora_remote_access_policy",
        "Risk Management Policy": "tools.dora_risk_management_policy",
        "Simulation Test Log": "tools.dora_simulation_test_log",
        "Third-Party Dependency Map": "tools.dora_third_party_dependency_map",
        "Third-Party Dependency Register": "tools.dora_third_party_dependency_register",
        "Third-Party Log": "tools.dora_third_party_log",
        "Third-Party Mapping": "tools.dora_third_party_mapping"
    },
        "DMA": {
        "Gatekeeper Service Log": "tools.dma_gatekeeper_service_log",
        "User Consent Overview": "tools.dma_user_consent_overview",
        "Interoperability Compliance Report": "tools.dma_interoperability_report",
        "Data Sharing Record": "tools.dma_data_sharing_record",
        "Complaint Handling Log": "tools.dma_complaint_log",
        "Self-Assessment Declaration": "tools.dma_self_assessment",
        "Third-Party Access Request Register": "tools.dma_third_party_access_register",
        "Transparency Obligation Record": "tools.dma_transparency_log"
    },
        "DSA": {
        "Content Moderation Log": "tools.dsa_content_moderation_log",
        "Transparency Report": "tools.dsa_transparency_report",
        "Complaint Handling Register": "tools.dsa_complaint_handling_log",
        "Risk Mitigation Measures": "tools.dsa_risk_mitigation_log",
        "Ad Targeting Policy Log": "tools.dsa_ad_targeting_log",
        "Trusted Flaggers Log": "tools.dsa_trusted_flaggers_log"
    }   
}


# === PDF Export Function ===
def export_pdf(title, answers, lang="en"):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_auto_page_break(auto=True, margin=15)

    pdf.set_font("Arial", 'B', 16)
    pdf.set_text_color(0, 51, 102)
    pdf.cell(0, 10, title, ln=True, align='C')
    pdf.ln(8)

    company = answers.get("What is the name of your organization?", "N/A")
    name = answers.get("What is your name?", "N/A")
    role = answers.get("What is your role?", "N/A")
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")

    pdf.set_font("Arial", 'I', 11)
    pdf.set_text_color(80, 80, 80)
    pdf.cell(0, 8, f"Organization: {company}", ln=True, align='C')
    pdf.cell(0, 8, f"Completed by: {name} ({role})", ln=True, align='C')
    pdf.cell(0, 8, f"Timestamp: {timestamp}", ln=True, align='C')
    pdf.ln(10)

    # Add a draft disclaimer
    pdf.set_font("Arial", 'B', 12)
    pdf.set_text_color(200, 0, 0)  # Red color
    pdf.multi_cell(0, 10, "DRAFT: This document is a draft generated by an AI assistant.\nIt is not legal advice and should be reviewed by a qualified professional.")
    pdf.ln(10)


    pdf.set_font("Arial", '', 12)
    pdf.set_text_color(0, 0, 0)
    for label, value in answers.items():
        pdf.set_font("Arial", 'B', 12)
        pdf.multi_cell(0, 10, f"{label}:", align='L')
        pdf.set_font("Arial", '', 12)
        pdf.multi_cell(0, 10, value)
        pdf.ln(5)

    filename = f"{title.replace(' ', '_')}_{int(time.time())}.pdf"
    pdf.output(filename)
    return filename

# === State Memory ===
conversation_state = {
    "step": 0,
    "questions": [],
    "answers": {},
    "title": "",
    "lang": "en",
    "tool": None
}

# === Load Tool Dynamically ===
def start_conversation(reg, doc):
    conversation_state["step"] = 0
    conversation_state["answers"] = {}
    conversation_state["title"] = f"{reg} - {doc}"

    module_path = DOCUMENT_TO_MODULE[reg][doc]
    module = importlib.import_module(module_path)
    conversation_state["questions"] = module.get_questions()
    conversation_state["tool"] = module

    first_question = conversation_state["questions"][0][1]
    return [{"role": "assistant", "content": f"📋 {first_question}"}]

# === Continue Chat ===
def continue_conversation(user_input):
    messages = [{"role": "user", "content": user_input}]
    step = conversation_state["step"]

    if step < len(conversation_state["questions"]):
        key, question = conversation_state["questions"][step]
        conversation_state["answers"][question] = user_input
        conversation_state["step"] += 1

        try:
            if user_input and len(user_input.strip()) > 3:
                conversation_state["lang"] = detect(user_input)
        except:
            conversation_state["lang"] = "en"

    if conversation_state["step"] < len(conversation_state["questions"]):
        next_question = conversation_state["questions"][conversation_state["step"]][1]
        messages.append({"role": "assistant", "content": f"🧠 {next_question}"})
        return messages, None
    else:
        pdf_path = export_pdf(conversation_state["title"], conversation_state["answers"], lang=conversation_state["lang"])
        messages.append({"role": "assistant", "content": "✅ Done! Your document is ready to download."})
        return messages, pdf_path

# === Gradio Interface ===
with gr.Blocks(
    title="Regulatory Compliance Generator",
    css="""
    footer, 
    a[href*="gradio.app"],
    a[href*="huggingface.co"] {
        display: none !important;
    }

    button[aria-label="Share"],
    button[title="Share"],
    button[data-testid="share-button"] {
        pointer-events: none !important;
        opacity: 0.4 !important;
        cursor: default !important;
    }
    """
) as demo:
    gr.Markdown("""
    ### 🔒 GDPR Notice  
    This assistant does not store any personal data.  
    All responses are processed temporarily to generate your document.  
    You are responsible for the accuracy and legality of the information submitted.
    """)

    with gr.Row():
        reg_dropdown = gr.Dropdown(
            choices=list(DOCUMENT_TO_MODULE.keys()),
            label="Select Regulation",
            value=None,
            interactive=True,
            info="-- Please choose a regulation --"
        )
        doc_dropdown = gr.Dropdown(
            label="Select Document Type",
            choices=[],
            value=None,
            interactive=True,
            info="-- Select a document after choosing a regulation --"
        )

    def update_doc_dropdown(selected_reg):
        if selected_reg in DOCUMENT_TO_MODULE:
            return gr.update(choices=list(DOCUMENT_TO_MODULE[selected_reg].keys()), value=None)
        return gr.update(choices=[], value=None)

    reg_dropdown.change(update_doc_dropdown, inputs=reg_dropdown, outputs=doc_dropdown)

    start_btn = gr.Button("🚀 Start")
    chatbox = gr.Chatbot(label="🧑‍⚖️ Legalbot", value=[], type="messages", elem_classes=["chatbot"])
    user_input = gr.Textbox(label="Your answer", placeholder="Type your answer and press Enter")
    pdf_output = gr.File(label="Download PDF", visible=True)

    start_btn.click(start_conversation, inputs=[reg_dropdown, doc_dropdown], outputs=chatbox)
    user_input.submit(continue_conversation, inputs=user_input, outputs=[chatbox, pdf_output])
    user_input.submit(lambda: "", None, user_input)

demo.launch()