# %%
import gradio as gr
from gradio_calendar import Calendar
from datetime import datetime, timedelta
import os
from datasets import load_dataset, Dataset
from huggingface_hub import login
import pytz 

login(token=os.getenv("sitetoken"))
DATASET_NAME = "Mephesto1/Diabetic-tracker"

def get_fresh_data():
    try:
        return load_dataset(DATASET_NAME)["train"]
    except:
        return Dataset.from_dict({
            "sensor_site": [],
            "sensor_date": [],
            "pod_site": [],
            "pod_date": []
        })

# Initial load
user_data = get_fresh_data()


def save_to_dataset(sensor_site, sensor_date, pod_site, pod_date):
    print("🔄 Save function triggered")
    new_row = {
        "sensor_site": sensor_site,
        "sensor_date": sensor_date,
        "pod_site": pod_site,
        "pod_date": pod_date
    }

    print("Attempting to save:", new_row)

    try:
        # Reload fresh in case old reference is stale
        data = load_dataset(DATASET_NAME)["train"]
    except:
        data = Dataset.from_dict({
            "sensor_site": [],
            "sensor_date": [],
            "pod_site": [],
            "pod_date": []
        })

    updated = data.add_item(new_row)
    try:
        updated.push_to_hub(DATASET_NAME)
        print("✅ Push to hub succeeded.")
    except Exception as e:
        print("❌ Push to hub failed:", e)

    return f"✅ Saved! Sensor: {sensor_site}, Pod: {pod_site}"

def load_last_inputs():
    try:
        refreshed = load_dataset(DATASET_NAME)["train"]
    except:
        refreshed = Dataset.from_dict({
            "sensor_site": [],
            "sensor_date": [],
            "pod_site": [],
            "pod_date": []
        })

    if len(refreshed) > 0:
        # Sort by sensor_date to get the latest change
        sorted_rows = sorted(
            refreshed,
            key=lambda r: datetime.strptime(str(r["sensor_date"]), "%Y-%m-%d") if isinstance(r["sensor_date"], str) else r["sensor_date"]
        )
        last = sorted_rows[-1]
        sensor_site = last["sensor_site"].strip().title()
        pod_site = last["pod_site"].strip().title()
        if sensor_site not in [s["site"] for s in sites]:
            sensor_site = ""
        if pod_site not in [s["site"] for s in sites]:
            pod_site = ""
        return (sensor_site, last["sensor_date"], pod_site, last["pod_date"])

    return "", "", "", ""

def convert_date_format(date_input):
    if isinstance(date_input, datetime):
        return date_input.strftime("%Y-%m-%d")
    elif isinstance(date_input, str):
        return datetime.strptime(date_input, "%m/%d/%Y").strftime("%Y-%m-%d")
    else:
        raise ValueError("Unsupported date format")

sites = [
    {"site": "Right Stomach", "side": "right"},
    {"site": "Right Arm", "side": "right"},
    {"site": "Left Arm", "side": "left"},
    {"site": "Left Stomach", "side": "left"},
    {"site": "Left Leg", "side": "left"},
    {"site": "Right Leg", "side": "right"}
]

sensor_sequence = [
    {"site": "Left Arm", "side": "left"},
    {"site": "Left Leg", "side": "left"},
    {"site": "Right Leg", "side": "right"},
    {"site": "Right Arm", "side": "right"},
    {"site": "Left Arm", "side": "left"},
    {"site": "Left Leg", "side": "left"},
    {"site": "Right Leg", "side": "right"},
    {"site": "Right Arm", "side": "right"}
]

sensor_change_days = 10
pod_change_days = 3

pod_sites_map = {
    "Left Arm": ["Left Leg", "Left Stomach", "Left Arm"],
    "Left Leg": ["Left Arm", "Left Stomach", "Left Leg"],
    "Right Leg": ["Right Arm", "Right Stomach", "Right Leg"],
    "Right Arm": ["Right Leg", "Right Stomach", "Right Arm"]
}

# pod_site → compatible sensor_sites for visibility
pod_sensor_visibility = {
    "Left Leg": ["Left Arm"],
    "Left Stomach": ["Left Arm", "Right Arm"],
    "Left Arm": ["Left Leg"],
    "Right Leg": ["Right Arm"],
    "Right Stomach": ["Left Arm", "Right Arm"],
    "Right Arm": ["Right Leg"]
}


def generate_schedules(sensor_site, sensor_date, pod_site, pod_date, as_list=False):
    sensor_date = convert_date_format(sensor_date)
    pod_date = convert_date_format(pod_date)

    start_date = datetime.strptime(sensor_date, "%Y-%m-%d")
    num_changes = 50 // sensor_change_days
    initial_index = next((i for i, site in enumerate(sensor_sequence) if site["site"] == sensor_site), 0)

    sensor_schedule = []
    for i in range(num_changes):
        change_date = start_date + timedelta(days=i * sensor_change_days)
        site_info = sensor_sequence[(initial_index + i) % len(sensor_sequence)]
        sensor_schedule.append({
            "date": change_date.strftime("%Y-%m-%d"),
            "site": site_info["site"],
            "side": site_info["side"]
        })

    pod_start_date = datetime.strptime(pod_date, "%Y-%m-%d")
    pod_num_changes = 15 // pod_change_days
    pod_schedule = []
    site_rotation = {}
    last_used_pod_site = None  # 🧠 Track last site used to avoid repeats

    sensor_info_for_start = next(
        (s for s in reversed(sensor_schedule) if s["date"] <= pod_date),
        sensor_schedule[0]
    )
    available_pod_sites_start = pod_sites_map.get(sensor_info_for_start["site"], ["Left Stomach", "Right Stomach"])
    starting_pod_index = 0
    if pod_site in available_pod_sites_start:
        starting_pod_index = available_pod_sites_start.index(pod_site)

    for i in range(pod_num_changes):
        pod_change_date = pod_start_date + timedelta(days=i * pod_change_days)
        sensor_info = next(
            (s for s in reversed(sensor_schedule) if s["date"] <= pod_change_date.strftime("%Y-%m-%d")),
            sensor_schedule[0]
        )

        pod_date_obj = pod_change_date.date()
        next_sensor = next(
            (s for s in sensor_schedule if datetime.strptime(s["date"], "%Y-%m-%d").date() > pod_date_obj),
            None
        )

        if next_sensor:
            next_sensor_date = datetime.strptime(next_sensor["date"], "%Y-%m-%d").date()
            days_until_sensor_change = (next_sensor_date - pod_date_obj).days
        else:
            days_until_sensor_change = 999

        available_pod_sites = pod_sites_map.get(sensor_info["site"], ["Left Stomach", "Right Stomach"])

        if not available_pod_sites:
        # Fallback to default pod sites if none are compatible
            available_pod_sites = ["Left Stomach", "Right Stomach"]

        # ✅ Dual compatibility check
        if days_until_sensor_change <= 3:
            compatible_sites = [
                site for site in available_pod_sites
                if sensor_info["site"] in pod_sensor_visibility.get(site, []) and
                   next_sensor["site"] in pod_sensor_visibility.get(site, [])
            ]
            if compatible_sites:
                available_pod_sites = compatible_sites

        # 🔁 Avoid repeating pod site if multiple valid options exist
        if len(available_pod_sites) > 1 and last_used_pod_site in available_pod_sites:
            available_pod_sites = [site for site in available_pod_sites if site != last_used_pod_site]

        key = f"{sensor_info['site']}_{sensor_info['date']}"

        if key not in site_rotation:
            next_index = starting_pod_index
        else:
            prev_index = site_rotation[key]
            next_index = (prev_index + 1) % len(available_pod_sites)
        
        # 🧠 Make sure the index is within bounds
        if next_index >= len(available_pod_sites):
            next_index = 0  # fallback to first available

        selected_pod_site = available_pod_sites[next_index]

        site_rotation[key] = next_index
        last_used_pod_site = selected_pod_site  # 🧠 Store used site for next iteration

        pod_schedule.append({
            "date": pod_change_date.strftime("%Y-%m-%d"),
            "pod_site": selected_pod_site,
            "sensor_site": sensor_info["site"]
        })

    if as_list:
        return sensor_schedule, pod_schedule

    return sensor_schedule, pod_schedule

def get_today_message(sensor_sched_list, pod_sched_list):
    import pytz
    from datetime import datetime

    # Use Eastern Time for "today"
    eastern = pytz.timezone("US/Eastern")
    today = datetime.now(eastern).date()
    print("Today's local date:", today)

    # 🔄 Load last saved entries
    saved_sensor_site, saved_sensor_date, saved_pod_site, saved_pod_date = load_last_inputs()

    # Ensure dates are parsed cleanly
    try:
        saved_sensor_date_obj = datetime.strptime(saved_sensor_date, "%Y-%m-%d").date()
    except Exception:
        saved_sensor_date_obj = None

    try:
        saved_pod_date_obj = datetime.strptime(saved_pod_date, "%Y-%m-%d").date()
    except Exception:
        saved_pod_date_obj = None

    # Flags to avoid suggesting actions already completed today
    already_changed_sensor_today = saved_sensor_date_obj == today
    already_changed_pod_today = saved_pod_date_obj == today

    # Log matches
    for s in sensor_sched_list:
        print("Sensor date in schedule:", s['date'], "→ Match:", datetime.strptime(s['date'], "%Y-%m-%d").date() == today)

    for p in pod_sched_list:
        print("Pod date in schedule:", p['date'], "→ Match:", datetime.strptime(p['date'], "%Y-%m-%d").date() == today)

    # Determine if actions are scheduled (and not already performed)
    sensor_msg = next((
        f"sensor and place it on the {s['site']}"
        for s in sensor_sched_list
        if datetime.strptime(s['date'], "%Y-%m-%d").date() == today and not already_changed_sensor_today
    ), None)

    pod_msg = next((
        f"pod and place it on the {p['pod_site']} (sensor on {p['sensor_site']})"
        for p in pod_sched_list
        if datetime.strptime(p['date'], "%Y-%m-%d").date() == today and not already_changed_pod_today
    ), None)

    # Custom responses based on what’s already done
    if already_changed_sensor_today and already_changed_pod_today:
        return "✅ You already changed both the **sensor** and **pod** today. All set!"

    if already_changed_sensor_today and pod_msg:
        return f"🟢 You changed the sensor today. Now it's time to change the **{pod_msg}**."

    if already_changed_pod_today and sensor_msg:
        return f"🟢 You changed the pod today. Now it's time to change the **{sensor_msg}**."

    if sensor_msg and pod_msg:
        return f"🟡 Looks like you need to change both the **{sensor_msg}** and **{pod_msg}** today!"

    if sensor_msg:
        return f"🟢 Looks like you need to change the **{sensor_msg}** today!"

    if pod_msg:
        return f"🟢 Looks like you need to change the **{pod_msg}** today!"

    return "✅ No changes needed today!"


def format_combined_schedule(sensor_sched, pod_sched):
    events = []
    for s in sensor_sched:
        date_obj = datetime.strptime(s["date"], "%Y-%m-%d")
        events.append({
            "date": date_obj,
            "type": "sensor",
            "text": f"Sensor: {s['site']} on {date_obj.strftime('%-m/%-d/%Y (%A)')}"
        })
    for p in pod_sched:
        date_obj = datetime.strptime(p["date"], "%Y-%m-%d")
        events.append({
            "date": date_obj,
            "type": "pod",
            "text": f"Pod on {date_obj.strftime('%-m/%-d/%Y (%A)')} at {p['pod_site']}"
        })
    events.sort(key=lambda e: e["date"])
    return "\n".join([e["text"] for e in events])

def save_and_refresh(sensor_site, sensor_date, pod_site, pod_date):
    gif = "butterfly_cute.gif"
    
    # 💾 Save user selections
    save_to_dataset(sensor_site, sensor_date, pod_site, pod_date)

    refreshed = get_fresh_data()
    if len(refreshed) > 0:
        sorted_rows = sorted(
            refreshed,
            key=lambda r: datetime.strptime(str(r["sensor_date"]), "%Y-%m-%d") if isinstance(r["sensor_date"], str) else r["sensor_date"]
        )
        last = sorted_rows[-1]
        last_sensor_site = last["sensor_site"].strip().title()
        last_sensor_date = last["sensor_date"]
        last_pod_site = last["pod_site"].strip().title()
        last_pod_date = last["pod_date"]
    else:
        last_sensor_site, last_sensor_date, last_pod_site, last_pod_date = "", "", "", ""

    # 📅 Regenerate updated schedules
    sensor_sched, pod_sched = generate_schedules(
        last_sensor_site, last_sensor_date, last_pod_site, last_pod_date, as_list=True
    )

    # 🧾 Format results for display
    combined = format_combined_schedule(sensor_sched, pod_sched)
    return gif, combined

# 🎛️ Interface Layout
with gr.Blocks() as demo:
    gr.Markdown("### Hello, let’s plan your pod and sensor change! 🐾🦋")

    last_sensor_site, last_sensor_date, last_pod_site, last_pod_date = load_last_inputs()
    dropdown_sites = [site["site"] for site in sites]

    gr.Image(value="chocolate_lab_butterfly_pic.png", show_label=False)

    sensor_sched_raw, pod_sched_raw = generate_schedules(
        last_sensor_site, last_sensor_date, last_pod_site, last_pod_date, as_list=True
    )
    gr.Markdown(get_today_message(sensor_sched_raw, pod_sched_raw))

    sensor_site = gr.Dropdown(dropdown_sites, label="Sensor Site", value=last_sensor_site)
    sensor_date = Calendar(type="datetime", label="Sensor Change Date", value=last_sensor_date)
    pod_site = gr.Dropdown(dropdown_sites, label="Pod Site", value=last_pod_site)
    pod_date = Calendar(type="datetime", label="Pod Change Date", value=last_pod_date)

    paw_button = gr.Button("🐾 Press your paw here to save!")
    butterfly_display = gr.Image(show_label=False)
    schedule_box = gr.Textbox(label="Upcoming Schedule", lines=12, value=format_combined_schedule(sensor_sched_raw, pod_sched_raw), interactive=False)

    paw_button.click(
        fn=save_and_refresh,
        inputs=[sensor_site, sensor_date, pod_site, pod_date],
        outputs=[butterfly_display, schedule_box]
    )

if __name__ == "__main__":
    demo.launch()