callbackmanager.py

import gradio as gr
from meldrx import MeldRxAPI  
import json
import os
import tempfile
from datetime import datetime

# Import PDF utilities
from pdfutils import PDFGenerator, generate_discharge_summary

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 get_auth_url(self) -> str:
        return self.api.get_authorization_url()

    def set_auth_code(self, code: str) -> str:
        self.auth_code = code
        if self.api.authenticate_with_code(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 code."

    def get_patient_data(self) -> str:
        if not self.access_token:
            return "Not authenticated. Please provide a valid authorization code first."
        patients = self.api.get_patients()
        if patients is not None:
            return json.dumps(patients, indent=2) if patients else "No patient data returned."
        return "Failed to retrieve patient data."

    def get_patient_documents(self, patient_id: str = None):
        """Fetch patient documents from MeldRx"""
        if not self.access_token:
            return "Not authenticated. Please provide a valid authorization code first."
        
        try:
            # This would call the actual MeldRx API to get documents for a specific patient
            # For demonstration, we'll return mock document data
            return [
                {
                    "doc_id": "doc123",
                    "type": "clinical_note",
                    "date": "2023-01-16",
                    "author": "Dr. Sample Doctor",
                    "content": "Patient presented with symptoms of respiratory distress...",
                },
                {
                    "doc_id": "doc124",
                    "type": "lab_result",
                    "date": "2023-01-17",
                    "author": "Lab System",
                    "content": "CBC results: WBC 7.5, RBC 4.2, Hgb 14.1...",
                }
            ]
        except Exception as e:
            return f"Error retrieving patient documents: {str(e)}"
        
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"""
    **Patient Discharge Form**
    - Name: {first_name} {middle_initial} {last_name}
    - Date of Birth: {dob}, Age: {age}, Sex: {sex}
    - Address: {address}, {city}, {state}, {zip_code}
    - Doctor: {doctor_first_name} {doctor_middle_initial} {doctor_last_name}
    - Hospital/Clinic: {hospital_name}
    - Doctor Address: {doctor_address}, {doctor_city}, {doctor_state}, {doctor_zip}
    - Admission Date: {admission_date}, Source: {referral_source}, Method: {admission_method}
    - Discharge Date: {discharge_date}, Reason: {discharge_reason}
    - Date of Death: {date_of_death}
    - Diagnosis: {diagnosis}
    - Procedures: {procedures}
    - Medications: {medications}
    - Prepared By: {preparer_name}, {preparer_job_title}
    """

    
    return form

def process_fhir_bundle(fhir_bundle):
    """Extract and structure patient data from FHIR bundle"""
    patients = []
    for entry in fhir_bundle.get("entry", []):
        resource = entry.get("resource", {})
        if resource.get("resourceType") == "Patient":
            patients.append(resource)
    return patients

def create_patient_stripe(patient):
    """Create a UI stripe for an individual patient with enhanced functionality"""
    # Safely extract name components
    official_name = next(
        (n for n in patient.get("name", []) if n.get("use") == "official"),
        {}
    )
    given_names = " ".join(official_name.get("given", ["Unknown"]))
    family_name = official_name.get("family", "Unknown")
    
    # Extract demographic information
    gender = patient.get("gender", "unknown").capitalize()
    birth_date = patient.get("birthDate", "Unknown")
    patient_id = patient.get("id", "Unknown")
    
    # Extract address information
    address = patient.get("address", [{}])[0]
    city = address.get("city", "Unknown")
    state = address.get("state", "")
    postal_code = address.get("postalCode", "")

    with gr.Row(variant="panel") as stripe:
        with gr.Column(scale=3):
            gr.Markdown(f"""
            **{given_names} {family_name}**  
            *{gender} | Born: {birth_date}*  
            {city}, {state} {postal_code}
            """)
            gr.Textbox(patient_id, label="Patient ID", visible=False)
        
        with gr.Column(scale=1):
            status = gr.Dropdown(
                choices=["Pending", "Reviewed", "Approved", "Rejected"],
                value="Pending",
                label="Review Status"
            )
        
        with gr.Column(scale=1):
            generate_btn = gr.Button("Generate Report", visible=False)
            probability = gr.Label("Risk Probability", visible=False)
            documents_btn = gr.Button("View Documents", visible=True)

    # Dynamic visibility and document handling
    def update_components(selected_status):
        visible = selected_status in ["Reviewed", "Approved"]
        return (
            gr.Button(visible=visible),
            gr.Label(visible=visible, value={"Low": 0.3, "Medium": 0.6, "High": 0.9}.get(selected_status, 0.5))
        )

    def show_documents(patient_id):
        return CALLBACK_MANAGER.get_patient_documents(patient_id)

    status.change(
        update_components,
        inputs=status,
        outputs=[generate_btn, probability]
    )
    
    documents_btn.click(
        fn=show_documents,
        inputs=patient_id,
        outputs=gr.JSON(label="Patient Documents")
    )

    return stripe


def generate_pdf_from_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
):
    """Generate a PDF discharge form using the provided data"""
    
    # Create PDF generator
    pdf_gen = PDFGenerator()
    
    # Format data for PDF generation
    patient_info = {
        "first_name": first_name,
        "last_name": last_name,
        "dob": dob,
        "age": age,
        "sex": sex,
        "mobile": "",  # Not collected in the form
        "address": address,
        "city": city,
        "state": state,
        "zip": zip_code
    }
    
    discharge_info = {
        "date_of_admission": admission_date,
        "date_of_discharge": discharge_date,
        "source_of_admission": referral_source,
        "mode_of_admission": admission_method,
        "discharge_against_advice": "Yes" if discharge_reason == "Discharge Against Advice" else "No"
    }
    
    diagnosis_info = {
        "diagnosis": diagnosis,
        "operation_procedure": procedures,
        "treatment": "",  # Not collected in the form
        "follow_up": ""   # Not collected in the form
    }
    
    medication_info = {
        "medications": [medications] if medications else [],
        "instructions": ""  # Not collected in the form
    }
    
    prepared_by = {
        "name": preparer_name,
        "title": preparer_job_title,
        "signature": ""  # Not collected in the form
    }
    
    # Generate PDF
    pdf_buffer = pdf_gen.generate_discharge_form(
        patient_info,
        discharge_info,
        diagnosis_info,
        medication_info,
        prepared_by
    )
    
    # Create temporary file to save the PDF
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.pdf')
    temp_file.write(pdf_buffer.read())
    temp_file_path = temp_file.name
    temp_file.close()
    
    return temp_file_path

def generate_pdf_from_meldrx(patient_data):
    """Generate a PDF using patient data from MeldRx"""
    if isinstance(patient_data, str):
        # If it's a string (error message or JSON string), try to parse it
        try:
            patient_data = json.loads(patient_data)
        except:
            return None, "Invalid patient data format"
    
    if not patient_data:
        return None, "No patient data available"
    
    try:
        # For demonstration, we'll use the first patient in the list if it's a list
        if isinstance(patient_data, list) and len(patient_data):
            patient = patient_data[0]
        else:
            patient = patient_data
            
        # Extract patient info
        patient_info = {
            "name": f"{patient.get('name', {}).get('given', [''])[0]} {patient.get('name', {}).get('family', '')}",
            "dob": patient.get('birthDate', 'Unknown'),
            "patient_id": patient.get('id', 'Unknown'),
            "admission_date": datetime.now().strftime("%Y-%m-%d"),  # Mock data
            "physician": "Dr. Provider"  # Mock data
        }
        
        # Mock LLM-generated content
        llm_content = {
            "diagnosis": "Diagnosis information would be generated by LLM",
            "treatment": "Treatment summary would be generated by LLM",
            "medications": "Medication list would be generated by LLM",
            "follow_up": "Follow-up instructions would be generated by LLM",
            "special_instructions": "Special instructions would be generated by LLM"
        }
        
        # Create discharge summary
        output_dir = tempfile.mkdtemp()
        pdf_path = generate_discharge_summary(patient_info, llm_content, output_dir)
        
        return pdf_path, "PDF generated successfully"
        
    except Exception as e:
        return None, f"Error generating PDF: {str(e)}"

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

with gr.Blocks() as demo:
    gr.Markdown("# Patient Discharge Form with MeldRx Integration")
    
    with gr.Tab("Authenticate with MeldRx"):
        gr.Markdown("## SMART on FHIR Authentication")
        auth_url_output = gr.Textbox(label="Authorization URL", value=CALLBACK_MANAGER.get_auth_url(), interactive=False)
        gr.Markdown("Copy the URL above, open it in a browser, log in, and paste the 'code' from the redirect URL below.")
        auth_code_input = gr.Textbox(label="Authorization Code")
        auth_submit = gr.Button("Submit Code")
        auth_result = gr.Textbox(label="Authentication Result")
        
        patient_data_button = gr.Button("Fetch Patient Data")
        patient_data_output = gr.Textbox(label="Patient Data")
        
        # Add button to generate PDF from MeldRx data
        meldrx_pdf_button = gr.Button("Generate PDF from MeldRx Data")
        meldrx_pdf_status = gr.Textbox(label="PDF Generation Status")
        meldrx_pdf_download = gr.File(label="Download Generated PDF")
        
        auth_submit.click(fn=CALLBACK_MANAGER.set_auth_code, inputs=auth_code_input, outputs=auth_result)
        patient_data_button.click(fn=CALLBACK_MANAGER.get_patient_data, inputs=None, outputs=patient_data_output)
        
        # Add functionality for PDF generation from MeldRx data
        meldrx_pdf_button.click(
            fn=generate_pdf_from_meldrx,
            inputs=patient_data_output,
            outputs=[meldrx_pdf_download, meldrx_pdf_status]
        )

    with gr.tab("Show Patients Log"):
        gr.Markdown("## Patient Dashboard")

        # Store patient data in a state variable
        patient_data_state = gr.State()

         # Pagination controls
        with gr.Row() as pagination_row:
            prev_btn = gr.Button("Previous Page", visible=False)
            next_btn = gr.Button("Next Page", visible=False)
        
        # Patient stripes container
        patient_stripes = gr.Column()
        
        # Metadata display
        metadata_md = gr.Markdown()
        
        # Refresh button
        refresh_btn = gr.Button("Refresh Data")

        def update_patient_display(patient_data_json):
            try:
                fhir_bundle = json.loads(patient_data_json)
                patients = process_fhir_bundle(fhir_bundle)
                
                # Update pagination controls
                has_prev = any(link["relation"] == "previous" for link in fhir_bundle.get("link", []))
                has_next = any(link["relation"] == "next" for link in fhir_bundle.get("link", []))
                
                # Create stripes
                stripes = []
                for patient in patients:
                    stripes.append(create_patient_stripe(patient))
                
                return (
                    gr.Row.update(visible=has_prev or has_next),
                    gr.Column.update(stripes),
                    gr.Markdown.update(value=f"""
                    **Bundle Metadata**  
                    Type: {fhir_bundle.get('type', 'Unknown')}  
                    Total Patients: {fhir_bundle.get('total', 0)}  
                    Last Updated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
                    """),
                    {"patients": patients, "bundle": fhir_bundle}
                )
            except Exception as e:
                return (
                    gr.Row.update(visible=False),
                    gr.Column.update([]),
                    gr.Markdown.update(value=f"Error loading patient data: {str(e)}"),
                    None
                )

        # Connect data flow
        patient_data_output.change(
            fn=update_patient_display,
            inputs=patient_data_output,
            outputs=[pagination_row, patient_stripes, metadata_md, patient_data_state]
        )
        
        refresh_btn.click(
            fn=lambda: CALLBACK_MANAGER.get_patient_data(),
            outputs=patient_data_output
        )
    
    with gr.Tab("Discharge Form"):
        gr.Markdown("## Patient Details")
        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("## Primary Healthcare Professional Details")
        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="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("## Admission and Discharge Details")
        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("## Diagnosis & Procedures")
        diagnosis = gr.Textbox(label="Diagnosis")
        procedures = gr.Textbox(label="Operation & Procedures")
        gr.Markdown("## Medication Details")
        medications = gr.Textbox(label="Medication on Discharge")
        gr.Markdown("## Prepared By")
        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")
            submit_pdf = gr.Button("Generate PDF")
        
        # Output areas
        form_output = gr.Markdown()
        pdf_output = gr.File(label="Download PDF")
        
        # 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
        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
        )

demo.launch()