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

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import PDF utilities
from 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
# Assuming your AI analysis functions are in the same script or imported
# For now, let's define placeholder AI analysis functions for Gradio context
def analyze_dicom_file_with_ai(dicom_file):
    return "DICOM Analysis Report (Placeholder - Real AI integration needed)"

def analyze_hl7_file_with_ai(hl7_file):
    return "HL7 Analysis Report (Placeholder - Real AI integration needed)"

def analyze_cda_xml_file_with_ai(cda_xml_file):
    return "CCDA/XML Analysis Report (Placeholder - Real AI integration needed)"

def analyze_pdf_file_with_ai(pdf_file):
    return "PDF Analysis Report (Placeholder - Real AI integration needed)"

def analyze_csv_file_with_ai(csv_file):
    return "CSV Analysis Report (Placeholder - Real AI integration needed)"


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:
        """Fetch patient data from MeldRx"""
        try:
            if not self.access_token:
                logger.warning("Not authenticated when getting patient data")
                return "Not authenticated. Please provide a valid authorization code first."

            # For demo purposes, if there's no actual API connected, return mock data
            # Remove this in production and use the real API call
            if not hasattr(self.api, 'get_patients') or self.api.get_patients is None:
                logger.info("Using mock patient data (no API connection)")
                # Return mock FHIR bundle with patient data
                mock_data = {
                    "resourceType": "Bundle",
                    "type": "searchset",
                    "total": 2,
                    "link": [],
                    "entry": [
                        {
                            "resource": {
                                "resourceType": "Patient",
                                "id": "patient1",
                                "name": [
                                    {
                                        "use": "official",
                                        "family": "Smith",
                                        "given": ["John"]
                                    }
                                ],
                                "gender": "male",
                                "birthDate": "1970-01-01",
                                "address": [
                                    {
                                        "city": "Boston",
                                        "state": "MA",
                                        "postalCode": "02108"
                                    }
                                ]
                            }
                        },
                        {
                            "resource": {
                                "resourceType": "Patient",
                                "id": "patient2",
                                "name": [
                                    {
                                        "use": "official",
                                        "family": "Johnson",
                                        "given": ["Jane"]
                                    }
                                ],
                                "gender": "female",
                                "birthDate": "1985-05-15",
                                "address": [
                                    {
                                        "city": "Cambridge",
                                        "state": "MA",
                                        "postalCode": "02139"
                                    }
                                ]
                            }
                        }
                    ]
                }
                return json.dumps(mock_data, indent=2)

            # Real implementation with API call
            logger.info("Calling Meldrx API to get patients")
            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."
        except Exception as e:
            error_msg = f"Error in get_patient_data: {str(e)}"
            logger.error(error_msg)
            return f"Error retrieving patient data: {str(e)}"

    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 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)}"

def generate_discharge_paper_one_click():
    """One-click function to fetch patient data and generate discharge paper."""
    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)
        pdf_path, status_message = generate_pdf_from_meldrx(patient_data)
        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)}"


# Create a simplified interface to avoid complex component interactions
CALLBACK_MANAGER = CallbackManager(
    redirect_uri="https://multitransformer-discharge-guard.hf.space/callback",
    client_secret=None
)

# Create the UI
with gr.Blocks() as demo:
    gr.Markdown("# Patient Discharge Form with MeldRx & Medical File Analysis")

    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", lines=10)

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

    with gr.Tab("Patient Dashboard"):
        gr.Markdown("## Patient Data")
        dashboard_output = gr.HTML("<p>Fetch patient data from the Authentication tab first.</p>")

        refresh_btn = gr.Button("Refresh Data")

        # Simple function to update dashboard based on fetched data
        def update_dashboard():
            try:
                data = CALLBACK_MANAGER.get_patient_data()
                if data.startswith("Not authenticated") or data.startswith("Failed") or data.startswith("Error"):
                    return f"<p>{data}</p>"

                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 for each patient
                    html = "<h3>Patients</h3>"
                    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
                        html += f"""
                        <div style="border: 1px solid #ddd; padding: 10px; margin: 10px 0; border-radius: 5px;">
                            <h4>{given} {family}</h4>
                            <p><strong>Gender:</strong> {gender}</p>
                            <p><strong>Birth Date:</strong> {birth_date}</p>
                            <p><strong>ID:</strong> {patient.get("id", "Unknown")}</p>
                        </div>
                        """

                    return html
                except Exception as e:
                    return f"<p>Error parsing patient data: {str(e)}</p>"
            except Exception as e:
                return f"<p>Error fetching patient data: {str(e)}</p>"

    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
        )

    with gr.Tab("Medical File Analysis"):
        gr.Markdown("## Analyze Medical Files with DocuNexus AI")
        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")

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

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

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

            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") # Redundant


        # Connect AI Analysis Buttons (using placeholder functions for now)
        analyze_dicom_button.click(
            lambda file: analyze_dicom_file_with_ai(file.name) if file else "No DICOM file uploaded",
            inputs=dicom_file, outputs=dicom_ai_output
        )
        analyze_hl7_button.click(
            lambda file: analyze_hl7_file_with_ai(file.name) if file else "No HL7 file uploaded",
            inputs=hl7_file, outputs=hl7_ai_output
        )
        analyze_xml_button.click(
            lambda file: analyze_cda_xml_file_with_ai(file.name) if file else "No XML file uploaded", # Using CCDA/XML analyzer for generic XML for now
            inputs=xml_file, outputs=xml_ai_output
        )
        analyze_ccda_button.click(
            lambda file: analyze_cda_xml_file_with_ai(file.name) if file else "No CCDA file uploaded", # Using CCDA/XML analyzer
            inputs=ccda_file, outputs=ccda_ai_output
        )
        analyze_ccd_button.click( # Redundant button, but kept for UI if needed
            lambda file: analyze_cda_xml_file_with_ai(file.name) if file else "No CCD file uploaded", # Using CCDA/XML analyzer
            inputs=ccd_file, outputs=ccd_ai_output
        )

    with gr.Tab("One-Click Discharge Paper"): # New Tab for One-Click Discharge Paper
        gr.Markdown("## One-Click Medical Discharge Paper Generation")
        one_click_pdf_button = gr.Button("Generate Discharge Paper (One-Click)")
        one_click_pdf_status = gr.Textbox(label="Discharge Paper Generation Status")
        one_click_pdf_download = gr.File(label="Download Discharge Paper")

        one_click_pdf_button.click(
            generate_discharge_paper_one_click,
            inputs=[],
            outputs=[one_click_pdf_download, one_click_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
    )

    # 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]
    )

    # 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(share=True)