app.py

import gradio as gr
from utils.meldrx import MeldRxAPI
import json
import os
import tempfile
from datetime import datetime
import traceback
import logging
from huggingface_hub import InferenceClient  # Import InferenceClient
from urllib.parse import urlparse, parse_qs  # Import URL parsing utilities
from utils.callbackmanager import CallbackManager
from utils.meldrx import MeldRxAPI
from utils.prompts import system_instructions
from old.extractcode import extract_code_from_url
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import PDF utilities
from utils.pdfutils import PDFGenerator, generate_discharge_summary

# Import necessary libraries for new file types and AI analysis functions
import pydicom  # For DICOM
import hl7  # For HL7
from xml.etree import ElementTree  # For XML and CCDA
from pypdf import PdfReader  # For PDF
import csv  # For CSV
import io  # For IO operations
from PIL import Image  # For image handling
from utils.callbackmanager import extract_auth_code_from_url, extract_code_from_url
from utils.generators import  generate_pdf_from_meldrx, generate_ai_discharge_content, generate_pdf_from_meldrx_with_ai_content, extract_section,  generate_pdf_from_form, generate_discharge_summary, generate_ai_discharge_content, analyze_dicom_file_with_ai, analyze_hl7_file_with_ai, analyze_cda_xml_file_with_ai, analyze_pdf_file_with_ai, analyze_csv_file_with_ai, generate_pdf_from_form , generate_ai_discharge_content , extract_section , generate_pdf_from_meldrx_with_ai_content 


# Initialize Inference Client - Ensure YOUR_HF_TOKEN is set in environment variables or replace with your actual token
HF_TOKEN = os.getenv("HF_TOKEN")  # Or replace with your actual token string
if not HF_TOKEN:
    raise ValueError(
        "HF_TOKEN environment variable not set. Please set your Hugging Face API token."
    )
client = InferenceClient(api_key=HF_TOKEN)
model_name = "meta-llama/Llama-3.3-70B-Instruct"  # Specify the model to use


def display_form(first_name, last_name, middle_initial, dob, age, sex, address, city, state, zip_code, doctor_first_name, doctor_last_name, doctor_middle_initial, hospital_name, doctor_address, doctor_city, doctor_state, doctor_zip, admission_date, referral_source, admission_method, discharge_date, discharge_reason, date_of_death, diagnosis, procedures, medications, preparer_name, preparer_job_title,):
    form = f"""
    <div style='color:#00FFFF; font-family: monospace;'>
    **Patient Discharge Form** <br>
    - Name: {first_name} {middle_initial} {last_name} <br>
    - Date of Birth: {dob}, Age: {age}, Sex: {sex} <br>
    - Address: {address}, {city}, {state}, {zip_code} <br>
    - Doctor: {doctor_first_name} {doctor_middle_initial} {doctor_last_name} <br>
    - Hospital/Clinic: {hospital_name} <br>
    - Doctor Address: {doctor_address}, {doctor_city}, {doctor_state}, {doctor_zip} <br>
    - Admission Date: {admission_date}, Source: {referral_source}, Method: {admission_method} <br>
    - Discharge Date: {discharge_date}, Reason: {discharge_reason} <br>
    - Date of Death: {date_of_death} <br>
    - Diagnosis: {diagnosis} <br>
    - Procedures: {procedures} <br>
    - Medications: {medications} <br>
    - Prepared By: {preparer_name}, {preparer_job_title}
    </div>
    """
    return form



CALLBACK_MANAGER = CallbackManager(
    redirect_uri="https://multitransformer-discharge-guard.hf.space/callback",
    client_secret=None,
)


# class CallbackManager:
#     def __init__(self, redirect_uri: str, client_secret: str = None):
#         client_id = os.getenv("APPID")
#         if not client_id:
#             raise ValueError("APPID environment variable not set.")
#         workspace_id = os.getenv("WORKSPACE_URL")
#         if not workspace_id:
#             raise ValueError("WORKSPACE_URL environment variable not set.")
#         self.api = MeldRxAPI(client_id, client_secret, workspace_id, redirect_uri)
#         self.auth_code = None
#         self.access_token = None

#     def handle_callback(self, callback_url: str) -> str:
#         """Handles the callback URL and extracts the code automatically."""
#         self.auth_code = extract_code_from_url(callback_url)
#         if not self.auth_code:
#             return "No authentication code found in URL."

#         if self.api.authenticate_with_code(self.auth_code):
#             self.access_token = self.api.access_token
#             return f"Authentication successful! Access Token: {self.access_token[:10]}... (truncated)"
#         return "Authentication failed. Please check the authorization code."
    
def generate_discharge_paper_one_click():
    """One-click function to fetch patient data and generate discharge paper with AI Content."""
    patient_data_str = CALLBACK_MANAGER.get_patient_data()
    if (
        patient_data_str.startswith("Not authenticated")
        or patient_data_str.startswith("Failed")
        or patient_data_str.startswith("Error")
    ):
        return None, patient_data_str  # Return error message if authentication or data fetch fails

    try:
        patient_data = json.loads(patient_data_str)

        # --- AI Content Generation for Discharge Summary ---
        # This is a placeholder - Replace with actual AI call using InferenceClient and patient_data to generate content
        ai_generated_content = generate_ai_discharge_content(
            patient_data
        )  # Placeholder AI function

        if not ai_generated_content:
            return None, "Error: AI content generation failed."

        # --- PDF Generation with AI Content ---
        pdf_path, status_message = generate_pdf_from_meldrx_with_ai_content(
            patient_data, ai_generated_content
        )  # Function to generate PDF with AI content

        if pdf_path:
            return pdf_path, status_message
        else:
            return None, status_message  # Return status message if PDF generation fails

    except json.JSONDecodeError:
        return None, "Error: Patient data is not in valid JSON format."
    except Exception as e:
        return None, f"Error during discharge paper generation: {str(e)}"



# Define the cyberpunk theme - using a dark base and neon accents
cyberpunk_theme = gr.themes.Monochrome(
    primary_hue="cyan",
    secondary_hue="pink",
    neutral_hue="slate",
    font=["Source Code Pro", "monospace"], # Retro monospace font
    font_mono=["Source Code Pro", "monospace"]
)

# Create the UI with the cyberpunk theme
with gr.Blocks(theme=cyberpunk_theme) as demo: # Apply the theme here
    gr.Markdown("<h1 style='color:#00FFFF; text-shadow: 0 0 5px #00FFFF;'>Discharge Guard <span style='color:#FF00FF; text-shadow: 0 0 5px #FF00FF;'>Cyber</span></h1>") # Cyberpunk Title

    with gr.Tab("Authenticate with MeldRx", elem_classes="cyberpunk-tab"): # Optional: Class for tab styling
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>SMART on FHIR Authentication</h2>") # Neon Tab Header
        auth_url_output = gr.Textbox(label="Authorization URL", value=CALLBACK_MANAGER.get_auth_url(), interactive=False)
        gr.Markdown("<p style='color:#A9A9A9;'>Copy the URL above, open it in a browser, log in, and paste the <span style='color:#00FFFF;'>entire redirected URL</span> from your browser's address bar below.</p>") # Subdued instructions with neon highlight
        redirected_url_input = gr.Textbox(label="Redirected URL") # New textbox for redirected URL
        extract_code_button = gr.Button("Extract Authorization Code", elem_classes="cyberpunk-button") # Cyberpunk button style
        extracted_code_output = gr.Textbox(label="Extracted Authorization Code", interactive=False) # Textbox to show extracted code

        auth_code_input = gr.Textbox(label="Authorization Code (from above, or paste manually if extraction fails)", interactive=True) # Updated label to be clearer
        auth_submit = gr.Button("Submit Code for Authentication", elem_classes="cyberpunk-button") # Cyberpunk button style
        auth_result = gr.HTML(label="Authentication Result") # Use HTML for styled result

        patient_data_button = gr.Button("Fetch Patient Data", elem_classes="cyberpunk-button") # Cyberpunk button style
        patient_data_output = gr.Textbox(label="Patient Data", lines=10)

        # Add button to generate PDF from MeldRx data (No AI)
        meldrx_pdf_button = gr.Button("Generate PDF from MeldRx Data (No AI)", elem_classes="cyberpunk-button") # Renamed button
        meldrx_pdf_status = gr.Textbox(label="PDF Generation Status (No AI)") # Renamed status
        meldrx_pdf_download = gr.File(label="Download Generated PDF (No AI)") # Renamed download

        def process_redirected_url(redirected_url):
            """Processes the redirected URL to extract and display the authorization code."""
            auth_code, error_message = extract_auth_code_from_url(redirected_url)
            if auth_code:
                return auth_code, "<span style='color:#00FF7F;'>Authorization code extracted!</span>" # Neon Green Success
            else:
                return "", f"<span style='color:#FF4500;'>Could not extract authorization code.</span> {error_message or ''}" # Neon Orange Error


        extract_code_button.click(
            fn=process_redirected_url,
            inputs=redirected_url_input,
            outputs=[extracted_code_output, auth_result],# Reusing auth_result for extraction status
        )

        auth_submit.click(
            fn=CALLBACK_MANAGER.set_auth_code,
            inputs=extracted_code_output,  # Using extracted code as input for authentication
            outputs=auth_result,
        )

    with gr.Tab("Patient Dashboard", elem_classes="cyberpunk-tab"): # Optional: Class for tab styling
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Patient Data</h2>") # Neon Tab Header
        dashboard_output = gr.HTML("<p style='color:#A9A9A9;'>Fetch patient data from the Authentication tab first.</p>") # Subdued placeholder text

        refresh_btn = gr.Button("Refresh Data", elem_classes="cyberpunk-button") # Cyberpunk button style

        # Simple function to update dashboard based on fetched data
        def update_dashboard():
            try:
                data = CALLBACK_MANAGER.get_patient_data()
                if (
                    data.startswith("<span style='color:#FF8C00;'>Not authenticated")
                    or data.startswith("<span style='color:#DC143C;'>Failed")
                    or data.startswith("<span style='color:#FF6347;'>Error")
                ):
                    return f"<p style='color:#FF8C00;'>{data}</p>" # Show auth errors in orange

                try:
                    # Parse the data
                    patients_data = json.loads(data)
                    patients = []

                    # Extract patients from bundle
                    for entry in patients_data.get("entry", []):
                        resource = entry.get("resource", {})
                        if resource.get("resourceType") == "Patient":
                            patients.append(resource)

                    # Generate HTML card
                    html = "<h3 style='color:#00FFFF; text-shadow: 0 0 2px #00FFFF;'>Patients</h3>" # Neon Sub-header
                    for patient in patients:
                        # Extract name
                        name = patient.get("name", [{}])[0]
                        given = " ".join(name.get("given", ["Unknown"]))
                        family = name.get("family", "Unknown")

                        # Extract other details
                        gender = patient.get("gender", "unknown").capitalize()
                        birth_date = patient.get("birthDate", "Unknown")

                        # Generate HTML card with cyberpunk styling
                        html += f"""
                        <div style="border: 1px solid #00FFFF; padding: 10px; margin: 10px 0; border-radius: 5px; background-color: #222; box-shadow: 0 0 5px #00FFFF;">
                            <h4 style='color:#00FFFF;'>{given} {family}</h4>
                            <p style='color:#A9A9A9;'><strong>Gender:</strong> <span style='color:#00FFFF;'>{gender}</span></p>
                            <p style='color:#A9A9A9;'><strong>Birth Date:</strong> <span style='color:#00FFFF;'>{birth_date}</span></p>
                            <p style='color:#A9A9A9;'><strong>ID:</strong> <span style='color:#00FFFF;'>{patient.get("id", "Unknown")}</span></p>
                        </div>
                        """

                    return html
                except Exception as e:
                    return f"<p style='color:#FF6347;'>Error parsing patient data: {str(e)}</p>" # Tomato Error
            except Exception as e:
                return f"<p style='color:#FF6347;'>Error fetching patient data: {str(e)}</p>" # Tomato Error


    with gr.Tab("Discharge Form", elem_classes="cyberpunk-tab"): # Optional: Class for tab styling
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Patient Details</h2>") # Neon Tab Header
        with gr.Row():
            first_name = gr.Textbox(label="First Name")
            last_name = gr.Textbox(label="Last Name")
            middle_initial = gr.Textbox(label="Middle Initial")
        with gr.Row():
            dob = gr.Textbox(label="Date of Birth")
            age = gr.Textbox(label="Age")
            sex = gr.Textbox(label="Sex")
        address = gr.Textbox(label="Address")
        with gr.Row():
            city = gr.Textbox(label="City")
            state = gr.Textbox(label="State")
            zip_code = gr.Textbox(label="Zip Code")
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Primary Healthcare Professional Details</h2>") # Neon Sub-header
        with gr.Row():
            doctor_first_name = gr.Textbox(label="Doctor's First Name")
            doctor_last_name = gr.Textbox(label="Doctor's Last Name")
            doctor_middle_initial = gr.Textbox(label="Doctor's Middle Initial")
        hospital_name = gr.Textbox(label="Hospital/Clinic Name")
        doctor_address = gr.Textbox(label="Address")
        with gr.Row():
            doctor_city = gr.Textbox(label="City")
            doctor_state = gr.Textbox(label="State")
            doctor_zip = gr.Textbox(label="Zip Code")
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Admission and Discharge Details</h2>") # Neon Sub-header
        with gr.Row():
            admission_date = gr.Textbox(label="Date of Admission")
            referral_source = gr.Textbox(label="Source of Referral")
        admission_method = gr.Textbox(label="Method of Admission")
        with gr.Row():
            discharge_date = gr.Textbox(label="Date of Discharge")
            discharge_reason = gr.Radio(
                ["Treated", "Transferred", "Discharge Against Advice", "Patient Died"],
                label="Discharge Reason",
            )
        date_of_death = gr.Textbox(label="Date of Death (if applicable)")
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Diagnosis & Procedures</h2>") # Neon Sub-header
        diagnosis = gr.Textbox(label="Diagnosis")
        procedures = gr.Textbox(label="Operation & Procedures")
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Medication Details</h2>") # Neon Sub-header
        medications = gr.Textbox(label="Medication on Discharge")
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Prepared By</h2>") # Neon Sub-header
        with gr.Row():
            preparer_name = gr.Textbox(label="Name")
            preparer_job_title = gr.Textbox(label="Job Title")

        # Add buttons for both display form and generate PDF
        with gr.Row():
            submit_display = gr.Button("Display Form", elem_classes="cyberpunk-button") # Cyberpunk button style
            submit_pdf = gr.Button("Generate PDF (No AI)", elem_classes="cyberpunk-button") # Renamed button to clarify no AI and styled

        # Output areas
        form_output = gr.HTML() # Use HTML to render styled form
        pdf_output = gr.File(label="Download PDF (No AI)")  # Renamed output to clarify no AI

        # Connect the display form button
        submit_display.click(
            display_form,
            inputs=[ first_name, last_name, middle_initial, dob, age, sex, address, city, state, zip_code, doctor_first_name, doctor_last_name, doctor_middle_initial, hospital_name, doctor_address, doctor_city, doctor_state, doctor_zip, admission_date, referral_source, admission_method, discharge_date, discharge_reason, date_of_death, diagnosis, procedures, medications, preparer_name, preparer_job_title,],
            outputs=form_output
        )

        # Connect the generate PDF button (No AI version)
        submit_pdf.click(
            generate_pdf_from_form,
            inputs=[ first_name, last_name, middle_initial, dob, age, sex, address, city, state, zip_code, doctor_first_name, doctor_last_name, doctor_middle_initial, hospital_name, doctor_address, doctor_city, doctor_state, doctor_zip, admission_date, referral_source, admission_method, discharge_date, discharge_reason, date_of_death, diagnosis, procedures, medications, preparer_name, preparer_job_title,],
            outputs=pdf_output
        )

    with gr.Tab("Medical File Analysis", elem_classes="cyberpunk-tab"): # Optional: Class for tab styling
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>Analyze Medical Files with Discharge Guard AI</h2>") # Neon Tab Header
        with gr.Column():
            dicom_file = gr.File(
                file_types=[".dcm"], label="Upload DICOM File (.dcm)"
            )
            dicom_ai_output = gr.Textbox(label="DICOM Analysis Report", lines=5)
            analyze_dicom_button = gr.Button("Analyze DICOM with AI", elem_classes="cyberpunk-button") # Cyberpunk button style

            hl7_file = gr.File(
                file_types=[".hl7"], label="Upload HL7 File (.hl7)"
            )
            hl7_ai_output = gr.Textbox(label="HL7 Analysis Report", lines=5)
            analyze_hl7_button = gr.Button("Analyze HL7 with AI", elem_classes="cyberpunk-button") # Cyberpunk button style

            xml_file = gr.File(
                file_types=[".xml"], label="Upload XML File (.xml)"
            )
            xml_ai_output = gr.Textbox(label="XML Analysis Report", lines=5)
            analyze_xml_button = gr.Button("Analyze XML with AI", elem_classes="cyberpunk-button") # Cyberpunk button style

            ccda_file = gr.File(
                file_types=[".xml", ".cda", ".ccd"], label="Upload CCDA File (.xml, .cda, .ccd)"
            )
            ccda_ai_output = gr.Textbox(label="CCDA Analysis Report", lines=5)
            analyze_ccda_button = gr.Button("Analyze CCDA with AI", elem_classes="cyberpunk-button") # Cyberpunk button style

            ccd_file = gr.File(
                file_types=[".ccd"],
                label="Upload CCD File (.ccd)",
            )  # Redundant, as CCDA also handles .ccd, but kept for clarity
            ccd_ai_output = gr.Textbox(
                label="CCD Analysis Report", lines=5
            )  # Redundant
            analyze_ccd_button = gr.Button("Analyze CCD with AI", elem_classes="cyberpunk-button") # Cyberpunk button style # Redundant
            pdf_file = gr.File(
                file_types=[".pdf"], label="Upload PDF File (.pdf)"
            )
            pdf_ai_output = gr.Textbox(label="PDF Analysis Report", lines=5)
            analyze_pdf_button = gr.Button("Analyze PDF with AI", elem_classes="cyberpunk-button") # Cyberpunk button style

            csv_file = gr.File(
                file_types=[".csv"], label="Upload CSV File (.csv)"
            )
            csv_ai_output = gr.Textbox(label="CSV Analysis Report", lines=5)
            analyze_csv_button = gr.Button("Analyze CSV with AI", elem_classes="cyberpunk-button") # Cyberpunk button style

        # Connect AI Analysis Buttons - using REAL AI functions now
        analyze_dicom_button.click(
            analyze_dicom_file_with_ai,  # Call REAL AI function
            inputs=dicom_file,
            outputs=dicom_ai_output
        )
        analyze_hl7_button.click(
            analyze_hl7_file_with_ai,  # Call REAL AI function
            inputs=hl7_file,
            outputs=hl7_ai_output
        )
        analyze_xml_button.click(
            analyze_cda_xml_file_with_ai,  # Call REAL AI function
            inputs=xml_file,
            outputs=xml_ai_output
        )
        analyze_ccda_button.click(
            analyze_cda_xml_file_with_ai,  # Call REAL AI function
            inputs=ccda_file,
            outputs=ccda_ai_output
        )
        analyze_ccd_button.click(  # Redundant button, but kept for UI if needed
            analyze_cda_xml_file_with_ai,  # Call REAL AI function
            inputs=ccd_file,
            outputs=ccd_ai_output
        )
        analyze_pdf_button.click(
            analyze_pdf_file_with_ai, inputs=pdf_file, outputs=pdf_ai_output
        )
        analyze_csv_button.click(
            analyze_csv_file_with_ai, inputs=csv_file, outputs=csv_ai_output
        )

    with gr.Tab(
        "One-Click Discharge Paper (AI)", elem_classes="cyberpunk-tab"
    ):  # New Tab for One-Click Discharge Paper with AI, styled
        gr.Markdown("<h2 style='color:#00FFFF; text-shadow: 0 0 3px #00FFFF;'>One-Click Medical Discharge Paper Generation with AI Content</h2>") # Neon Tab Header
        one_click_ai_pdf_button = gr.Button(
            "Generate Discharge Paper with AI (One-Click)", elem_classes="cyberpunk-button"
        )  # Updated button label and styled
        one_click_ai_pdf_status = gr.Textbox(
            label="Discharge Paper Generation Status (AI)"
        )  # Updated status label
        one_click_ai_pdf_download = gr.File(
            label="Download Discharge Paper (AI)"
        )  # Updated download label

        one_click_ai_pdf_button.click(
            generate_discharge_paper_one_click,  # Use the one-click function that now calls AI
            inputs=[],
            outputs=[one_click_ai_pdf_download, one_click_ai_pdf_status],
        )

    # Connect the patient data buttons
    patient_data_button.click(
        fn=CALLBACK_MANAGER.get_patient_data,
        inputs=None,
        outputs=patient_data_output
    )

    # Connect refresh button to update dashboard
    refresh_btn.click(
        fn=update_dashboard, inputs=None, outputs=dashboard_output
    )

    # Corrected the button click function name here to `generate_pdf_from_meldrx` (No AI PDF)
    meldrx_pdf_button.click(
        fn=generate_pdf_from_meldrx,
        inputs=patient_data_output,
        outputs=[meldrx_pdf_download, meldrx_pdf_status]
    )

    # Connect patient data updates to dashboard
    patient_data_button.click(
        fn=update_dashboard, inputs=None, outputs=dashboard_output
    )

# Launch with sharing enabled for public access
demo.launch(ssr_mode=False)