home.py

import math
import os
from datetime import datetime
from typing import Dict, Any, List, Tuple

import gradio as gr
from gradio_modal import Modal
import pandas as pd

# ----- DB wiring -----
# Expects your SheamiDB class to be in modules/db.py
# from modules.db import SheamiDB
# For illustration, we lazy-import at runtime to avoid import errors if path differs.
from dotenv import load_dotenv

from common import get_db

from modules.models import SheamiUser
from plot_helper import (
    MAX_CHARTS_IN_PAGE,
    NUM_COLS,
    build_trend_figure,
    render_vitals_plot_layout,
    reset_latest_vitals_labels,
    reset_trends,
    reset_vitals_plots,
    update_trends,
)
from ui import (
    empty_state_component,
    flatten_vitals,
    hide_tabs_if_no_patient_selected,
    render_latest_vitals_card_layout,
    render_top_menu_bar,
    clear_component,
    close_side_bar,
    disable_component,
    enable_component,
    get_gradio_block,
    hide_component,
    open_side_bar,
    render_logo_small,
    save_vitals_readings,
    show_component,
    show_no_data_found_if_dataframe_empty,
    show_no_data_found_if_none,
)

load_dotenv(override=True)

# ----- Data shaping helpers -----

def get_status_icon(status):
    if status == "completed":
        status = "✅"
    if status == "inprogress":
        status = "⌛"
    if status == "failed":
        status = "❌"
    return status


# Example: flatten your MongoDB document
def flatten_run_stats(docs):
    rows = []
    for doc in docs:
        latest_milestone_status = get_status_icon(
            doc.get("milestones", [{}])[-1].get("status", "")
        )
        rows.append(
            {
                "id": str(doc.get("_id", "")),
                "upload_date": doc.get("created_at", ""),
                # "user_id": str(doc.get("user_id", "")),
                # "patient_id": str(doc.get("patient_id", "")),
                "source_files": ", ".join(doc.get("source_file_names", [])),
                "action": doc.get("action", ""),
                "status": get_status_icon(doc.get("status")),
                "progress": f"{doc.get("steps_completed", 0)}/{doc.get("steps_total", 0)}",
                # For milestones, show latest milestone name & status
                "latest_milestone": doc.get("milestones", [{}])[-1].get(
                    "milestone", ""
                ),
                "milestone_status": latest_milestone_status,
            }
        )
    df = pd.DataFrame(rows)
    df = df.fillna("-")  # 👈 normalize missing values
    return df


def _fmt(dt: Any) -> str:
    if isinstance(dt, (datetime,)):
        return dt.strftime("%Y-%m-%d %H:%M")
    if isinstance(dt, str):
        return dt
    return ""


def flatten_reports_v2(reports: List[Dict[str, Any]]) -> pd.DataFrame:
    """Flatten lab_results from new parsed_data_v2 structure into a table."""
    rows = []
    for r in reports:
        rid = str(r.get("_id", ""))
        uploaded_at = _fmt(r.get("uploaded_at"))
        file_name = r.get("file_name", "")

        parsed_v2 = r.get("parsed_data_v2") or {}
        lab_results = parsed_v2.get("lab_results", [])

        if not lab_results:
            rows.append(
                {
                    "report_id": rid,
                    "uploaded_at": uploaded_at,
                    "file_name": file_name,
                    "test_name": "",
                    "result_value": "",
                    "test_unit": "",
                    "ref_min": "",
                    "ref_max": "",
                    "ref_raw": "",
                    "test_date": "",
                    "inferred_range": "",
                }
            )
        else:

            def format_result(t: dict, section_name: str = None):
                ref = t.get("test_reference_range", {})
                if not ref:
                    ref = {}
                test_name = t.get("test_name", "")
                if section_name:
                    test_name = f"{section_name} | {test_name}"
                return {
                    "report_id": rid,
                    "uploaded_at": uploaded_at,
                    "file_name": file_name,
                    "test_name": test_name,
                    "result_value": t.get("result_value", ""),
                    "test_unit": t.get("test_unit", ""),
                    "ref_min": ref.get("min", "-"),
                    "ref_max": ref.get("max", "-"),
                    "ref_raw": ref.get("raw", "-"),
                    "test_date": t.get("test_date", ""),
                    "inferred_range": t.get("inferred_range", ""),
                }

            for t in lab_results:
                sub_results = t.get("sub_results", {})
                section_name = t.get("section_name", {})
                if sub_results:
                    for sub_result in sub_results:
                        rows.append(format_result(sub_result, section_name))
                else:
                    rows.append(format_result(t))
    # if not rows:
    #     rows = [
    #         {
    #             "report_id": "",
    #             "uploaded_at": "",
    #             "file_name": "",
    #             "test_name": "",
    #             "result_value": "",
    #             "test_unit": "",
    #             "ref_min": "",
    #             "ref_max": "",
    #             "ref_raw": "",
    #             "test_date": "",
    #             "inferred_range": "",
    #         }
    #     ]
    df = pd.DataFrame(
        data=rows,
        columns=[
            "report_id",
            "uploaded_at",
            "file_name",
            "test_name",
            "result_value",
            "test_unit",
            "ref_min",
            "ref_max",
            "ref_raw",
            "test_date",
            "inferred_range",
        ],
    )
    df = df.fillna("-")  # 👈 normalize missing values
    return df


import tempfile
import os


async def download_final_report(report_id: str) -> str | None:
    """
    Fetch the PDF for a given final_report and return a temporary file path.
    """
    if not report_id:
        return None

    db = get_db()
    pdf_bytes = await db.get_final_report_pdf(report_id)
    if not pdf_bytes:
        return None

    # Write to a temporary file
    tmpdir = tempfile.mkdtemp()
    file_path = os.path.join(tmpdir, f"{report_id}.pdf")
    with open(file_path, "wb") as f:
        f.write(pdf_bytes)

    return file_path


def trends_index(trends: List[Dict[str, Any]]) -> List[str]:
    names = sorted({t.get("test_name", "") for t in trends if t.get("test_name")})
    return names


# ----- App state + loaders -----


async def load_user(email: str) -> Tuple[Dict[str, Any], List[Tuple[str, str]]]:
    """Return (user_dict, patient_choices[(label, value), ...])"""
    if not email:
        return {}, []
    db = get_db()
    user = await db.get_user_by_email(email)
    if not user:
        return {}, []
    # Preload patients
    patients = await db.get_patients_by_user(str(user["_id"]))
    choices = [
        (
            ("🔒" if p.get("locked", False) else "✎ ") + p.get("name", str(p["_id"])),
            str(p["_id"]),
        )
        for p in patients
    ]
    return user, choices


def format_patient_info(patient):
    locked = patient.get("locked", False)
    dob = patient.get("dob")
    if isinstance(dob, str):
        dob_str = dob.split(" ")[0]  # keep only date
    elif dob:
        dob_str = dob.strftime("%Y-%m-%d")
    else:
        dob_str = "-"
    return (
        ("🔒" if locked else "👤")
        + f'**{patient.get("name", "-")}** | **{patient.get("gender", "-")}** | Born: {dob_str}'
    )


async def load_patient_bundle(
    patient_id: str,
):
    """
    Given patient_id, return:
    - reports_df
    - test_names (for dropdown)
    - meta dict with patient basics
    - final_reports list
    - run_stats list
    - vitals history
    - latest vitals
    - vitals plots
    """
    if not patient_id:
        return pd.DataFrame(), [], {}, []
    db = get_db()

    # patient
    patient = await db.get_patient_by_id(patient_id)

    # related
    reports = await db.get_reports_by_patient(patient_id)
    trends = await db.get_trends_by_patient(patient_id)
    finals = await db.get_final_reports_by_patient(patient_id)
    run_stats = await db.get_run_stats_by_patient(patient_id)
    latest_vitals = await render_latest_vitals_card_layout(patient_id)
    vitals_history = await db.get_vitals_by_patient(patient_id)
    vitals_plots = await render_vitals_plot_layout(patient_id)

    reports_df = flatten_reports_v2(reports)
    test_names = trends_index(trends)

    # print("test_names = ", test_names)
    meta = {
        "Patient": patient.get("name", ""),
        "Gender": patient.get("gender", ""),
        "DOB": patient.get("dob", ""),
        "Created": _fmt(patient.get("created_at")),
    }
    selected_patient_info = format_patient_info(patient)
    return (
        reports_df,
        test_names,
        meta,
        selected_patient_info,
        finals,
        run_stats,
        vitals_history,
        latest_vitals,
        vitals_plots,
    )


def format_final_reports(finals: List[Dict[str, Any]]) -> pd.DataFrame:
    rows = []
    for fr in finals:
        rows.append(
            {
                "final_report_id": str(fr.get("_id", "")),
                "generated_at": _fmt(fr.get("generated_at")),
                # "summary": fr.get("summary", ""),
                # "recommendations": "; ".join(fr.get("recommendations", [])),
            }
        )
    return pd.DataFrame(columns=["final_report_id", "generated_at"], data=rows)


async def add_patient_ui(user_email, name, dob, gender):
    db = get_db()
    yield "⌛ Saving ...", gr.update()
    user = await db.get_user_by_email(user_email)
    if not user:
        yield "User not found", gr.update()
        return
    pid = await db.add_patient(user["_id"], name, dob, gender)
    if pid:
        (user_dict, patients) = await load_user(user_email)
        yield f"✅ Patient {name} added (ID: {pid})", gr.update(
            choices=patients, value=pid
        )
        return
    else:
        yield f"❌ Error adding patient", gr.update()
        return


async def edit_patient_ui(user_email, patient_id, name, dob, gender):
    db = get_db()
    yield "⌛ Saving ...", None, None
    success = await db.update_patient(
        patient_id, {"name": name, "dob": dob, "gender": gender}
    )
    if success:
        patient = await db.get_patient_by_id(patient_id)
        (user_dict, patients) = await load_user(user_email)
        yield "✅ Updated", format_patient_info(patient), gr.update(
            choices=patients, value=patient_id
        )
    else:
        yield "❌ Patient not found", None, None
    return


async def delete_patient_ui(user_email, patient_id):
    db = get_db()
    yield "⌛ Deleting ...", gr.update(), gr.update()
    last_delete_log = ""
    async for delete_log in db.delete_patient(patient_id):
        yield delete_log, gr.update(), gr.update()
        last_delete_log = delete_log
    success = "✅" in last_delete_log
    if success:
        (user, patients) = await load_user(user_email)
        yield "✅ Deleted", gr.update(choices=patients, value=None), hide_component()
    else:
        yield "❌ Patient not found", gr.update(), gr.update()
    return


async def get_selected_patient_name_ui(selected_patient):
    db = get_db()
    selected_patient_info = await db.get_patient_by_id(selected_patient)
    print("selected_patient = ", selected_patient_info)
    return selected_patient_info["name"]


async def get_selected_patient_details_ui(selected_patient):
    db = get_db()
    selected_patient_info = await db.get_patient_by_id(selected_patient)
    print("selected_patient = ", selected_patient_info)
    return (
        selected_patient_info["name"],
        selected_patient_info["dob"],
        selected_patient_info["gender"],
    )


# ----- Gradio UI (Sidebar Layout) -----
def build_home_page(
    logged_in_user: SheamiUser,
    selected_patient_info: gr.Markdown,
    delete_patient_btn: gr.Button,
    edit_patient_btn: gr.Button,
    upload_reports_btn: gr.Button,
    add_vitals_btn: gr.Button,
):
    with gr.Column() as sheami_secure_container:
        (
            sheami_logo_container,
            email_in,
            # selected_patient_info,
            # delete_patient_btn,
            # edit_patient_btn,
            # upload_reports_btn,
            # add_vitals_btn,
        ) = render_top_menu_bar(
            logged_in_user=logged_in_user
        )  # create component once
        ### Modals
        with Modal(visible=False) as add_patient_modal:
            render_logo_small()
            gr.Markdown("# ➕ Add Patient")
            new_name = gr.Textbox(label="Name")
            new_dob = gr.DateTime(
                label="Date of birth", include_time=False, type="datetime"
            )
            new_gender = gr.Dropdown(["M", "F", "O"], label="Gender")
            with gr.Row():
                gr.Column()
                save_add_patient_btn = gr.Button("Add", scale=0, variant="primary")
                gr.Column()
            add_out = gr.Markdown(label="Status")

        with Modal(visible=False) as edit_patient_modal:
            render_logo_small()
            edit_name = gr.Textbox(label="Edit Name")
            edit_dob = gr.DateTime(
                label="Edit DOB", include_time=False, type="datetime"
            )
            edit_gender = gr.Dropdown(["M", "F"], label="Edit Gender")
            save_edit_patient_btn = gr.Button(
                "✏️ Save Changes", size="lg", variant="primary"
            )
            edit_out = gr.Markdown(label="Status")

        with Modal(visible=False) as delete_patient_modal:
            render_logo_small()
            gr.Markdown("## Are you sure you want to delete the selected patient?")
            selected_patient_for_delete = gr.Text(
                "<selected patient here>", show_label=False
            )
            with gr.Row():
                gr.Column()
                confirm_delete_patient_btn = gr.Button(
                    "Confirm Delete", variant="stop", size="lg", scale=0
                )
                gr.Column()
            delete_out = gr.Markdown(label="Status")

        with Modal(visible=False) as upload_reports_modal:
            user_email_state = gr.State("")
            patient_id_state = gr.State("")
            # show them if you want
            shown_email = gr.Textbox(
                label="User Email", interactive=False, visible=False
            )
            shown_patient = gr.Textbox(
                label="Patient ID", interactive=False, visible=False
            )

        with gr.Sidebar(open=False) as sheami_sidebar:  # Sidebar
            # render_logo()
            load_patients_btn = gr.Button(
                "📂 Load Patients", variant="huggingface", visible=False
            )

            with gr.Row(equal_height=True):
                # gr.Markdown("### 👥 Patients")
                add_patient_btn = gr.Button("➕ Add Patient", variant="huggingface")
            patient_list = gr.Radio(
                label="👥 Patients",
                choices=[],
                interactive=True,
                elem_classes="pro-radio",
            )
            gr.Markdown("---")
            gr.Markdown("🔒 Sample Patient (read-only)")
            gr.Markdown("✎  Editable Patient")

        with gr.Row():
            with gr.Column(scale=3):  # Main area
                # render_logo_small()
                meta_box = gr.JSON(label="Patient Details", visible=False)
                with gr.Row() as no_patient_selected_container:
                    no_patient_selected_html = empty_state_component(
                        message="Please choose a patient from the sidebar on the left.",
                        title="No patient selected",
                        icon="🩺",
                    )

                with gr.Tabs(visible=False) as home_tabs_container:
                    with gr.Tab("📉 Trends", id=0):
                        # num_cols = gr.Slider(1, 4, value=2, step=1, label="Columns", interactive=True)
                        no_trends_message = empty_state_component(
                            title="No trends available yet",
                            message="Upload your lab reports to unlock deeper insights.",
                            icon="📉",
                        )
                        with gr.Row(visible=False) as trends_pagination_container:
                            gr.Column()
                            prev_btn = gr.Button("⬅️ Prev", scale=0, interactive=False)
                            page_info = gr.Textbox(
                                value="Page 0 / 0",
                                interactive=False,
                                show_label=False,
                                container=False,
                                scale=0,
                            )
                            next_btn = gr.Button("Next ➡️", scale=0, interactive=False)
                            gr.Column()
                        test_names_state = gr.State([])
                        num_cols = gr.State(NUM_COLS)
                        trend_charts = []
                        rows = math.ceil(
                            MAX_CHARTS_IN_PAGE / NUM_COLS
                        )  # default num_cols_per_row columns per row
                        for r in range(rows):
                            with gr.Row():
                                for c in range(
                                    NUM_COLS
                                ):  # num_cols_per_row cols per row
                                    trend_charts.append(gr.Plot(visible=False))
                        # Hidden states
                        current_page = gr.State(0)
                        all_figures = gr.State([])

                        next_btn.click(
                            lambda page, figs: page + 1,
                            [current_page, all_figures],
                            current_page,
                        ).then(
                            update_trends,
                            [patient_list, current_page, num_cols],
                            trend_charts
                            + [current_page, page_info, prev_btn, next_btn],
                        )

                        prev_btn.click(
                            lambda page: max(page - 1, 0), current_page, current_page
                        ).then(
                            update_trends,
                            [patient_list, current_page, num_cols],
                            trend_charts
                            + [current_page, page_info, prev_btn, next_btn],
                        )
                    with gr.Tab("🩺 Vitals", id=1):
                        with Modal(visible=False) as vitals_modal:
                            gr.Markdown("## Enter Vitals")
                            with gr.Row():
                                vital_date_in = gr.DateTime(
                                    value=datetime.now(),
                                    label="Date",
                                    include_time=False,
                                    type="datetime",
                                    interactive=True,
                                )
                                vitals_height = gr.Number(label="Height (cm)")
                                vitals_weight = gr.Number(label="Weight (kg)")
                                vitals_bp_sys = gr.Number(label="BP Systolic")
                                vitals_bp_dia = gr.Number(label="BP Diastolic")
                                vitals_fasting_glucose = gr.Number(
                                    label="Fasting Glucose (mg/dL)"
                                )
                                vitals_pbs = gr.Number(
                                    label="Postprandial Glucose (mg/dL)"
                                )

                                vitals_spo2 = gr.Number(label="SpO₂ (%)")

                            gr.Markdown("#### Add Custom Reading")
                            with gr.Row():
                                vitals_custom_name = gr.Textbox(label="Name")
                                vitals_custom_value = gr.Number(label="Value")
                                vitals_custom_unit = gr.Textbox(
                                    label="Unit", placeholder="e.g., bpm, cm"
                                )

                            save_vitals_btn = gr.Button(
                                "Save Readings", variant="huggingface", scale=0
                            )
                            save_vitals_status = gr.Markdown()
                        no_vitals_message = empty_state_component(
                            icon="🩺",
                            title="No vitals recorded yet",
                            message="Please add your first entry to begin tracking.",
                        )

                        with gr.Column(visible=False) as vitals_container:
                            # Latest Vitals
                            gr.Markdown("## Vitals - Latest")
                            gr.Markdown("---")
                            with gr.Row():
                                latest_vitals_cards = [
                                    gr.Label(visible=False) for _ in range(20)
                                ]
                            with gr.Row():
                                # Historical Vitals
                                with gr.Column(scale=1):
                                    gr.Markdown("## Vitals - History")
                                    gr.Markdown("---")
                                    vitals_df = gr.DataFrame(
                                        headers=[
                                            "Date",
                                            "Name",
                                            "Value",
                                            "Unit",
                                            "AI Status",
                                        ],
                                        interactive=False,
                                    )
                                with gr.Column(scale=2):
                                    gr.Markdown("## Vitals - Trends")
                                    gr.Markdown("---")
                                    vitals_plots = []
                                    rows = math.ceil(
                                        20 / 2
                                    )  # default num_cols_per_row columns per row
                                    for r in range(rows):
                                        with gr.Row():
                                            for c in range(
                                                2
                                            ):  # num_cols_per_row cols per row
                                                vitals_plots.append(
                                                    gr.Plot(visible=True)
                                                )
                    with gr.Tab("🧪 Tests", id=2):
                        no_tests_message = empty_state_component(
                            title="No tests available yet",
                            message="Upload your lab reports to get started.",
                            icon="🧪",
                        )

                        tests_df = gr.DataFrame(
                            headers=[
                                "report_id",
                                "uploaded_at",
                                "file_name",
                                "test_name",
                                "result_value",
                                "test_unit",
                                "reference_range",
                                "test_date",
                                "inferred_range",
                            ],
                            row_count=(0, "dynamic"),
                            wrap=True,
                            interactive=False,
                            visible=False,
                        )
                    with gr.Tab("📊 Downloads", id=3):
                        no_downloads_message = empty_state_component(
                            title="No reports available for download",
                            message="Upload your lab reports to generate insights.",
                            icon="📊",
                        )
                        with gr.Column(visible=False) as downloads_container:
                            gr.Markdown("Click on any row to view/download the report")
                            downloads_df = gr.DataFrame(
                                headers=[
                                    "final_report_id",
                                    "generated_at",
                                    # "summary",
                                    # "recommendations",
                                ],
                                row_count=(0, "dynamic"),
                                wrap=True,
                                interactive=False,
                                visible=True,
                            )
                            gr.Markdown("Click on any row to view/download the report")
                        with Modal(visible=False) as final_report_download_modal:
                            final_report_download_status = gr.Markdown("")
                            render_logo_small()
                            with gr.Row():
                                gr.Column(scale=2)
                                report_id_in = gr.Textbox(
                                    label="Report ID",
                                    placeholder="Paste Report ID to download",
                                    visible=False,
                                )

                                download_final_report_btn = gr.DownloadButton(
                                    "⬇️ Download PDF",
                                    variant="huggingface",
                                    scale=1,
                                    visible=False,
                                )
                                gr.Column(scale=2)

                            html_report_component = gr.HTML()

                        # Function triggered when a row is selected
                        async def on_row_select(evt: gr.SelectData):
                            df_selected_row = evt.row_value
                            print("Selected Row = ", df_selected_row)
                            # df_row is a dict representing the selected row
                            if df_selected_row is None:
                                yield gr.update(
                                    visible=False
                                ), None, None, None, "No Row Selected"
                                return
                            report_id = df_selected_row[0]  # first column is report_id
                            # Return modal visibility + the file to bind to download button
                            yield (
                                show_component(),
                                str(report_id),
                                hide_component(),
                                None,
                                "⌛ Fetching data ...",
                            )  # final_report_download_modal, report_id_in, download_btn, html_report, status
                            report_file_name = await download_final_report(report_id)
                            print("report_file_name = ", report_file_name)
                            report_html_content = await get_db().get_final_report_html(
                                report_id
                            )
                            # Return modal visibility + the file to bind to download button
                            yield (
                                gr.update(visible=True),
                                str(report_id),
                                (
                                    gr.update(value=report_file_name, visible=True)
                                    if report_file_name
                                    else hide_component()
                                ),
                                report_html_content,
                                "",
                            )  # final_report_download_modal, report_id_in, download_btn, html_report
                            return

                        downloads_df.select(
                            on_row_select,
                            outputs=[
                                final_report_download_modal,
                                report_id_in,
                                download_final_report_btn,
                                html_report_component,
                                final_report_download_status,
                            ],
                        ).then(close_side_bar, outputs=[sheami_sidebar])
                        download_final_report_btn.click(
                            fn=download_final_report,
                            inputs=[report_id_in],
                            outputs=download_final_report_btn,
                        )
                    with gr.Tab("⬆️ Uploads", id=4):
                        no_uploads_message = empty_state_component(
                            title="No reports uploaded so far",
                            message="Upload your lab reports to generate insights.",
                            icon="⬆️",
                        )

                        with gr.Column(visible=False) as uploads_container:
                            gr.Markdown("Click on any row to view more details")
                            with Modal(visible=False) as run_stats_modal:
                                gr.Markdown("## Run Stats Details")
                                run_stats_container = gr.Markdown()

                            uploads_df = gr.DataFrame(
                                headers=[
                                    "id",
                                    "upload_date",
                                    "source_files",
                                    "action",
                                    "status",
                                    "progress",
                                    "latest_milestone",
                                    "milestone_status",
                                ],
                                row_count=(0, "dynamic"),
                                wrap=True,
                                interactive=False,
                            )
                            gr.Markdown("Click on any row to view more details")

                        async def on_run_stats_select(evt: gr.SelectData):
                            run_stat_id = evt.row_value[0]  # assuming 'id' is column 0
                            run_stats_details = await get_db().get_run_stats_by_id(
                                run_stat_id
                            )
                            # print(run_stats_details)
                            markdown = ""
                            for i, file_name in enumerate(
                                run_stats_details["source_file_names"]
                            ):
                                text = f"\n### {file_name}\n\n```{(
                                        run_stats_details.get(
                                            "source_file_contents", []
                                        )[i]
                                        if "source_file_contents"
                                        in run_stats_details
                                        else ""
                                    )}```"
                                text = text.replace("\\n", "\n")
                                markdown += text
                            return show_component(), gr.update(value=markdown)

                        uploads_df.select(
                            on_run_stats_select,
                            outputs=[run_stats_modal, run_stats_container],
                        )

        # ---- Events ----
        async def on_load(email):
            user, patient_choices = await load_user(email)
            return gr.update(choices=patient_choices), {}, "⚠ No patient selected"

        load_patients_btn.click(
            on_load,
            inputs=[email_in],
            outputs=[patient_list, meta_box, selected_patient_info],
        )

        async def on_patient_select(patient_id):
            if patient_id:
                (
                    reports_df_val,
                    test_names,
                    meta,
                    selected_patient_info,
                    finals,
                    run_stats,
                    vitals_history,
                    latest_vitals,
                    vitals_plots,
                ) = await load_patient_bundle(patient_id)
                (*trends, page, page_info, prev_btn, next_btn) = await update_trends(
                    patient_id=patient_id, page=0, num_cols=NUM_COLS
                )
                return (
                    meta,
                    selected_patient_info,
                    reports_df_val,
                    test_names[:MAX_CHARTS_IN_PAGE],
                    gr.update(
                        headers=["final_report_id", "generated_at"],
                        value=format_final_reports(finals),
                    ),
                    gr.update(
                        headers=["user_id", "patient_id"],
                        value=flatten_run_stats(run_stats),
                    ),
                    gr.update(
                        value=flatten_vitals(vitals_history),
                    ),
                    *latest_vitals,
                    *vitals_plots,
                    *trends,
                    page,
                    page_info,
                    prev_btn,
                    next_btn,
                )
            else:
                (*trends, page, page_info, prev_btn, next_btn) = await reset_trends()
                (*vitals_plots,) = reset_vitals_plots()
                (*vitals_labels,) = reset_latest_vitals_labels()
                return (
                    None,
                    "⚠ No patient selected",
                    None,
                    [],
                    gr.update(
                        headers=["final_report_id", "generated_at"],
                        value=[],
                    ),
                    gr.update(
                        headers=["user_id", "patient_id"],
                        value=[],
                    ),
                    gr.update(
                        headers=["patient_id"],
                        value=[],
                    ),
                    *vitals_labels,
                    *vitals_plots,
                    *trends,
                    page,
                    page_info,
                    prev_btn,
                    next_btn,
                )

        def toggle_open_state_if_none(selected_item):
            return gr.update(open=True if not selected_item else False)

        def toggle_interactive_status(selected_patient, selected_patient_info):
            return (
                gr.update(
                    interactive=(
                        True
                        if selected_patient and "🔒" not in selected_patient_info
                        else False
                    )
                ),  # edit_patient_btn
                gr.update(
                    interactive=(
                        True
                        if selected_patient and "🔒" not in selected_patient_info
                        else False
                    )
                ),  # delete_patient_btn
                gr.update(
                    interactive=(
                        True
                        if selected_patient and "🔒" not in selected_patient_info
                        else False
                    )
                ),  # upload_reports_btn
                gr.update(
                    interactive=(
                        True
                        if selected_patient and "🔒" not in selected_patient_info
                        else False
                    )
                ),  # add_vitals_btn
            )

        patient_list.change(
            toggle_open_state_if_none, inputs=[patient_list], outputs=[sheami_sidebar]
        ).then(
            hide_tabs_if_no_patient_selected,
            inputs=[patient_list],
            outputs=[home_tabs_container, no_patient_selected_container],
        ).then(
            on_patient_select,
            inputs=[patient_list],
            outputs=[
                meta_box,
                selected_patient_info,
                tests_df,
                test_names_state,
                downloads_df,
                uploads_df,
                vitals_df,
                *latest_vitals_cards,
                *vitals_plots,
                *trend_charts,
                current_page,
                page_info,
                prev_btn,
                next_btn,
            ],
        ).then(
            toggle_interactive_status,
            inputs=[patient_list, selected_patient_info],
            outputs=[
                edit_patient_btn,
                delete_patient_btn,
                upload_reports_btn,
                add_vitals_btn,
            ],
        )

        test_names_state.change(
            show_no_data_found_if_none,
            inputs=[test_names_state],
            outputs=[no_trends_message, trends_pagination_container],
        )
        vitals_df.change(
            show_no_data_found_if_dataframe_empty,
            inputs=[vitals_df],
            outputs=[no_vitals_message, vitals_container],
        )
        tests_df.change(
            show_no_data_found_if_dataframe_empty,
            inputs=[tests_df],
            outputs=[no_tests_message, tests_df],
        )
        downloads_df.change(
            show_no_data_found_if_dataframe_empty,
            inputs=[downloads_df],
            outputs=[no_downloads_message, downloads_container],
        )
        uploads_df.change(
            show_no_data_found_if_dataframe_empty,
            inputs=[uploads_df],
            outputs=[no_uploads_message, uploads_container],
        )

        add_patient_btn.click(close_side_bar, outputs=[sheami_sidebar]).then(
            clear_component, outputs=[add_out]
        ).then(show_component, outputs=[add_patient_modal])
        save_add_patient_btn.click(
            disable_component, outputs=save_add_patient_btn
        ).then(
            add_patient_ui,
            inputs=[email_in, new_name, new_dob, new_gender],
            outputs=[add_out, patient_list],
        ).then(
            hide_component, outputs=[add_patient_modal]
        )

        delete_patient_btn.click(close_side_bar, outputs=[sheami_sidebar]).then(
            clear_component, outputs=[delete_out]
        ).then(disable_component, outputs=confirm_delete_patient_btn).then(
            show_component, outputs=[delete_patient_modal]
        ).then(
            get_selected_patient_name_ui,
            inputs=[patient_list],
            outputs=[selected_patient_for_delete],
        ).then(
            enable_component, outputs=confirm_delete_patient_btn
        )

        confirm_delete_patient_btn.click(
            disable_component, outputs=confirm_delete_patient_btn
        ).then(
            delete_patient_ui,
            inputs=[email_in, patient_list],
            outputs=[delete_out, patient_list, delete_patient_modal],
        )

        edit_patient_btn.click(enable_component, outputs=save_edit_patient_btn).then(
            close_side_bar, outputs=[sheami_sidebar]
        ).then(clear_component, outputs=[edit_out]).then(
            show_component, outputs=[edit_patient_modal]
        ).then(
            get_selected_patient_details_ui,
            inputs=[patient_list],
            outputs=[edit_name, edit_dob, edit_gender],
        )

        save_edit_patient_btn.click(
            disable_component, outputs=save_edit_patient_btn
        ).then(
            edit_patient_ui,
            inputs=[email_in, patient_list, edit_name, edit_dob, edit_gender],
            outputs=[edit_out, selected_patient_info, patient_list],
        ).then(
            hide_component, outputs=[edit_patient_modal]
        )

        add_vitals_btn.click(enable_component, outputs=save_vitals_btn).then(
            lambda: gr.update(selected=1), outputs=[home_tabs_container]
        ).then(close_side_bar, outputs=[sheami_sidebar]).then(
            clear_component, outputs=[save_vitals_status]
        ).then(
            show_component, outputs=[vitals_modal]
        )
        save_vitals_btn.click(disable_component, outputs=save_vitals_btn).then(
            save_vitals_readings,
            inputs=[
                patient_list,
                vital_date_in,
                vitals_height,
                vitals_weight,
                vitals_bp_sys,
                vitals_bp_dia,
                vitals_fasting_glucose,
                vitals_pbs,
                vitals_spo2,
                vitals_custom_name,
                vitals_custom_value,
                vitals_custom_unit,
                email_in,
            ],
            outputs=[
                save_vitals_status,
                vitals_df,
                *latest_vitals_cards,
                *vitals_plots,
            ],
        ).then(enable_component, outputs=save_vitals_btn)

        # open modal and set states
        def show_upload_reports_modal(user_email, patient_id):
            return [
                user_email,
                patient_id,
                gr.update(value=user_email),
                gr.update(value=patient_id),
                gr.update(visible=True),
            ]

        # load upload_reports_modal with the upload screen
        get_gradio_block(
            container=upload_reports_modal,
            user_email_state=user_email_state,
            patient_id_state=patient_id_state,
            fn_callback=on_patient_select,
            fn_callback_inputs=[patient_list],
            fn_callback_outputs=[
                meta_box,
                selected_patient_info,
                tests_df,
                test_names_state,
                downloads_df,
                uploads_df,
                vitals_df,
                *latest_vitals_cards,
                *vitals_plots,
                *trend_charts,
                current_page,
                page_info,
                prev_btn,
                next_btn,
            ],
        )

        upload_reports_btn.click(close_side_bar, outputs=[sheami_sidebar]).then(
            show_upload_reports_modal,
            inputs=[email_in, patient_list],
            outputs=[
                user_email_state,
                patient_id_state,
                shown_email,
                shown_patient,
                upload_reports_modal,
            ],
        )
        # Run the same function on load automatically (simulate click effect)
        gr.on(
            fn=on_load,
            inputs=[email_in],
            outputs=[patient_list, meta_box, selected_patient_info],
        ).then(open_side_bar, outputs=[sheami_sidebar])
        # sheami_sidebar.expand(fn=hide_component, outputs=[sheami_logo_container])
        # sheami_sidebar.collapse(fn=show_component, outputs=[sheami_logo_container])

        def enable_add_patient_if_logged_in(email):
            if email:
                return enable_component()
            else:
                return disable_component()

        email_in.change(
            enable_add_patient_if_logged_in,
            inputs=[email_in],
            outputs=[add_patient_btn],
        )

    return sheami_secure_container