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

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

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 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

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")
        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)
    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")
        submit = gr.Button("Generate Form")
        output = gr.Markdown()
        submit.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=output
        )

demo.launch()