import json
import os
import re
from html import escape
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load knowledge base
with open("knowledge_base.json", "r") as f:
    KNOWLEDGE_BASE = json.load(f)

# Model configuration. Keep the default tiny so free CPU Spaces stay responsive.
MODEL_ID = os.environ.get("MODEL_ID", "HuggingFaceTB/SmolLM2-135M-Instruct")
TRUST_REMOTE_CODE = MODEL_ID.startswith("openbmb/")
PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "1") == "1"
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "420"))
ACTIVITY_CONTEXT_LIMIT = int(os.environ.get("ACTIVITY_CONTEXT_LIMIT", "10"))

APP_CSS = """
    :root {
        --ink: #151a45;
        --muted: #6f7390;
        --purple: #6954e8;
        --purple-soft: #f0ecff;
        --green: #1d9b6c;
        --green-soft: #effaf4;
        --gold: #a66d12;
        --gold-soft: #fff8e8;
        --red: #c9493d;
        --red-soft: #fff1ed;
        --blue: #265fc5;
        --blue-soft: #eef7ff;
        --line: rgba(105, 84, 232, 0.18);
        --shadow: 0 18px 50px rgba(70, 58, 130, 0.10);
    }

    body,
    .gradio-container {
        background:
            radial-gradient(circle at 12% 4%, rgba(255, 223, 143, 0.34), transparent 28%),
            radial-gradient(circle at 86% 12%, rgba(198, 188, 255, 0.34), transparent 30%),
            linear-gradient(135deg, #fffdf8 0%, #fbf8ff 52%, #f5fbff 100%);
        color: var(--ink);
        font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
    }

    .gradio-container {
        max-width: 1120px !important;
        margin: 0 auto !important;
        padding: 32px 28px 44px !important;
    }

    .hero {
        display: flex;
        align-items: center;
        justify-content: space-between;
        gap: 22px;
        margin-bottom: 22px;
    }

    .brand-lockup {
        display: flex;
        align-items: center;
        gap: 16px;
    }

    .brand-mark,
    .panel-icon,
    .section-icon {
        display: inline-flex;
        align-items: center;
        justify-content: center;
        width: 48px;
        height: 48px;
        border-radius: 50%;
        background: linear-gradient(145deg, #ede8ff, #fff3cf);
        color: var(--purple);
        font-weight: 900;
        box-shadow: inset 0 0 0 1px rgba(105, 84, 232, 0.14);
    }

    .brand-mark svg,
    .panel-icon svg,
    .section-icon svg {
        width: 23px;
        height: 23px;
        fill: none;
        stroke: currentColor;
        stroke-width: 2;
        stroke-linecap: round;
        stroke-linejoin: round;
    }

    .brand-title {
        font-size: 36px;
        line-height: 1;
        font-weight: 900;
        letter-spacing: 0;
        margin: 0 0 7px;
        color: var(--ink);
    }

    .brand-copy,
    .panel-subtitle,
    .plan-copy,
    .card-copy,
    .tip-strip {
        color: var(--muted);
        font-size: 14px;
        line-height: 1.55;
    }

    .idea-pill {
        min-width: 300px;
        padding: 18px 22px;
        border-radius: 8px;
        background: linear-gradient(135deg, rgba(243, 239, 255, 0.95), rgba(255, 255, 255, 0.86));
        border: 1px solid rgba(105, 84, 232, 0.12);
        box-shadow: var(--shadow);
        color: var(--ink);
    }

    .idea-pill strong {
        display: block;
        margin-bottom: 2px;
    }

    .form-panel {
        padding: 22px 22px 18px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.78);
        border: 1px solid rgba(105, 84, 232, 0.14);
        box-shadow: var(--shadow);
        margin-bottom: 18px;
    }

    .form-panel,
    .form-panel > .wrap,
    .form-panel > div {
        background: rgba(255, 255, 255, 0.76) !important;
    }

    .form-panel .block,
    .form-panel .form,
    .form-panel .prose,
    .form-panel .gradio-html,
    .form-panel .gradio-markdown {
        background: transparent !important;
        border: 0 !important;
        box-shadow: none !important;
    }

    .form-panel.family-panel {
        border-color: rgba(255, 154, 108, 0.28);
    }

    .form-panel.session-panel {
        border-color: rgba(105, 84, 232, 0.22);
        background: linear-gradient(180deg, rgba(255, 255, 255, 0.86), rgba(249, 246, 255, 0.72));
    }

    .panel-heading {
        display: flex;
        align-items: center;
        gap: 14px;
        margin-bottom: 18px;
        padding-bottom: 14px;
        border-bottom: 1px solid rgba(105, 84, 232, 0.12);
    }

    .panel-title {
        font-size: 20px;
        font-weight: 850;
        color: var(--purple);
        margin: 0 0 3px;
    }

    .panel-subtitle {
        margin: 0;
    }

    .section-hero {
        display: grid;
        grid-template-columns: minmax(0, 1fr) auto;
        align-items: center;
        min-height: 132px;
        padding: 28px 34px;
        margin-bottom: 20px;
        border-radius: 8px;
        background:
            radial-gradient(circle at 88% 26%, rgba(122, 93, 246, 0.26), transparent 32%),
            linear-gradient(135deg, #f2edff, #fff7f1);
        border: 1px solid rgba(105, 84, 232, 0.12);
        overflow: hidden;
    }

    .section-hero h2 {
        margin: 0 0 8px;
        font-size: 26px;
        color: var(--ink);
    }

    .toy-scene {
        position: relative;
        width: 170px;
        height: 102px;
    }

    .bear-big,
    .bear-small,
    .block-a,
    .block-b,
    .block-c {
        position: absolute;
        border-radius: 50%;
    }

    .bear-big {
        width: 92px;
        height: 92px;
        right: 72px;
        bottom: 2px;
        background: #a58cf0;
        box-shadow: inset -13px -14px 0 rgba(84, 61, 190, 0.12);
    }

    .bear-big::before,
    .bear-small::before {
        content: "";
        position: absolute;
        inset: 24px 27px 0;
        width: 26px;
        height: 20px;
        border-radius: 50%;
        background: #f8dfc4;
    }

    .bear-small {
        width: 54px;
        height: 54px;
        right: 10px;
        bottom: 0;
        background: #f2c878;
    }

    .block-a,
    .block-b,
    .block-c {
        border-radius: 8px;
        bottom: 4px;
        width: 32px;
        height: 32px;
    }

    .block-a { right: 80px; background: #ff896f; }
    .block-b { right: 49px; background: #77c9db; }
    .block-c { right: 64px; bottom: 36px; background: #f3be57; }

    .custom-output {
        display: flex;
        flex-direction: column;
        gap: 18px;
    }

    .plan-card {
        position: relative;
        padding: 24px;
        border-radius: 8px;
        border: 1px solid var(--line);
        background: rgba(255, 255, 255, 0.84);
        box-shadow: var(--shadow);
        overflow: hidden;
    }

    .plan-card.assessment { background: linear-gradient(135deg, #fff, #fbf8ff); }
    .plan-card.recommended { background: linear-gradient(135deg, var(--green-soft), #fff); border-color: rgba(29, 155, 108, 0.22); }
    .plan-card.alternative { background: linear-gradient(135deg, var(--gold-soft), #fff); border-color: rgba(166, 109, 18, 0.22); }
    .plan-card.fallback { background: linear-gradient(135deg, var(--red-soft), #fff); border-color: rgba(201, 73, 61, 0.18); }

    .card-heading {
        display: flex;
        align-items: center;
        gap: 13px;
        margin-bottom: 16px;
    }

    .card-heading h3 {
        margin: 0;
        font-size: 22px;
        color: var(--ink);
    }

    .recommended .card-heading h3 { color: var(--green); }
    .alternative .card-heading h3 { color: var(--gold); }
    .fallback .card-heading h3 { color: var(--red); }
    .assessment-grid,
    .metric-grid {
        display: grid;
        grid-template-columns: repeat(3, minmax(0, 1fr));
        gap: 12px;
    }

    .fact,
    .metric {
        padding: 13px 14px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.72);
        border: 1px solid rgba(105, 84, 232, 0.10);
    }

    .fact-label,
    .metric-label {
        display: block;
        margin-bottom: 5px;
        color: var(--muted);
        font-size: 12px;
        font-weight: 750;
        text-transform: uppercase;
        letter-spacing: 0;
    }

    .fact-value,
    .metric-value {
        color: var(--ink);
        font-size: 14px;
        font-weight: 700;
        line-height: 1.45;
    }

    .timeline {
        display: grid;
        gap: 13px;
    }

    .timeline-row {
        display: grid;
        grid-template-columns: 112px minmax(0, 1fr);
        gap: 14px;
        align-items: stretch;
    }

    .time-chip {
        display: flex;
        align-items: center;
        justify-content: center;
        min-height: 58px;
        border-radius: 8px;
        background: rgba(29, 155, 108, 0.10);
        color: var(--green);
        border: 1px solid rgba(29, 155, 108, 0.18);
        font-weight: 850;
        font-size: 13px;
    }

    .timeline-item {
        padding: 15px 17px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.80);
        border: 1px solid rgba(29, 155, 108, 0.12);
    }

    .timeline-title {
        font-weight: 850;
        margin-bottom: 5px;
        color: var(--ink);
    }

    .timeline-reason {
        color: var(--muted);
        font-size: 13px;
        line-height: 1.45;
    }

    .plan-title {
        margin: -4px 0 14px 62px;
        color: var(--muted);
        font-weight: 700;
    }

    .step-list {
        display: grid;
        gap: 10px;
        margin-top: 12px;
    }

    .step-item {
        display: grid;
        grid-template-columns: 88px minmax(0, 1fr);
        gap: 12px;
        align-items: start;
        padding: 14px 16px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.72);
        border: 1px solid rgba(105, 84, 232, 0.10);
    }

    .step-label {
        color: var(--purple);
        font-size: 12px;
        font-weight: 850;
        line-height: 1.35;
        text-transform: uppercase;
    }

    .step-action {
        color: var(--ink);
        font-size: 14px;
        font-weight: 650;
        line-height: 1.45;
    }

    .plan-note,
    .metric-note {
        color: var(--muted);
        font-size: 13px;
        line-height: 1.5;
    }

    .plan-note {
        margin-top: 13px;
    }

    .metric-note {
        margin-top: 6px;
    }

    .tradeoff-row {
        display: grid;
        grid-template-columns: repeat(3, minmax(0, 1fr));
        gap: 10px;
        margin-top: 14px;
    }

    .tradeoff-chip {
        padding: 10px 12px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.76);
        border: 1px solid rgba(105, 84, 232, 0.10);
    }

    .tradeoff-label {
        display: block;
        color: var(--muted);
        font-size: 11px;
        font-weight: 800;
        text-transform: uppercase;
    }

    .tradeoff-value {
        display: block;
        margin-top: 4px;
        color: var(--ink);
        font-weight: 850;
    }

    .tip-strip {
        margin-top: 18px;
        padding: 15px 18px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.70);
        border: 1px solid rgba(105, 84, 232, 0.18);
        text-align: center;
    }

    #generate-btn {
        border-radius: 8px !important;
        min-height: 52px;
        font-weight: 850;
        background: linear-gradient(135deg, #7d60f2, #5b42d5) !important;
        border: 0 !important;
        box-shadow: 0 14px 28px rgba(91, 66, 213, 0.24);
    }

    input[type="checkbox"]:checked,
    .progress-bar,
    .progress-bar > div {
        accent-color: var(--purple) !important;
        background-color: var(--purple) !important;
    }

    #preferences-input textarea {
        min-height: 96px !important;
        resize: vertical;
    }

    #plan-output {
        margin-bottom: 34px;
    }

    #plan-output .progress,
    #plan-output .progress-bar,
    #plan-output [role="progressbar"] {
        margin-bottom: 24px !important;
    }

    .footer-note {
        margin-top: 72px;
        padding-top: 18px;
        border-top: 1px solid rgba(105, 84, 232, 0.14);
        color: var(--muted);
        font-size: 13px;
        text-align: center;
    }

    @media (max-width: 860px) {
        .gradio-container {
            max-width: 100% !important;
            padding: 18px 14px 30px !important;
            overflow-x: hidden;
        }

        .hero,
        .section-hero {
            grid-template-columns: 1fr;
            flex-direction: column;
            align-items: flex-start;
            gap: 14px;
        }

        .brand-lockup {
            align-items: flex-start;
        }

        .brand-title {
            font-size: 28px;
            line-height: 1.08;
        }

        .idea-pill {
            min-width: 0;
            width: 100%;
            padding: 14px 16px;
        }

        .form-panel,
        .plan-card,
        .section-hero {
            padding: 18px !important;
        }

        .panel-heading,
        .card-heading {
            align-items: flex-start;
        }

        .section-icon,
        .panel-icon,
        .brand-mark {
            width: 40px;
            height: 40px;
            flex: 0 0 40px;
        }

        .assessment-grid,
        .metric-grid,
        .tradeoff-row,
        .timeline-row {
            grid-template-columns: 1fr;
        }

        .plan-title {
            margin-left: 0;
        }

        .step-item {
            grid-template-columns: 1fr;
        }

        .toy-scene {
            display: none;
        }

        #generate-btn {
            width: 100%;
        }

        #preferences-input textarea {
            min-height: 112px !important;
        }
    }
"""

def patch_transformers_for_minicpm():
    """Restore a helper MiniCPM remote code imports from older Transformers."""
    import transformers.utils.import_utils as import_utils

    if not hasattr(import_utils, "is_torch_fx_available"):
        import_utils.is_torch_fx_available = lambda: False

def load_model():
    print(f"Loading model {MODEL_ID}...")
    if TRUST_REMOTE_CODE:
        patch_transformers_for_minicpm()

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        trust_remote_code=TRUST_REMOTE_CODE
    )

    # Detect if CUDA is available
    if torch.cuda.is_available():
        torch_dtype = torch.float16
        device_map = "auto"
    else:
        torch_dtype = torch.float32
        device_map = None

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        trust_remote_code=TRUST_REMOTE_CODE,
        torch_dtype=torch_dtype,
        device_map=device_map,
        use_safetensors=True
    )

    if device_map is None:
        model = model.to("cpu")

    model.eval()
    print("Model loaded.")
    return tokenizer, model

# Global model and tokenizer
TOKENIZER = None
MODEL = None

def get_model():
    global TOKENIZER, MODEL
    if TOKENIZER is None or MODEL is None:
        TOKENIZER, MODEL = load_model()
    return TOKENIZER, MODEL

# Prompt engineering
SYSTEM_PROMPT = """You are Parent Co-Pilot, an AI family situation planner. Your job is to help parents navigate difficult real-world situations by generating realistic, constraint-aware plans.

You are NOT an activity recommendation engine. You are a family situation planning system.

The user's goal is: "Help me get through the next X amount of time."
Activities are merely tools used to achieve that goal.

When given a family situation, you must:
1. Analyze the situation and identify primary and secondary constraints
2. Prioritize what matters most right now
3. Generate a realistic, timeline-based plan
4. Generate a concrete alternative mini-plan with different tradeoffs
5. Generate a concrete emergency fallback plan for when the primary plan fails
6. Provide tradeoff analysis that explains parent effort, interruption risk, and cleanup burden

Rules:
- Be realistic about parent energy. Do not suggest high-energy activities when parent is exhausted.
- Treat sick-child, sick-parent, and both-sick situations differently.
- If a child is sick, prioritize comfort, rest, hydration, low stimulation, and easy supervision.
- If a parent is sick, minimize parent movement, setup, cleanup, and verbal coaching.
- If both parent and child are sick, plan for survival mode: rest, fluids, comfort, safe screen time if needed, and no outdoor or physical activities.
- Use the provided activity library as ingredients, but combine them intelligently.
- Choose activities by reasoning from the household, child age, time horizon, environment, parent energy, current situation, and preferences. Do not choose activities because they appear early in the list.
- Treat the activity list as unordered candidate ingredients.
- Do not make screen time the default recommended plan. Use it mainly as an alternative or emergency support unless the parent is sick, both are sick, or the user explicitly asks for it.
- Plans must account for child ages and household structure.
- Always include a lower-effort alternative and an emergency fallback.
- The alternative plan must include specific steps and must not be another emergency fallback.
- The emergency fallback must include specific steps for immediate stabilization.
- Output MUST be valid JSON only. No markdown, no explanations outside the JSON."""

def normalize_situations(situations, situation_detail):
    selected = situations or []
    if isinstance(selected, str):
        selected = [selected]
    detail = situation_detail.strip() if situation_detail else ""
    return ", ".join(selected + ([detail] if detail else []))

def includes_any(text, terms):
    return any(term in text for term in terms)

def icon_svg(name):
    icons = {
        "house-heart": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M8.62 13.8A2.25 2.25 0 1 1 12 10.836a2.25 2.25 0 1 1 3.38 2.966l-2.626 2.856a.998.998 0 0 1-1.507 0z"/><path d="M3 10a2 2 0 0 1 .709-1.528l7-6a2 2 0 0 1 2.582 0l7 6A2 2 0 0 1 21 10v9a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2z"/></svg>""",
        "users": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M16 21v-2a4 4 0 0 0-4-4H6a4 4 0 0 0-4 4v2"/><path d="M16 3.128a4 4 0 0 1 0 7.744"/><path d="M22 21v-2a4 4 0 0 0-3-3.87"/><circle cx="9" cy="7" r="4"/></svg>""",
        "clock": """<svg viewBox="0 0 24 24" aria-hidden="true"><circle cx="12" cy="12" r="10"/><path d="M12 6v6l4 2"/></svg>""",
        "calendar-clock": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M16 14v2.2l1.6 1"/><path d="M16 2v4"/><path d="M21 7.5V6a2 2 0 0 0-2-2H5a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h3.5"/><path d="M3 10h5"/><path d="M8 2v4"/><circle cx="16" cy="16" r="6"/></svg>""",
        "compass": """<svg viewBox="0 0 24 24" aria-hidden="true"><circle cx="12" cy="12" r="10"/><path d="m16.24 7.76-1.804 5.411a2 2 0 0 1-1.265 1.265L7.76 16.24l1.804-5.411a2 2 0 0 1 1.265-1.265z"/></svg>""",
        "clipboard-list": """<svg viewBox="0 0 24 24" aria-hidden="true"><rect width="8" height="4" x="8" y="2" rx="1" ry="1"/><path d="M16 4h2a2 2 0 0 1 2 2v14a2 2 0 0 1-2 2H6a2 2 0 0 1-2-2V6a2 2 0 0 1 2-2h2"/><path d="M12 11h4"/><path d="M12 16h4"/><path d="M8 11h.01"/><path d="M8 16h.01"/></svg>""",
        "check": """<svg viewBox="0 0 24 24" aria-hidden="true"><circle cx="12" cy="12" r="10"/><path d="m9 12 2 2 4-4"/></svg>""",
        "arrow-right-left": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="m16 3 4 4-4 4"/><path d="M20 7H4"/><path d="m8 21-4-4 4-4"/><path d="M4 17h16"/></svg>""",
        "shuffle": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M18 4h3v3"/><path d="M3 7h3c3.5 0 5 10 8.5 10H21"/><path d="m18 20 3-3-3-3"/><path d="M3 17h3c1.2 0 2.2-.8 3.1-2"/></svg>""",
        "alert": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M12 9v4"/><path d="M12 17h.01"/><path d="M10.3 3.9 2.6 18a2 2 0 0 0 1.7 3h15.4a2 2 0 0 0 1.7-3L13.7 3.9a2 2 0 0 0-3.4 0z"/></svg>""",
        "siren": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M7 18v-6a5 5 0 1 1 10 0v6"/><path d="M5 21a1 1 0 0 0 1 1h12a1 1 0 0 0 1-1v-1a2 2 0 0 0-2-2H7a2 2 0 0 0-2 2z"/><path d="M21 12h1"/><path d="M18.5 4.5 18 5"/><path d="M2 12h1"/><path d="M12 2v1"/><path d="m4.929 4.929.707.707"/><path d="M12 12v6"/></svg>""",
        "life-buoy": """<svg viewBox="0 0 24 24" aria-hidden="true"><circle cx="12" cy="12" r="10"/><path d="m4.93 4.93 4.24 4.24"/><path d="m14.83 9.17 4.24-4.24"/><path d="m14.83 14.83 4.24 4.24"/><path d="m9.17 14.83-4.24 4.24"/><circle cx="12" cy="12" r="4"/></svg>""",
        "scale": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M12 3v18"/><path d="m19 8 3 8a5 5 0 0 1-6 0zV7"/><path d="M3 7h1a17 17 0 0 0 8-2 17 17 0 0 0 8 2h1"/><path d="m5 8 3 8a5 5 0 0 1-6 0zV7"/><path d="M7 21h10"/></svg>""",
        "radar": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M19.07 4.93A10 10 0 0 0 6.99 3.34"/><path d="M4 6h.01"/><path d="M2.29 9.62A10 10 0 1 0 21.31 8.35"/><path d="M16.24 7.76A6 6 0 1 0 8.23 16.67"/><path d="M12 18h.01"/><path d="M17.99 11.66A6 6 0 0 1 15.77 16.67"/><circle cx="12" cy="12" r="2"/><path d="m13.41 10.59 5.66-5.66"/></svg>""",
        "sparkles": """<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M11.017 2.814a1 1 0 0 1 1.966 0l1.051 5.558a2 2 0 0 0 1.594 1.594l5.558 1.051a1 1 0 0 1 0 1.966l-5.558 1.051a2 2 0 0 0-1.594 1.594l-1.051 5.558a1 1 0 0 1-1.966 0l-1.051-5.558a2 2 0 0 0-1.594-1.594l-5.558-1.051a1 1 0 0 1 0-1.966l5.558-1.051a2 2 0 0 0 1.594-1.594z"/><path d="M20 2v4"/><path d="M22 4h-4"/><circle cx="4" cy="20" r="2"/></svg>""",
    }
    return icons.get(name, icons["sparkles"])

def icon_badge(name):
    return f'<span class="section-icon">{icon_svg(name)}</span>'

def build_planning_prompt(household, time_horizon, situations, situation_detail, energy, notes):
    situation = normalize_situations(situations, situation_detail)
    situation_text = " ".join([
        situation,
        notes or "",
        household.get("environment", ""),
        household.get("preferences", "")
    ]).lower()

    # Filter relevant activities from knowledge base
    relevant = []
    for act in KNOWLEDGE_BASE["activities"]:
        # Basic filtering
        tags = act["tags"]
        include = True

        # Energy filter
        if energy == "Exhausted" and act["energy"] in ["high", "medium"]:
            include = False
        if energy == "Tired" and act["energy"] == "high":
            include = False

        child_sick = includes_any(situation_text, ["sick child", "sick kid", "child is sick", "kid is sick"])
        parent_sick = includes_any(situation_text, ["sick parent", "parent is sick", "i am sick", "i'm sick"])
        both_sick = "both sick" in situation_text or (child_sick and parent_sick)
        avoid_messy = includes_any(situation_text, ["messy", "mess", "dirty", "hands dirty", "paint", "tactile"])

        # Sick-day filter
        if child_sick or parent_sick or both_sick or "ill" in situation_text:
            if "outdoor" in tags or "physical" in tags or "energy_outlet" in tags:
                include = False
            if "rest" in tags or "calming" in tags or "quiet" in tags:
                include = True  # Re-include quiet activities
            if parent_sick or both_sick:
                if act["energy"] in ["high", "medium"] or "high_cleanup" in tags or "medium_cleanup" in tags:
                    include = False
                if "independent" in tags or "screen" in tags or "rest" in tags:
                    include = True
            if "outdoor" in tags or "physical" in tags or "energy_outlet" in tags:
                include = False
            if "medium_cleanup" in tags or "high_cleanup" in tags:
                include = False

        # Outdoor situation filter
        if "no outdoor" in situation_text or "too hot" in situation_text or "rainy" in situation_text or "cold" in situation_text:
            if "outdoor" in tags:
                include = False

        if avoid_messy:
            if "medium_cleanup" in tags or "high_cleanup" in tags or "tactile" in tags or "sensory" in tags:
                include = False

        # Solo parent filter - prioritize independent activities
        if "solo" in situation_text:
            if "independent" in tags or "screen" in tags or "quiet" in tags:
                include = True

        # Time filter
        duration = act["duration"]
        # If time horizon is short, skip very long activities
        if time_horizon in ["15 minutes", "30 minutes", "1 hour"]:
            if "60" in duration and "-" not in duration:
                include = False

        if include:
            relevant.append(f"- {act['name']} (tags: {', '.join(tags)}, energy: {act['energy']}, duration: {act['duration']})")

    # Keep the prompt compact enough for fast CPU generation.
    activity_context = "\n".join(relevant[:ACTIVITY_CONTEXT_LIMIT])

    user_prompt = f"""Household Profile:
- Structure: {household.get('structure', 'Not specified')}
- Children: {household.get('children', 'Not specified')}
- Environment: {household.get('environment', 'Not specified')}
- Budget: {household.get('budget', 'Not specified')}
- Preferences: {household.get('preferences', 'None')}

Session Inputs:
- Time Horizon: {time_horizon}
- Selected Situations: {", ".join(situations or []) if situations else "Not specified"}
- Situation Details: {situation_detail or "None"}
- Parent Energy: {energy}
- Additional Notes: {notes}

Sickness Guidance:
- Sick child: prioritize comfort, rest, hydration, low stimulation, easy supervision, and short transitions.
- Sick parent: prioritize parent rest, independent or screen-supported activities, low/no setup, and low/no cleanup.
- Both sick: survival-mode plan only; no outdoor activities, no physical games, no crafts with cleanup, no parent-intensive activities.
- Do not describe sick-day plans as high effort or high cleanup unless the user specifically requests that tradeoff.
- The activity list below is unordered. Choose based on the user's details, not the list order.
- Keep screen time out of the recommended plan unless it is clearly the most compassionate option for sickness, exhaustion, or an explicit request. It is acceptable in the alternative or emergency fallback.

Available Activities (unordered candidate ingredients):
{activity_context}

Generate a JSON response with exactly this structure:
{{
  "situation_assessment": {{
    "primary_goal": "...",
    "primary_constraint": "...",
    "secondary_constraints": "..."
  }},
  "recommended_plan": {{
    "title": "...",
    "timeline": [
      {{"time": "0:00-0:20", "activity": "...", "rationale": "..."}},
      {{"time": "0:20-0:50", "activity": "...", "rationale": "..."}}
    ],
    "total_effort": "Low/Medium/High",
    "tradeoffs": {{
      "best_when": "...",
      "parent_effort": "Low/Medium/High",
      "interruption_risk": "Low/Medium/High",
      "cleanup_burden": "Low/Medium/High"
    }}
  }},
  "alternative_plan": {{
    "title": "...",
    "steps": [
      {{"label": "0:00-0:10", "action": "..."}},
      {{"label": "0:10-0:40", "action": "..."}},
      {{"label": "0:40-end", "action": "..."}}
    ],
    "key_difference": "...",
    "tradeoffs": {{
      "best_when": "...",
      "parent_effort": "Low/Medium/High",
      "interruption_risk": "Low/Medium/High",
      "cleanup_burden": "Low/Medium/High"
    }}
  }},
  "emergency_fallback": {{
    "title": "...",
    "trigger": "...",
    "steps": [
      {{"label": "Stop", "action": "..."}},
      {{"label": "Stabilize", "action": "..."}},
      {{"label": "Simplify", "action": "..."}}
    ],
    "tradeoffs": {{
      "best_when": "...",
      "parent_effort": "Low/Medium/High",
      "interruption_risk": "Low/Medium/High",
      "cleanup_burden": "Low/Medium/High"
    }}
  }}
}}"""

    return user_prompt

def parse_json_output(text):
    """Extract JSON from model output."""
    # Try to find JSON between code blocks or braces
    text = text.strip()
    # Remove markdown code blocks
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:].strip()

    # Find first { and last }
    try:
        start = text.index("{")
        end = text.rindex("}") + 1
        json_str = text[start:end]
        return json.loads(json_str)
    except (ValueError, json.JSONDecodeError):
        return None

def horizon_minutes(time_horizon):
    horizons = {
        "15 minutes": 15,
        "30 minutes": 30,
        "1 hour": 60,
        "2 hours": 120,
        "Afternoon": 180,
        "Today": 360
    }
    return horizons.get(time_horizon, 60)

def fmt_minutes(minutes):
    hours, mins = divmod(minutes, 60)
    if hours:
        return f"{hours}:{mins:02d}"
    return f"0:{mins:02d}"

def build_timeline(time_horizon, names):
    total = horizon_minutes(time_horizon)
    setup = 5 if total <= 30 else 10
    remaining = max(total - setup, 10)
    first = max(5, remaining // 3)
    second = max(5, remaining // 3)

    boundaries = [
        (0, setup),
        (setup, setup + first),
        (setup + first, setup + first + second),
        (setup + first + second, total),
    ]
    activities = [
        ("Set expectations, water/snack, and one clear boundary", "A short setup reduces repeat questions and gives the plan a calmer start."),
        (names[0], "Start with a quiet, low-effort activity that does not require much parent coaching."),
        (names[1], "Rotate before boredom turns into interruptions."),
        (names[2], "End with the calmest option so the parent is not managing a messy transition."),
    ]

    timeline = []
    for (start, end), (activity, rationale) in zip(boundaries, activities):
        if end <= start:
            continue
        timeline.append({
            "time": f"{fmt_minutes(start)}-{fmt_minutes(end)}",
            "activity": activity,
            "rationale": rationale
        })
    return timeline

def build_local_plan(household_json, time_horizon, situations, situation_detail, energy, notes):
    household = json.loads(household_json) if household_json else {}
    situation_text = " ".join([
        normalize_situations(situations, situation_detail),
        notes or "",
        household.get("environment", ""),
        household.get("preferences", "")
    ]).lower()
    children = household.get("children", "your child")

    child_sick = includes_any(situation_text, ["sick child", "sick kid", "child is sick", "kid is sick"])
    parent_sick = includes_any(situation_text, ["sick parent", "parent is sick", "i am sick", "i'm sick"])
    both_sick = "both sick" in situation_text or (child_sick and parent_sick)
    no_outdoor = includes_any(situation_text, ["no outdoor", "apartment", "too hot", "rainy", "cold"])
    avoid_messy = includes_any(situation_text, ["messy", "mess", "dirty", "hands dirty", "paint", "tactile"])
    work_focus = includes_any(situation_text, ["work", "meeting", "focus", "call"])
    solo_parent = "solo" in situation_text
    preference_terms = {
        "drawing": ["drawing", "draw", "coloring", "color", "art"],
        "puzzles": ["puzzle", "puzzles", "logic", "matching", "maze"],
        "books": ["book", "read", "story"],
        "animals": ["animal", "doll", "figurine", "vet"],
        "quiet": ["quiet", "calm", "low stimulation", "low-stimulation"],
        "music": ["music", "audio", "podcast", "lullabies"],
    }

    def activity_allowed(act):
        tags = act["tags"]
        if no_outdoor and "outdoor" in tags:
            return False
        if avoid_messy and (
            "medium_cleanup" in tags
            or "high_cleanup" in tags
            or "tactile" in tags
            or "sensory" in tags
        ):
            return False
        if energy == "Exhausted" and (act["energy"] != "low" or "independent" not in tags):
            return False
        if energy == "Tired" and act["energy"] == "high":
            return False
        if child_sick or parent_sick or both_sick:
            if "outdoor" in tags or "physical" in tags or "energy_outlet" in tags:
                return False
            if "medium_cleanup" in tags or "high_cleanup" in tags:
                return False
        return True

    def activity_score(act):
        tags = act["tags"]
        haystack = f"{act['name']} {' '.join(tags)}".lower()
        score = 0

        if "low_cleanup" in tags:
            score += 3
        if "quiet" in tags:
            score += 3
        if "independent" in tags:
            score += 2
        if "indoor" in tags:
            score += 1
        if "rest" in tags or "calming" in tags:
            score += 2
        if "screen" in tags:
            score -= 5
            if parent_sick or both_sick or energy == "Exhausted":
                score += 3
        if work_focus and "independent" in tags:
            score += 3
        if solo_parent and "independent" in tags:
            score += 3
        if child_sick and ("rest" in tags or "calming" in tags or "quiet" in tags):
            score += 4
        if avoid_messy and "low_cleanup" in tags:
            score += 2
        if "multi_child" in tags and any(term in children.lower() for term in ["2", "two", "3", "three", "kids", "children"]):
            score += 2

        for terms in preference_terms.values():
            if any(term in situation_text for term in terms) and any(term in haystack for term in terms):
                score += 4

        return score

    allowed_activities = sorted(
        [act for act in KNOWLEDGE_BASE["activities"] if activity_allowed(act)],
        key=activity_score,
        reverse=True
    )
    quiet_activities = [
        act for act in allowed_activities
        if "quiet" in act["tags"] and act["energy"] == "low"
    ]
    independent_activities = [
        act for act in allowed_activities
        if "independent" in act["tags"] and act["energy"] == "low"
    ]
    calming_activities = [
        act for act in allowed_activities
        if "calming" in act["tags"] or "rest" in act["tags"]
    ]
    non_screen_quiet = [
        act for act in quiet_activities
        if "screen" not in act["tags"]
    ]
    non_screen_independent = [
        act for act in independent_activities
        if "screen" not in act["tags"]
    ]
    non_screen_calming = [
        act for act in calming_activities
        if "screen" not in act["tags"]
    ]
    def pick_activities(*groups, limit=4):
        selected = []
        seen = set()
        for group in groups:
            for act in group:
                if act["name"] in seen or "screen" in act["tags"]:
                    continue
                selected.append(act)
                seen.add(act["name"])
                if len(selected) >= limit:
                    return selected
        return selected

    if both_sick:
        primary_constraint = "Both parent and child need rest, comfort, and the lowest possible effort plan."
        activities = pick_activities(non_screen_calming, non_screen_quiet)
        effort = "Low"
        score = "1"
        interruption_risk = "Low"
    elif parent_sick:
        primary_constraint = "The parent is sick, so the plan must minimize setup, movement, coaching, and cleanup."
        activities = pick_activities(non_screen_independent, non_screen_calming, non_screen_quiet)
        effort = "Low"
        score = "2"
        interruption_risk = "Medium"
    elif child_sick:
        primary_constraint = "The child is sick, so comfort, rest, hydration, and easy supervision matter most."
        activities = pick_activities(non_screen_calming, non_screen_quiet)
        effort = "Low"
        score = "2"
        interruption_risk = "Medium"
    elif energy == "Exhausted":
        primary_constraint = "Parent energy and recovery are the limiting factors."
        activities = pick_activities(non_screen_independent, non_screen_calming, non_screen_quiet)
        effort = "Low"
        score = "2"
        interruption_risk = "Medium"
    elif "work" in situation_text or "meeting" in situation_text or "focus" in situation_text:
        primary_constraint = "The parent needs protected focus time with minimal interruptions."
        activities = pick_activities(non_screen_independent, non_screen_quiet)
        effort = "Low"
        score = "3"
        interruption_risk = "Medium"
    elif "solo" in situation_text:
        primary_constraint = "The plan needs to be sustainable without another adult taking over."
        activities = pick_activities(non_screen_independent, non_screen_quiet, non_screen_calming)
        effort = "Medium"
        score = "4"
        interruption_risk = "Medium"
    else:
        primary_constraint = "The family needs a realistic structure for the time block."
        activities = pick_activities(non_screen_quiet, non_screen_independent, non_screen_calming)
        effort = "Medium"
        score = "4"
        interruption_risk = "Medium"

    names = [act["name"] for act in activities[:4]]
    while len(names) < 4:
        names.append("Rest or quiet time on the couch")
    timeline = build_timeline(time_horizon, names)

    return {
        "situation_assessment": {
            "primary_goal": f"Get through the next {time_horizon} with a plan that fits {children}.",
            "primary_constraint": primary_constraint,
            "secondary_constraints": "Keep setup simple, preserve parent energy, and avoid cleanup surprises."
        },
        "recommended_plan": {
            "title": "Low-friction family reset",
            "timeline": timeline,
            "total_effort": effort,
            "tradeoffs": {
                "best_when": "Choose this when the parent has enough bandwidth for a tiny setup and wants the most balanced option.",
                "parent_effort": effort,
                "interruption_risk": interruption_risk,
                "cleanup_burden": "Low"
            }
        },
        "alternative_plan": {
            "title": "Even lower effort",
            "steps": [
                {"label": "0:00-0:05", "action": "Tell the child the plan is switching to the easiest version and set up water or a snack within reach."},
                {"label": "0:05-0:40", "action": f"Use {names[3]} as the main activity, with screen time allowed as a deliberate support tool if the parent needs uninterrupted recovery or work time."},
                {"label": "0:40-end", "action": "Keep the child in the same room or an easy-to-check spot, then close with a calm reset instead of a new activity."}
            ],
            "key_difference": "Trades enrichment for reliability and parent bandwidth.",
            "tradeoffs": {
                "best_when": "Choose this when parent energy is the real bottleneck or interruptions would be costly.",
                "parent_effort": "Low",
                "interruption_risk": "Low",
                "cleanup_burden": "Low"
            }
        },
        "emergency_fallback": {
            "title": "Stabilize first",
            "trigger": "Use this if there is crying, repeated interruption, escalating conflict, or the parent hits a wall.",
            "steps": [
                {"label": "Stop", "action": "Pause the current plan and lower the goal to safety, comfort, and one calm next step."},
                {"label": "Stabilize", "action": "Meet the immediate need first: bathroom, water, snack, medicine if already part of care, comfort item, or a quiet place to rest."},
                {"label": "Simplify", "action": "Switch to the safest low-effort option in the same room as the parent, even if that means screen time or rest instead of the original plan."}
            ],
            "tradeoffs": {
                "best_when": "Use this when the goal changes from having a good plan to getting everyone safe and regulated.",
                "parent_effort": "Lowest",
                "interruption_risk": "Lowest",
                "cleanup_burden": "Lowest"
            }
        },
        "metrics": {
            "parent_effort_score": score,
            "interruption_risk": interruption_risk,
            "cleanup_burden": "Low",
            "why_it_matters": "These scores help you choose the plan that protects parent bandwidth, not just the plan that looks best on paper.",
            "parent_effort_note": "How much setup, coaching, movement, and decision-making the parent has to do.",
            "interruption_note": "How likely the child is to need help or switch activities before the block ends.",
            "cleanup_note": "How much mess or reset work this creates after the plan."
        }
    }

def generate_plan(household_json, time_horizon, situations, situation_detail, energy, notes, progress=gr.Progress(track_tqdm=True)):
    try:
        progress(0.05, desc="Loading the planning model...")
        tokenizer, model = get_model()

        household = json.loads(household_json) if household_json else {}
        user_prompt = build_planning_prompt(household, time_horizon, situations, situation_detail, energy, notes)

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt}
        ]

        # Ensure pad_token is set for generation
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else "<pad>"
        if tokenizer.eos_token_id is None:
            tokenizer.eos_token_id = tokenizer.pad_token_id

        input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

        progress(0.35, desc="Reasoning through the family constraints...")
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id else tokenizer.pad_token_id
            )

        response_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        result = parse_json_output(response_text)

        if result is None:
            progress(0.85, desc="Polishing a grounded backup plan...")
            result = build_local_plan(household_json, time_horizon, situations, situation_detail, energy, notes)

        progress(1.0, desc="Plan ready.")
        return format_output(result)
    except Exception as exc:
        print(f"Model generation failed, using local fallback: {exc}")
        return format_output(build_local_plan(household_json, time_horizon, situations, situation_detail, energy, notes))

def normalize_steps(steps, fallback_steps):
    """Return display-ready action steps even when model output is sparse."""
    cleaned = []
    if isinstance(steps, list):
        for index, step in enumerate(steps[:4], start=1):
            if isinstance(step, dict):
                label = step.get("label") or step.get("time") or f"Step {index}"
                action = step.get("action") or step.get("activity") or step.get("description") or ""
            else:
                label = f"Step {index}"
                action = str(step)

            if str(action).strip():
                cleaned.append({"label": str(label), "action": str(action)})

    return cleaned or fallback_steps

def render_step_list(steps):
    return "\n".join(
        f"""
        <div class="step-item">
            <div class="step-label">{escape(step['label'])}</div>
            <div class="step-action">{escape(step['action'])}</div>
        </div>
        """
        for step in steps
    )

def normalize_tradeoffs(plan, defaults):
    tradeoffs = plan.get("tradeoffs", {})
    if not isinstance(tradeoffs, dict):
        tradeoffs = {}

    return {
        "best_when": str(tradeoffs.get("best_when") or defaults["best_when"]),
        "parent_effort": str(tradeoffs.get("parent_effort") or defaults["parent_effort"]),
        "interruption_risk": str(tradeoffs.get("interruption_risk") or defaults["interruption_risk"]),
        "cleanup_burden": str(tradeoffs.get("cleanup_burden") or defaults["cleanup_burden"]),
    }

def render_tradeoff_chips(tradeoffs):
    return f"""
    <div class="tradeoff-row">
        <div class="tradeoff-chip">
            <span class="tradeoff-label">Parent Effort</span>
            <span class="tradeoff-value">{escape(tradeoffs['parent_effort'])}</span>
        </div>
        <div class="tradeoff-chip">
            <span class="tradeoff-label">Interruption Risk</span>
            <span class="tradeoff-value">{escape(tradeoffs['interruption_risk'])}</span>
        </div>
        <div class="tradeoff-chip">
            <span class="tradeoff-label">Cleanup</span>
            <span class="tradeoff-value">{escape(tradeoffs['cleanup_burden'])}</span>
        </div>
    </div>
    """

def format_output(result):
    """Format the JSON result into custom HTML cards for display."""
    sa = result.get("situation_assessment", {})
    rp = result.get("recommended_plan", {})
    ap = result.get("alternative_plan", {})
    ef = result.get("emergency_fallback", {})
    m = result.get("metrics", {})

    timeline_items = []
    for item in rp.get("timeline", []):
        timeline_items.append(f"""
            <div class="timeline-row">
                <div class="time-chip">{escape(str(item.get('time', '')))}</div>
                <div class="timeline-item">
                    <div class="timeline-title">{escape(str(item.get('activity', '')))}</div>
                    <div class="timeline-reason">Rationale: {escape(str(item.get('rationale', '')))}</div>
                </div>
            </div>
        """)

    timeline_html = "\n".join(timeline_items) or """
        <div class="timeline-row">
            <div class="time-chip">Now</div>
            <div class="timeline-item">
                <div class="timeline-title">Start with the easiest stabilizing step</div>
                <div class="timeline-reason">Rationale: Reduce decision load before adding activities.</div>
            </div>
        </div>
    """

    alternative_steps = normalize_steps(
        ap.get("steps"),
        [
            {"label": "Switch", "action": ap.get("description", "Use the lowest-effort quiet option available and reduce the goal to calm, safe, and supervised.")},
            {"label": "Protect", "action": "Choose the version that costs less parent attention, even if it is less enriching."},
            {"label": "Reset", "action": "End with water, snack, bathroom, or rest before introducing anything new."}
        ]
    )
    fallback_steps = normalize_steps(
        ef.get("steps"),
        [
            {"label": "Stop", "action": ef.get("trigger", "Use this when the main plan starts failing or parent capacity drops suddenly.")},
            {"label": "Stabilize", "action": ef.get("description", "Meet the immediate need and lower the bar to safety, comfort, and one calm next step.")},
            {"label": "Simplify", "action": "Move to a safe quiet activity in the same room as the parent until things feel manageable again."}
        ]
    )

    alternative_steps_html = render_step_list(alternative_steps)
    fallback_steps_html = render_step_list(fallback_steps)
    recommended_tradeoffs = normalize_tradeoffs(
        rp,
        {
            "best_when": "Choose this when you want the most balanced plan for the time block.",
            "parent_effort": str(rp.get("total_effort", "Medium")),
            "interruption_risk": str(m.get("interruption_risk", "Medium")),
            "cleanup_burden": str(m.get("cleanup_burden", "Low")),
        }
    )
    alternative_tradeoffs = normalize_tradeoffs(
        ap,
        {
            "best_when": "Choose this when protecting parent bandwidth matters more than enrichment.",
            "parent_effort": "Low",
            "interruption_risk": "Low/Medium",
            "cleanup_burden": "Low",
        }
    )
    fallback_tradeoffs = normalize_tradeoffs(
        ef,
        {
            "best_when": "Use this when the plan is breaking down and the goal is immediate stabilization.",
            "parent_effort": "Lowest",
            "interruption_risk": "Lowest",
            "cleanup_burden": "Lowest",
        }
    )

    return f"""
    <div class="custom-output">
        <div class="section-hero">
            <div>
                <h2>We've got you.</h2>
                <div class="plan-copy">A realistic plan for the next parenting window, built around energy, constraints, and backup options.</div>
            </div>
            <div class="toy-scene" aria-hidden="true">
                <div class="bear-big"></div>
                <div class="bear-small"></div>
                <div class="block-a"></div>
                <div class="block-b"></div>
                <div class="block-c"></div>
            </div>
        </div>

        <section class="plan-card assessment">
            <div class="card-heading">
                {icon_badge("clipboard-list")}
                <h3>Situation Assessment</h3>
            </div>
            <div class="assessment-grid">
                <div class="fact">
                    <span class="fact-label">Primary Goal</span>
                    <div class="fact-value">{escape(str(sa.get('primary_goal', 'N/A')))}</div>
                </div>
                <div class="fact">
                    <span class="fact-label">Primary Constraint</span>
                    <div class="fact-value">{escape(str(sa.get('primary_constraint', 'N/A')))}</div>
                </div>
                <div class="fact">
                    <span class="fact-label">Secondary Constraints</span>
                    <div class="fact-value">{escape(str(sa.get('secondary_constraints', 'N/A')))}</div>
                </div>
            </div>
        </section>

        <section class="plan-card recommended">
            <div class="card-heading">
                {icon_badge("check")}
                <h3>Recommended Plan</h3>
            </div>
            <div class="plan-title">{escape(str(rp.get('title', 'Plan')))} (Effort: {escape(str(rp.get('total_effort', 'N/A')))})</div>
            <div class="timeline">{timeline_html}</div>
            {render_tradeoff_chips(recommended_tradeoffs)}
            <div class="plan-note"><strong>Best when:</strong> {escape(recommended_tradeoffs['best_when'])}</div>
        </section>

        <section class="plan-card alternative">
            <div class="card-heading">
                {icon_badge("arrow-right-left")}
                <h3>Alternative Plan</h3>
            </div>
            <div class="plan-title">{escape(str(ap.get('title', 'Lower-effort backup path')))}</div>
            <div class="step-list">{alternative_steps_html}</div>
            <div class="plan-note">
                <strong>Key difference:</strong> {escape(str(ap.get('key_difference', 'This option gives up some enrichment to protect parent bandwidth.')))}
            </div>
            {render_tradeoff_chips(alternative_tradeoffs)}
            <div class="plan-note"><strong>Best when:</strong> {escape(alternative_tradeoffs['best_when'])}</div>
        </section>

        <section class="plan-card fallback">
            <div class="card-heading">
                {icon_badge("siren")}
                <h3>Emergency Fallback</h3>
            </div>
            <div class="plan-title">{escape(str(ef.get('title', 'Immediate stabilization plan')))}</div>
            <div class="plan-note"><strong>Use when:</strong> {escape(str(ef.get('trigger', 'The primary plan is no longer working.')))}</div>
            <div class="step-list">{fallback_steps_html}</div>
            {render_tradeoff_chips(fallback_tradeoffs)}
            <div class="plan-note"><strong>Best when:</strong> {escape(fallback_tradeoffs['best_when'])}</div>
        </section>

        <div class="tip-strip"><strong>Tip:</strong> Keep one quiet backup option ready before the first transition gets wobbly.</div>
    </div>
    """

# Gradio UI
with gr.Blocks(
    title="Parent Co-Pilot"
) as demo:

    gr.HTML(f"""
    <div class="hero">
        <div class="brand-lockup">
            <div class="brand-mark">{icon_svg("sparkles")}</div>
            <div>
                <h1 class="brand-title">Parent Co-Pilot</h1>
                <div class="brand-copy">AI-powered parenting support for the next messy stretch.</div>
            </div>
        </div>
        <div class="idea-pill">
            <strong>Need a plan?</strong>
            Tell us the real-life constraints. We'll shape the next hour, day, or week.
        </div>
    </div>
    """)

    with gr.Group(elem_classes="form-panel family-panel"):
        gr.HTML(f"""
        <div class="panel-heading">
            {icon_badge("users")}
            <div>
                <div class="panel-title">Your Family & Preferences</div>
                <p class="panel-subtitle">Tell us what the plan needs to respect.</p>
            </div>
        </div>
        """)
        structure = gr.Dropdown(
            choices=["Single parent", "Two-parent household", "Shared custody", "Multigenerational", "Prefer not to say"],
            label="Household Structure",
            value="Two-parent household"
        )
        children = gr.Textbox(
            label="Children",
            placeholder="e.g., 2 children: ages 3 and 7",
            value="1 child: age 5"
        )
        environment = gr.Dropdown(
            choices=["Apartment", "House", "House with yard", "No outdoor space"],
            label="Environment",
            value="Apartment"
        )
        budget = gr.Dropdown(
            choices=["Free", "Low cost", "Flexible"],
            label="Budget Preference",
            value="Low cost"
        )
        preferences = gr.Textbox(
            label="Activity Preferences",
            placeholder="e.g., Loves animals, avoids messy activities",
            value="",
            lines=3,
            elem_id="preferences-input"
        )

    household_state = gr.State()

    def save_profile(structure, children, environment, budget, preferences):
        return json.dumps({
            "structure": structure,
            "children": children,
            "environment": environment,
            "budget": budget,
            "preferences": preferences
        })

    for comp in [structure, children, environment, budget, preferences]:
        comp.change(save_profile, [structure, children, environment, budget, preferences], household_state)

    # Initialize state
    demo.load(save_profile, [structure, children, environment, budget, preferences], household_state)

    with gr.Group(elem_classes="form-panel session-panel"):
        gr.HTML(f"""
        <div class="panel-heading">
            {icon_badge("calendar-clock")}
            <div>
                <div class="panel-title">Session Planner</div>
                <p class="panel-subtitle">Let's plan the actual window in front of you.</p>
            </div>
        </div>
        """)
        time_horizon = gr.Dropdown(
            choices=["15 minutes", "30 minutes", "1 hour", "2 hours", "Afternoon", "Today"],
            label="How long do you need help with?",
            value="1 hour"
        )
        situations = gr.CheckboxGroup(
            choices=[
                "Work day",
                "Summer break",
                "School holiday",
                "Sick child",
                "Sick parent",
                "Both sick",
                "Too hot outside",
                "Rainy day",
                "Solo parenting",
                "Schedule disruption"
            ],
            label="Current Situation",
            value=["Work day"]
        )
        situation_detail = gr.Textbox(
            label="Situation Details",
            placeholder="e.g., Need to focus for a meeting, child has a mild fever, spouse is traveling...",
            lines=2,
            value="Need to focus"
        )
        energy = gr.Dropdown(
            choices=["Fine", "Tired", "Exhausted"],
            label="Parent Energy",
            value="Tired"
        )
        notes = gr.Textbox(
            label="Optional Notes",
            placeholder="e.g., Need quiet activities, child loves animals, low cleanup...",
            lines=2,
            value="Need quiet activities"
        )

        generate_btn = gr.Button("Generate Plan", variant="primary", elem_id="generate-btn")

    output = gr.HTML(label="Plan", elem_id="plan-output")

    generate_btn.click(
        generate_plan,
        [household_state, time_horizon, situations, situation_detail, energy, notes],
        output,
        show_progress="full",
        show_progress_on=[output]
    )

    gr.HTML(f"""
    <div class="footer-note">
        Built by Clarissa Chen for the Build Small Hackathon. Current tiny model: {escape(MODEL_ID)}.
    </div>
    """)

if __name__ == "__main__":
    if PRELOAD_MODEL:
        get_model()
    demo.launch(
        css=APP_CSS,
        theme=gr.themes.Soft(primary_hue="violet", secondary_hue="slate", neutral_hue="slate")
    )