app.py

"""
Planning Summary Audio Analyzer - Hugging Face Spaces App
Analyzes audio recordings of planning conversations and generates 
a structured Word document planning summary report using Google's Gemini API.

CHANGELOG (corrections applied):
  1. Treatment preference: added "conditional_comfort_care" option for nuanced cases
  2. Beneficiary status: improved prompt guidance to distinguish account access from
     formal beneficiary designation
  3. Values vs. care preferences: clarified prompt so medical decision criteria are not
     conflated with life meaning/joy
  4. Personal items: added "no_specific_items" option (deliberate choice vs. indecision)
  5. Name spelling: prompt now flags uncertain proper noun spellings with [verify spelling]
  6. Next Steps section: driven by extracted data and topics discussed, not hardcoded
  7. Prompt includes "topics_discussed" field so the report only covers relevant sections
"""

import os
import re
import json
import time
import tempfile
from docx import Document
from docx.shared import Inches, Pt, Twips
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.enum.style import WD_STYLE_TYPE
from docx.oxml.ns import qn
from docx.oxml import OxmlElement

# Defer heavy/optional imports so core logic is testable without them
try:
    import gradio as gr
    HAS_GRADIO = True
except ImportError:
    HAS_GRADIO = False

try:
    import google.generativeai as genai
    HAS_GENAI = True
except ImportError:
    HAS_GENAI = False

# ============================================================================
# EXTRACTION PROMPT
# ============================================================================

EXTRACTION_PROMPT = """
You are analyzing a recorded conversation about advance care planning and end-of-life wishes.
Listen to the ENTIRE audio carefully and extract ALL relevant information.

CRITICAL INSTRUCTIONS FOR SINGLE-SELECT FIELDS:
- You MUST select exactly ONE option for each single-select field
- Use the EXACT string values specified (copy them exactly)
- If the conversation implies something even indirectly, make your best inference
- NEVER leave single-select fields as null - always pick the best match

IMPORTANT RULES FOR PROPER NOUNS:
- If a last name is spelled out letter by letter, use that exact spelling.
- If a last name is only spoken (not spelled), transcribe it phonetically and append
  [verify spelling] after it. Example: "Potoff [verify spelling]"
- First names that are spelled out should use the spelled version.

Return a JSON object with this EXACT structure:
```json
{
  "participant": {
    "name": "First and last name if mentioned",
    "conversation_date": "MM/DD/YYYY format if mentioned, or null",
    "facilitator": "Facilitator name and credentials, or 'Not discussed'",
    "location": "Location of conversation, or 'Not discussed'"
  },
  "topics_discussed": ["health", "financial", "funeral"],
  "health_care_wishes": {
    "primary_decision_maker": {
      "name": "Name",
      "relationship": "spouse/wife/husband/son/daughter/etc",
      "phone": "phone or null",
      "email": "email or null"
    },
    "backup_decision_maker": {
      "name": "Name or null",
      "relationship": "relationship or null",
      "phone": "phone or null"
    },
    "values_summary": "2-3 sentence summary of their values and care priorities",
    "advance_care_status": "MUST BE ONE OF: has_current_documents | has_documents_needs_update | no_documents",
    "treatment_preference": "MUST BE ONE OF: comfort_care_only | full_treatment_if_recovery | conditional_comfort_care | unsure",
    "treatment_details": "Details about specific conditions, CPR, ventilation, ICU preferences",
    "additional_notes": "Other notes or null"
  },
  "financial_planning": {
    "financial_primary": {
      "name": "Name",
      "relationship": "relationship",
      "phone": "phone or null",
      "email": "email or null"
    },
    "financial_backup": {
      "name": "Name or null",
      "relationship": "relationship or null",
      "phone": "phone or null"
    },
    "financial_conversation_summary": "Summary of who handles finances and how",
    "documents_in_place": {
      "financial_poa": false,
      "will": false,
      "living_trust": false,
      "tod_designations": false,
      "joint_ownership": false,
      "none": true
    },
    "legal_details": "Details about their legal/financial situation",
    "next_steps": {
      "review_update_poa_will": false,
      "identify_alternate": false,
      "contact_attorney": true,
      "seek_trust_advice": false,
      "other": "Other specific steps mentioned or null"
    },
    "beneficiary_status": "MUST BE ONE OF: all_current | need_to_update | unsure",
    "account_notes": "Notes about specific accounts, 401k, pension, retirement",
    "beneficiary_conversation": "Summary of beneficiary discussion",
    "has_info_list": "MUST BE ONE OF: yes_shared | yes_not_shared | not_created",
    "info_location": "Where files/info are stored",
    "organization_ideas": "Ideas about organizing documents",
    "shared_with_loved_ones": "MUST BE ONE OF: yes_written | yes_not_written | not_yet",
    "sharing_notes": "Notes about family discussions",
    "overall_wishes": "Summary of financial wishes",
    "specific_items_status": "MUST BE ONE OF: has_specific_items | no_specific_items | not_yet_decided",
    "specific_items": [
      {"item": "Description of item", "recipient": "Intended recipient"}
    ]
  },
  "funeral_plans": {
    "service_type": "MUST BE ONE OF: funeral | memorial | celebration_of_life | other | not_discussed",
    "service_type_other": "If other, describe",
    "body_preference": "MUST BE ONE OF: burial | cremation | donation | undecided | not_discussed",
    "body_details": "Location details like 'Ashes spread at favorite fishing lake'",
    "conversation_summary": "Summary of funeral wishes discussion",
    "preferred_location": "Where service should be held or null",
    "service_leader": "Who should lead or null",
    "music_readings": "Music and reading preferences or null",
    "appearance_clothing": "Dress code preferences or null",
    "charity_donations": "Charity for donations or null",
    "cost_planning": "MUST BE ONE OF: prepaid | family_aware | needs_discussion | not_discussed",
    "additional_notes": "Notes about life insurance, costs, etc."
  },
  "values_reflections": {
    "what_matters_most": "How they want to live and be remembered, based on legacy statements",
    "meaning_and_joy": "Hobbies, relationships, activities, and sources of happiness mentioned OUTSIDE of medical decision-making context",
    "want_remembered_for": "What they explicitly said they hope people remember about them"
  },
  "recommended_next_steps": {
    "create_healthcare_poa": false,
    "provide_poa_to_healthcare_team": false,
    "complete_financial_poa_will_trust": false,
    "review_update_beneficiaries": false,
    "create_financial_info_list": false,
    "discuss_wishes_with_loved_ones": false,
    "store_documents_safely": false,
    "review_plans_annually": false,
    "explore_funeral_preplanning": false,
    "other_steps": ["any other specific steps identified in conversation"]
  },
  "facilitator_summary": "Facilitator's closing summary and recommendations"
}
```

DECISION GUIDE FOR COMMON SCENARIOS:

topics_discussed:
- Listen for which topics the participant chose to focus on
- Only include "health", "financial", and/or "funeral" if they were actually discussed
- If the participant said they only want to discuss health and financial, do NOT include "funeral"

advance_care_status:
- If they say they haven't done paperwork/documents yet -> "no_documents"
- If they have old documents that need updating -> "has_documents_needs_update"
- If they have current, up-to-date documents -> "has_current_documents"

treatment_preference:
- If they UNCONDITIONALLY want only comfort/palliative care, no machines ever -> "comfort_care_only"
- If they want full treatment/CPR/ventilation IF there's hope of meaningful recovery -> "full_treatment_if_recovery"
- If their preference DEPENDS ON CONDITIONS such as cognitive function, prognosis,
  or quality of life (e.g. "treat me if I can still think clearly, but let me go
  if I'm cognitively impaired") -> "conditional_comfort_care"
- If they're unsure or need more information -> "unsure"

beneficiary_status:
- IMPORTANT: "all_current" means the participant explicitly confirmed that formal
  beneficiary designations (not just account access) are filed and up to date
- If they say their trusted person "has access" or "knows about" the accounts but
  did NOT explicitly confirm legal beneficiary designations are current -> "unsure"
- If they know some designations need updating -> "need_to_update"
- If they're not sure who's listed or need to check -> "unsure"

has_info_list:
- If they have files/info but haven't shared the location or it's disorganized -> "yes_not_shared"
- If their trusted person knows where everything is -> "yes_shared"
- If they haven't created any list -> "not_created"

shared_with_loved_ones:
- If they've talked but nothing is written down -> "yes_not_written"
- If they've discussed AND written it down -> "yes_written"
- If they haven't discussed wishes yet -> "not_yet"

specific_items_status:
- If they named specific items for specific people -> "has_specific_items"
- If they explicitly said no specific designations are needed (e.g. "everything goes
  to my spouse" or "nothing specific needs to go anywhere specific") -> "no_specific_items"
- If they haven't thought about it yet or are undecided -> "not_yet_decided"

service_type / body_preference / cost_planning:
- If funeral planning was NOT discussed at all, use "not_discussed" for these fields
- Celebration of life, casual gathering, party -> "celebration_of_life"
- Traditional funeral -> "funeral"
- Memorial service -> "memorial"
- Cremation, ashes spread somewhere -> "cremation"
- Burial in cemetery/ground -> "burial"
- Donate body to science -> "donation"
- If they mention life insurance will cover it or family knows about funding -> "family_aware"
- If they have a pre-paid funeral plan -> "prepaid"
- If costs haven't been discussed -> "needs_discussion"

values_reflections:
- "meaning_and_joy": ONLY include hobbies, relationships, passions, and activities
  that bring happiness. Do NOT include medical decision criteria like cognitive
  function preferences here. Those belong in treatment_details.
- "want_remembered_for": Use the participant's own words about how they want to be remembered.
- "what_matters_most": Summarize their overall philosophy about living and legacy.

recommended_next_steps:
- Set each to true ONLY if it is relevant based on what was discussed
- For example, if funeral planning was not discussed, do not set explore_funeral_preplanning to true
- If documents already exist and are current, do not set create_healthcare_poa to true
- Base these on gaps identified during the conversation

Listen for these key topics:
- Who would make healthcare decisions (usually spouse first, then adult child)
- Who would handle finances (often same people)
- Whether they have existing legal documents
- Their wishes about medical treatment and life support
- Funeral/memorial preferences
- Special items to give specific people
- What matters most to them, their values

Return ONLY valid JSON, no markdown formatting or explanation.
"""

# ============================================================================
# AUDIO ANALYSIS
# ============================================================================

def analyze_audio(audio_path: str, api_key: str) -> str:
    """Upload audio to Gemini and extract planning information."""
    if not HAS_GENAI:
        raise RuntimeError("google-generativeai package is required for audio analysis")
    genai.configure(api_key=api_key)

    audio_file = genai.upload_file(audio_path)

    # Wait for processing
    while audio_file.state.name == "PROCESSING":
        time.sleep(5)
        audio_file = genai.get_file(audio_file.name)

    if audio_file.state.name == "FAILED":
        raise ValueError(f"Audio processing failed: {audio_file.state.name}")

    model = genai.GenerativeModel('gemini-3-flash-preview')

    response = model.generate_content(
        [audio_file, EXTRACTION_PROMPT],
        generation_config=genai.GenerationConfig(
            temperature=0.1,
            max_output_tokens=8192
        )
    )

    # Cleanup uploaded file
    genai.delete_file(audio_file.name)

    return response.text

# ============================================================================
# JSON PARSING AND NORMALIZATION
# ============================================================================

def parse_json_response(response_text: str) -> dict | None:
    """Extract JSON from Gemini response."""
    cleaned = re.sub(r'```json\s*', '', response_text)
    cleaned = re.sub(r'```\s*', '', cleaned)

    json_match = re.search(r'\{[\s\S]*\}', cleaned)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            return None
    return None


def normalize_value(value, valid_options, default=None):
    """Normalize a value to match one of the valid options."""
    if value is None:
        return default

    val_str = str(value).lower().strip()
    val_normalized = val_str.replace(' ', '_').replace('-', '_')

    # Direct match
    for opt in valid_options:
        if val_normalized == opt.lower():
            return opt

    # Fuzzy matching rules
    matching_rules = {
        'no_documents': ['no_documents', 'none', 'not_completed', 'no documents', 'not yet completed'],
        'has_current_documents': ['has_current', 'current', 'up_to_date', 'have documents'],
        'has_documents_needs_update': ['needs_update', 'need_update', 'review', 'outdated'],
        'full_treatment_if_recovery': ['full_treatment', 'full treatment', 'aggressive', 'treatment if recovery'],
        'conditional_comfort_care': ['conditional', 'depends on', 'conditional_comfort', 'if cognitive', 'condition based'],
        'comfort_care_only': ['comfort_care_only', 'comfort care only', 'palliative only', 'no machines ever', 'only comfort'],
        'unsure': ['unsure', 'not sure', 'uncertain', 'undecided', 'need more info'],
        'need_to_update': ['need_to_update', 'needs update', 'update', 'outdated'],
        'all_current': ['all_current', 'all current', 'confirmed current', 'designations current'],
        'yes_not_shared': ['yes_not_shared', 'yes but', 'have but', 'not shared', 'disorganized'],
        'yes_shared': ['yes_shared', 'yes shared', 'knows where', 'shared'],
        'not_created': ['not_created', 'no list', 'none created', "haven't created", 'not yet created'],
        'yes_not_written': ['yes_not_written', 'discussed not written', 'talked but', 'verbal', 'not written'],
        'yes_written': ['yes_written', 'written', 'documented', 'written down'],
        'not_yet': ['not_yet', "haven't discussed", 'not discussed yet'],
        'has_specific_items': ['has_specific', 'yes_specific', 'has items', 'specific items'],
        'no_specific_items': ['no_specific', 'no specific', 'everything to', 'nothing specific', 'no designations'],
        'not_yet_decided': ['not_yet_decided', 'not decided', 'undecided', "haven't thought"],
        'celebration_of_life': ['celebration', 'celebration_of_life', 'party', 'gathering', 'casual'],
        'funeral': ['funeral', 'traditional'],
        'memorial': ['memorial', 'memorial_service'],
        'other': ['other'],
        'not_discussed': ['not_discussed', 'not discussed', 'skipped', 'not covered'],
        'cremation': ['cremation', 'cremate', 'ashes', 'cremated'],
        'burial': ['burial', 'bury', 'buried', 'cemetery', 'ground'],
        'donation': ['donation', 'donate', 'science', 'donate body'],
        'family_aware': ['family_aware', 'family aware', 'life insurance', 'insurance', 'covered'],
        'prepaid': ['prepaid', 'pre-paid', 'pre paid', 'paid'],
        'needs_discussion': ['needs_discussion', 'need to discuss'],
    }

    for opt in valid_options:
        if opt in matching_rules:
            for pattern in matching_rules[opt]:
                if pattern in val_str or pattern in val_normalized:
                    return opt

    return default


def normalize_data(data: dict) -> dict:
    """Normalize all single-select fields in the extracted data."""
    if not data:
        return data

    # Ensure topics_discussed exists
    if 'topics_discussed' not in data:
        data['topics_discussed'] = ['health', 'financial', 'funeral']

    # Health care wishes
    if 'health_care_wishes' in data:
        hcw = data['health_care_wishes']
        hcw['advance_care_status'] = normalize_value(
            hcw.get('advance_care_status'),
            ['has_current_documents', 'has_documents_needs_update', 'no_documents'],
            'no_documents'
        )
        hcw['treatment_preference'] = normalize_value(
            hcw.get('treatment_preference'),
            ['comfort_care_only', 'full_treatment_if_recovery', 'conditional_comfort_care', 'unsure'],
            'unsure'
        )

    # Financial planning
    if 'financial_planning' in data:
        fp = data['financial_planning']
        fp['beneficiary_status'] = normalize_value(
            fp.get('beneficiary_status'),
            ['all_current', 'need_to_update', 'unsure'],
            'unsure'
        )
        fp['has_info_list'] = normalize_value(
            fp.get('has_info_list'),
            ['yes_shared', 'yes_not_shared', 'not_created'],
            'yes_not_shared'
        )
        fp['shared_with_loved_ones'] = normalize_value(
            fp.get('shared_with_loved_ones'),
            ['yes_written', 'yes_not_written', 'not_yet'],
            'yes_not_written'
        )
        fp['specific_items_status'] = normalize_value(
            fp.get('specific_items_status'),
            ['has_specific_items', 'no_specific_items', 'not_yet_decided'],
            'not_yet_decided'
        )

    # Funeral plans
    if 'funeral_plans' in data:
        fun = data['funeral_plans']
        fun['service_type'] = normalize_value(
            fun.get('service_type'),
            ['funeral', 'memorial', 'celebration_of_life', 'other', 'not_discussed'],
            'not_discussed'
        )
        fun['body_preference'] = normalize_value(
            fun.get('body_preference'),
            ['burial', 'cremation', 'donation', 'undecided', 'not_discussed'],
            'not_discussed'
        )
        fun['cost_planning'] = normalize_value(
            fun.get('cost_planning'),
            ['prepaid', 'family_aware', 'needs_discussion', 'not_discussed'],
            'not_discussed'
        )

    # Ensure recommended_next_steps exists with sensible defaults
    if 'recommended_next_steps' not in data:
        data['recommended_next_steps'] = _infer_next_steps(data)

    return data


def _infer_next_steps(data: dict) -> dict:
    """Infer recommended next steps from the extracted data when the model
    does not return them explicitly."""
    topics = data.get('topics_discussed', [])
    health = data.get('health_care_wishes', {})
    financial = data.get('financial_planning', {})

    steps = {
        "create_healthcare_poa": False,
        "provide_poa_to_healthcare_team": False,
        "complete_financial_poa_will_trust": False,
        "review_update_beneficiaries": False,
        "create_financial_info_list": False,
        "discuss_wishes_with_loved_ones": False,
        "store_documents_safely": False,
        "review_plans_annually": False,
        "explore_funeral_preplanning": False,
        "other_steps": []
    }

    if 'health' in topics:
        status = health.get('advance_care_status', '')
        if status in ('no_documents', 'has_documents_needs_update'):
            steps['create_healthcare_poa'] = True
            steps['provide_poa_to_healthcare_team'] = True
        steps['store_documents_safely'] = True
        steps['review_plans_annually'] = True

    if 'financial' in topics:
        docs = financial.get('documents_in_place', {})
        if is_true(docs.get('none')) or not any(
            is_true(docs.get(k)) for k in
            ['financial_poa', 'will', 'living_trust', 'tod_designations', 'joint_ownership']
        ):
            steps['complete_financial_poa_will_trust'] = True

        ben = financial.get('beneficiary_status', '')
        if ben in ('unsure', 'need_to_update'):
            steps['review_update_beneficiaries'] = True

        info = financial.get('has_info_list', '')
        if info in ('not_created', 'yes_not_shared'):
            steps['create_financial_info_list'] = True

        shared = financial.get('shared_with_loved_ones', '')
        if shared in ('not_yet', 'yes_not_written'):
            steps['discuss_wishes_with_loved_ones'] = True

        steps['store_documents_safely'] = True
        steps['review_plans_annually'] = True

    if 'funeral' in topics:
        steps['explore_funeral_preplanning'] = True

    return steps

# ============================================================================
# WORD DOCUMENT GENERATION
# ============================================================================

def get_value(data, *keys, default="Not discussed"):
    """Safely get nested dictionary values."""
    result = data
    for key in keys:
        if isinstance(result, dict) and key in result:
            result = result[key]
        else:
            return default
    if result is None:
        return default
    return result if result else default


def cb(checked: bool) -> str:
    """Return Unicode checkbox."""
    return "\u2612" if checked else "\u2610"


def is_true(val) -> bool:
    """Check if a value is truthy."""
    if val is None:
        return False
    if isinstance(val, bool):
        return val
    if isinstance(val, str):
        return val.lower() in ('true', 'yes', '1')
    return bool(val)


def set_cell_shading(cell, color):
    """Set cell background color."""
    shading_elm = OxmlElement('w:shd')
    shading_elm.set(qn('w:fill'), color)
    cell._tc.get_or_add_tcPr().append(shading_elm)


def add_checkbox_paragraph(doc, checked, text, indent_level=0):
    """Add a paragraph with checkbox."""
    p = doc.add_paragraph()
    p.paragraph_format.left_indent = Inches(0.25 * indent_level)
    p.paragraph_format.space_before = Pt(2)
    p.paragraph_format.space_after = Pt(2)
    run = p.add_run(f"{cb(checked)} {text}")
    run.font.name = 'Arial'
    run.font.size = Pt(10)
    return p


def add_field_label(doc, text):
    """Add a field label paragraph."""
    p = doc.add_paragraph()
    p.paragraph_format.space_before = Pt(8)
    p.paragraph_format.space_after = Pt(2)
    run = p.add_run(text)
    run.font.name = 'Arial'
    run.font.size = Pt(10)
    run.bold = True
    return p


def add_field_value(doc, text):
    """Add a field value paragraph."""
    p = doc.add_paragraph()
    p.paragraph_format.space_after = Pt(4)
    run = p.add_run(text)
    run.font.name = 'Arial'
    run.font.size = Pt(10)
    return p


def add_section_header(doc, text):
    """Add a section header."""
    p = doc.add_paragraph()
    p.paragraph_format.space_before = Pt(16)
    p.paragraph_format.space_after = Pt(8)
    run = p.add_run(text)
    run.font.name = 'Arial'
    run.font.size = Pt(14)
    run.bold = True
    return p


def add_sub_header(doc, text):
    """Add a sub-header."""
    p = doc.add_paragraph()
    p.paragraph_format.space_before = Pt(12)
    p.paragraph_format.space_after = Pt(4)
    run = p.add_run(text)
    run.font.name = 'Arial'
    run.font.size = Pt(11)
    run.bold = True
    return p


def generate_docx(data: dict, output_path: str) -> str:
    """Generate the planning summary Word document."""
    doc = Document()

    # Set default font
    style = doc.styles['Normal']
    style.font.name = 'Arial'
    style.font.size = Pt(10)

    # Set page margins (0.75 inch)
    for section in doc.sections:
        section.top_margin = Inches(0.6)
        section.bottom_margin = Inches(0.6)
        section.left_margin = Inches(0.75)
        section.right_margin = Inches(0.75)

    topics = data.get('topics_discussed', ['health', 'financial', 'funeral'])

    # ===== TITLE =====
    title = doc.add_paragraph()
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    title_run = title.add_run("My Planning Summary")
    title_run.font.name = 'Arial'
    title_run.font.size = Pt(16)
    title_run.bold = True

    participant = data.get('participant', {})

    # Participant info
    p = doc.add_paragraph()
    p.add_run("Participant Name: ").bold = True
    p.add_run(get_value(participant, 'name'))

    p = doc.add_paragraph()
    p.add_run("Date of Conversation: ").bold = True
    p.add_run(get_value(participant, 'conversation_date'))

    p = doc.add_paragraph()
    p.add_run("Facilitator: ").bold = True
    p.add_run(get_value(participant, 'facilitator'))

    p = doc.add_paragraph()
    p.add_run("Location: ").bold = True
    p.add_run(get_value(participant, 'location'))

    # Intro text
    intro = doc.add_paragraph()
    intro.paragraph_format.space_before = Pt(8)
    intro.paragraph_format.space_after = Pt(12)
    intro_run = intro.add_run(
        "This summary captures the main points of our conversation about planning for the future. "
        "It is not a legal document, but it can help you share your wishes and guide next steps."
    )
    intro_run.font.size = Pt(9)
    intro_run.font.color.rgb = None  # Use default color

    # ===== HEALTH & CARE WISHES =====
    if 'health' in topics:
        add_section_header(doc, "My Health & Care Wishes")

        health = data.get('health_care_wishes', {})

        add_field_label(doc, "Who would you trust to make health decisions for you if you could not speak for yourself?")

        primary = health.get('primary_decision_maker', {})
        add_field_value(doc, f"Primary: {get_value(primary, 'name')} \u2013 {get_value(primary, 'relationship')}")
        add_field_value(doc, f"Phone: {get_value(primary, 'phone')} | Email: {get_value(primary, 'email')}")

        backup = health.get('backup_decision_maker', {})
        add_field_value(doc, f"Back-up: {get_value(backup, 'name')} \u2013 {get_value(backup, 'relationship')}")
        add_field_value(doc, f"Phone: {get_value(backup, 'phone')}")

        add_field_label(doc, "Summary of what you shared about your values and care priorities:")
        add_field_value(doc, get_value(health, 'values_summary'))

        add_field_label(doc, "Current advance-care planning status:")
        status = health.get('advance_care_status', '')
        add_checkbox_paragraph(doc, status == 'has_current_documents', "I have a health-care power of attorney or living will (up-to-date)")
        add_checkbox_paragraph(doc, status == 'has_documents_needs_update', "I have documents but need to review/update")
        add_checkbox_paragraph(doc, status == 'no_documents', "I have not yet completed these documents")

        add_field_label(doc, "Treatment preferences we discussed:")
        treatment = health.get('treatment_preference', '')
        add_checkbox_paragraph(doc, treatment == 'comfort_care_only', "Comfort care only (no machines or resuscitation under any circumstances)")
        add_checkbox_paragraph(doc, treatment == 'full_treatment_if_recovery', "Full medical treatment if recovery is possible")
        add_checkbox_paragraph(doc, treatment == 'conditional_comfort_care', "Conditional: treatment depends on prognosis or quality of life (see details below)")
        add_checkbox_paragraph(doc, treatment == 'unsure', "Unsure / would like more information")

        add_field_label(doc, "Additional details related to treatment preferences that were discussed:")
        add_field_value(doc, get_value(health, 'treatment_details'))

        add_field_label(doc, "Notes from conversation:")
        add_field_value(doc, get_value(health, 'additional_notes'))

    # ===== FINANCIAL PLANNING =====
    if 'financial' in topics:
        doc.add_page_break()
        add_section_header(doc, "My Financial Planning & Legacy")

        financial = data.get('financial_planning', {})

        add_sub_header(doc, "A. Trusted Person(s)")
        add_field_label(doc, "Who do you trust to manage your finances if you are unable?")

        fin_primary = financial.get('financial_primary', {})
        add_field_value(doc, f"Primary: {get_value(fin_primary, 'name')} \u2013 {get_value(fin_primary, 'relationship')}")
        add_field_value(doc, f"Phone: {get_value(fin_primary, 'phone')} | Email: {get_value(fin_primary, 'email')}")

        fin_backup = financial.get('financial_backup', {})
        add_field_value(doc, f"Back-up: {get_value(fin_backup, 'name')} \u2013 {get_value(fin_backup, 'relationship')}")
        add_field_value(doc, f"Phone: {get_value(fin_backup, 'phone')}")

        add_field_label(doc, "Conversation summary:")
        add_field_value(doc, get_value(financial, 'financial_conversation_summary'))

        add_sub_header(doc, "B. Legal and Financial Readiness")
        add_field_label(doc, "Documents currently in place (check all that apply):")

        docs = financial.get('documents_in_place', {})
        add_checkbox_paragraph(doc, is_true(docs.get('financial_poa')), "Financial Power of Attorney (POA)")
        add_checkbox_paragraph(doc, is_true(docs.get('will')), "Will")
        add_checkbox_paragraph(doc, is_true(docs.get('living_trust')), "Living Trust")
        add_checkbox_paragraph(doc, is_true(docs.get('tod_designations')), "Transfer on Death (TOD) designations")
        add_checkbox_paragraph(doc, is_true(docs.get('joint_ownership')), "Joint ownership of key accounts")
        add_checkbox_paragraph(doc, is_true(docs.get('none')), "None completed yet")

        add_field_label(doc, "Details from conversation:")
        add_field_value(doc, get_value(financial, 'legal_details'))

        add_field_label(doc, "Next steps (as discussed):")
        next_steps = financial.get('next_steps', {})
        add_checkbox_paragraph(doc, is_true(next_steps.get('review_update_poa_will')), "Review or update existing POA or Will")
        add_checkbox_paragraph(doc, is_true(next_steps.get('identify_alternate')), "Identify alternate decision-maker")
        add_checkbox_paragraph(doc, is_true(next_steps.get('contact_attorney')), "Contact attorney or legal aid for document preparation")
        add_checkbox_paragraph(doc, is_true(next_steps.get('seek_trust_advice')), "Seek advice on creating or updating a Trust")
        other_steps = next_steps.get('other')
        if other_steps and other_steps not in ('null', None, 'Not discussed', ''):
            add_checkbox_paragraph(doc, True, f"Other: {other_steps}")

        add_sub_header(doc, "C. Beneficiaries and Account Management")
        add_field_label(doc, "Status of major accounts and beneficiaries:")

        ben_status = financial.get('beneficiary_status', '')
        add_checkbox_paragraph(doc, ben_status == 'all_current', "All beneficiaries current and reflect my wishes")
        add_checkbox_paragraph(doc, ben_status == 'need_to_update', "Need to review or update some")
        add_checkbox_paragraph(doc, ben_status == 'unsure', "Unsure / need help locating information")

        add_field_label(doc, "Notes about specific accounts (bank, insurance, retirement):")
        add_field_value(doc, get_value(financial, 'account_notes'))

        add_field_label(doc, "Conversation summary:")
        add_field_value(doc, get_value(financial, 'beneficiary_conversation'))

        add_sub_header(doc, "D. Organizing Financial Information")
        add_field_label(doc, "Do you have a list of key information (accounts, passwords, insurance details)?")

        info_status = financial.get('has_info_list', '')
        add_checkbox_paragraph(doc, info_status == 'yes_shared', "Yes \u2013 my trusted person knows where it is")
        add_checkbox_paragraph(doc, info_status == 'yes_not_shared', "Yes \u2013 but not shared or outdated")
        add_checkbox_paragraph(doc, info_status == 'not_created', "Not yet created")

        add_field_label(doc, "Where this information can be found:")
        add_field_value(doc, get_value(financial, 'info_location'))

        add_field_label(doc, "Additional organization ideas shared in conversation:")
        add_field_value(doc, get_value(financial, 'organization_ideas'))

        add_sub_header(doc, "E. Talking with Loved Ones")
        add_field_label(doc, "Have you shared your financial wishes with loved ones or trusted helpers?")

        shared = financial.get('shared_with_loved_ones', '')
        add_checkbox_paragraph(doc, shared == 'yes_written', "Yes \u2013 discussed and written down")
        add_checkbox_paragraph(doc, shared == 'yes_not_written', "Yes \u2013 discussed, not written")
        add_checkbox_paragraph(doc, shared == 'not_yet', "Not yet")

        add_field_label(doc, "Notes from discussion:")
        add_field_value(doc, get_value(financial, 'sharing_notes'))

        add_sub_header(doc, "F. Overall Financial Wishes and Legacy Intentions")
        add_field_label(doc, "Summary of what matters most to you about how your financial affairs are handled:")
        add_field_value(doc, get_value(financial, 'overall_wishes'))

        add_field_label(doc, "Are there personal or sentimental items you want to designate for specific people?")
        items_status = financial.get('specific_items_status', 'not_yet_decided')
        add_checkbox_paragraph(doc, items_status == 'has_specific_items', "Yes (list below)")
        add_checkbox_paragraph(doc, items_status == 'no_specific_items', "No specific designations needed")
        add_checkbox_paragraph(doc, items_status == 'not_yet_decided', "Not yet decided")

        items = financial.get('specific_items', [])
        if items and items_status == 'has_specific_items':
            add_field_label(doc, "Items / recipients:")
            for item in items:
                add_field_value(doc, f"\u2013 {item.get('item', 'Item')} \u2192 {item.get('recipient', 'Recipient')}")

    # ===== FUNERAL & MEMORIAL =====
    if 'funeral' in topics:
        doc.add_page_break()
        add_section_header(doc, "My Funeral & Memorial Plans")

        funeral = data.get('funeral_plans', {})

        add_field_label(doc, "Type of service you prefer:")
        service_type = funeral.get('service_type', '')
        other_service = funeral.get('service_type_other') or '____________________'
        add_checkbox_paragraph(doc, service_type == 'funeral', "Funeral service")
        add_checkbox_paragraph(doc, service_type == 'memorial', "Memorial service")
        add_checkbox_paragraph(doc, service_type == 'celebration_of_life', "Celebration of Life")
        add_checkbox_paragraph(doc, service_type == 'other', f"Other: {other_service}")

        add_field_label(doc, "Body preference:")
        body_pref = funeral.get('body_preference', '')
        body_details = funeral.get('body_details') or ''

        burial_loc = body_details if body_pref == 'burial' and body_details else '____________________'
        crem_details = body_details if body_pref == 'cremation' and body_details else '____________________'

        add_checkbox_paragraph(doc, body_pref == 'burial', f"Burial (location: {burial_loc})")
        add_checkbox_paragraph(doc, body_pref == 'cremation', f"Cremation ({crem_details})")
        add_checkbox_paragraph(doc, body_pref == 'donation', "Donation to science")
        add_checkbox_paragraph(doc, body_pref == 'undecided', "Undecided")

        add_field_label(doc, "Summary of conversation:")
        add_field_value(doc, get_value(funeral, 'conversation_summary'))

        add_field_label(doc, "Special requests or details:")

        p = doc.add_paragraph()
        p.add_run("Preferred location: ").bold = True
        p.add_run(get_value(funeral, 'preferred_location'))

        p = doc.add_paragraph()
        p.add_run("Leader of service: ").bold = True
        p.add_run(get_value(funeral, 'service_leader'))

        p = doc.add_paragraph()
        p.add_run("Music/Readings: ").bold = True
        p.add_run(get_value(funeral, 'music_readings'))

        p = doc.add_paragraph()
        p.add_run("Appearance/Clothing: ").bold = True
        p.add_run(get_value(funeral, 'appearance_clothing'))

        p = doc.add_paragraph()
        p.add_run("Charities for donations: ").bold = True
        p.add_run(get_value(funeral, 'charity_donations'))

        add_field_label(doc, "Funeral cost planning:")
        cost = funeral.get('cost_planning', '')
        add_checkbox_paragraph(doc, cost == 'prepaid', "Pre-paid plan")
        add_checkbox_paragraph(doc, cost == 'family_aware', "Family aware of funding")
        add_checkbox_paragraph(doc, cost == 'needs_discussion', "Needs discussion")

        add_field_label(doc, "Additional notes:")
        add_field_value(doc, get_value(funeral, 'additional_notes'))
    else:
        # Funeral not discussed: add a brief note instead of the full section
        doc.add_page_break()
        add_section_header(doc, "My Funeral & Memorial Plans")
        add_field_value(doc, "Funeral and memorial planning was not discussed in this session.")

    # ===== VALUES & REFLECTIONS =====
    add_section_header(doc, "My Values & Life Reflections")

    values = data.get('values_reflections', {})

    add_field_label(doc, "What matters most to me about how I live and am remembered:")
    add_field_value(doc, get_value(values, 'what_matters_most'))

    add_field_label(doc, "What gives my life meaning and joy:")
    add_field_value(doc, get_value(values, 'meaning_and_joy'))

    add_field_label(doc, "What I hope my family and friends remember most:")
    add_field_value(doc, get_value(values, 'want_remembered_for'))

    # ===== NEXT STEPS (data-driven) =====
    add_section_header(doc, "Next Steps & Resources")
    add_field_label(doc, "From today's conversation, the next steps we identified:")

    rec = data.get('recommended_next_steps', {})

    step_labels = [
        ('create_healthcare_poa', "Update or create a Health Care Power of Attorney (POA) or Living Will"),
        ('provide_poa_to_healthcare_team', "Provide copies of my Health Care POA and Living Will to my health care team"),
        ('complete_financial_poa_will_trust', "Complete or update Financial Power of Attorney / Will / Trust"),
        ('review_update_beneficiaries', "Review and update beneficiaries on insurance, retirement, and bank accounts"),
        ('create_financial_info_list', "Create or update a list of key financial information and tell my trusted person where it's stored"),
        ('discuss_wishes_with_loved_ones', "Talk with my loved ones about my wishes for health, finances, and funeral planning"),
        ('store_documents_safely', "Store all important documents safely in a clearly labeled folder or binder at home"),
        ('review_plans_annually', "Review all plans annually or after major life events"),
        ('explore_funeral_preplanning', "Explore funeral or memorial pre-planning options"),
    ]

    for key, label in step_labels:
        checked = is_true(rec.get(key, False))
        add_checkbox_paragraph(doc, checked, label)

    # Handle additional custom steps
    other_steps = rec.get('other_steps', [])
    if isinstance(other_steps, list):
        for step_text in other_steps:
            if step_text and step_text not in ('null', None, ''):
                add_checkbox_paragraph(doc, True, step_text)
    elif isinstance(other_steps, str) and other_steps not in ('null', None, ''):
        add_checkbox_paragraph(doc, True, other_steps)

    add_field_label(doc, "Facilitator Summary or Recommendations:")
    add_field_value(doc, get_value(data, 'facilitator_summary'))

    doc.save(output_path)
    return output_path

# ============================================================================
# MAIN PROCESSING FUNCTION
# ============================================================================

def process_audio(audio_file):
    """Main function to process audio and generate Word document."""
    if audio_file is None:
        return None, "Please record or upload an audio file.", None

    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return None, "API key not configured. Please set GEMINI_API_KEY in Space secrets.", None

    try:
        # Analyze audio
        raw_response = analyze_audio(audio_file, api_key)

        # Parse response
        data = parse_json_response(raw_response)
        if not data:
            return None, "Failed to parse the AI response. Please try again.", None

        # Normalize data
        data = normalize_data(data)

        # Generate Word document
        participant_name = get_value(data, 'participant', 'name', default='Unknown')
        safe_name = re.sub(r'[^a-zA-Z0-9]', '_', participant_name)

        output_dir = tempfile.gettempdir()
        output_filename = os.path.join(output_dir, f"Planning_Summary_{safe_name}.docx")

        generate_docx(data, output_filename)

        # Return results
        json_output = json.dumps(data, indent=2)
        status = f"Successfully generated planning summary for {participant_name}"

        return output_filename, status, json_output

    except Exception as e:
        return None, f"Error: {str(e)}", None


def on_recording_stop(audio_data):
    """
    Called when recording stops. Automatically triggers processing.
    audio_data is a tuple of (sample_rate, audio_array) from microphone recording.
    """
    if audio_data is None:
        return None, "No audio recorded.", None

    # Save the recorded audio to a temporary file
    import numpy as np
    from scipy.io import wavfile

    sample_rate, audio_array = audio_data

    # Create temporary wav file
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, f"recording_{int(time.time())}.wav")

    # Ensure audio is in the right format
    if audio_array.dtype != np.int16:
        # Normalize and convert to int16
        if audio_array.dtype == np.float32 or audio_array.dtype == np.float64:
            audio_array = (audio_array * 32767).astype(np.int16)
        else:
            audio_array = audio_array.astype(np.int16)

    wavfile.write(temp_path, sample_rate, audio_array)

    # Process the audio
    docx_file, status, json_data = process_audio(temp_path)

    return docx_file, status, json_data


def process_uploaded_file(audio_file):
    """Process an uploaded audio file."""
    if audio_file is None:
        return None, "Please upload an audio file.", None

    return process_audio(audio_file)

# ============================================================================
# GRADIO INTERFACE
# ============================================================================

if HAS_GRADIO:
    # Custom theme with neutral colors
    custom_theme = gr.themes.Base(
        primary_hue=gr.themes.colors.slate,
        secondary_hue=gr.themes.colors.gray,
        neutral_hue=gr.themes.colors.gray,
    ).set(
        button_primary_background_fill="#1a1a1a",
        button_primary_background_fill_hover="#333333",
        button_primary_text_color="white",
        block_label_text_color="#374151",
        block_title_text_color="#111827",
    )

    with gr.Blocks(title="Advance Care Planning") as demo:
        gr.Markdown("""
        # Advance Care Planning

        Record or upload an audio conversation to generate a structured Word document summary report.
        """)

        with gr.Tabs():
            with gr.TabItem("Record Audio"):
                gr.Markdown("""
                **Instructions:** Click the microphone button to start recording. Click again to stop.
                The recording will be automatically analyzed when you stop.
                """)

                with gr.Row():
                    with gr.Column(scale=1):
                        audio_recorder = gr.Audio(
                            label="Recording",
                            sources=["microphone"],
                            type="numpy",
                            interactive=True
                        )

                    with gr.Column(scale=1):
                        record_status = gr.Textbox(label="Status", interactive=False)
                        record_docx_output = gr.File(label="Download Word Document")

                with gr.Accordion("View Extracted Data (JSON)", open=False):
                    record_json_output = gr.Code(label="Extracted Data", language="json")

                # Auto-process when recording stops
                audio_recorder.stop_recording(
                    fn=on_recording_stop,
                    inputs=[audio_recorder],
                    outputs=[record_docx_output, record_status, record_json_output]
                )

            with gr.TabItem("Upload Audio"):
                with gr.Row():
                    with gr.Column(scale=1):
                        audio_upload = gr.Audio(
                            label="Upload Audio Recording",
                            type="filepath",
                            sources=["upload"]
                        )

                        upload_btn = gr.Button("Analyze & Generate Word Doc", variant="primary")

                    with gr.Column(scale=1):
                        upload_status = gr.Textbox(label="Status", interactive=False)
                        upload_docx_output = gr.File(label="Download Word Document")

                with gr.Accordion("View Extracted Data (JSON)", open=False):
                    upload_json_output = gr.Code(label="Extracted Data", language="json")

                upload_btn.click(
                    fn=process_uploaded_file,
                    inputs=[audio_upload],
                    outputs=[upload_docx_output, upload_status, upload_json_output]
                )

        gr.Markdown("""
        ---
        **Notes:**
        - Supported audio formats: MP3, WAV, M4A, and other common formats
        - The generated Word document is a summary document, not a legal document
        """)

if __name__ == "__main__":
    if HAS_GRADIO:
        demo.launch(theme=custom_theme)
    else:
        print("Gradio not installed. Core logic is available for import.")