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

import os
import re
import json
import time
import tempfile
import gradio as gr
import google.generativeai as genai
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

# ============================================================================
# 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
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",
    "facilitator": "Facilitator name and credentials",
    "location": "Location of conversation"
  },
  "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",
      "relationship": "relationship",
      "phone": "phone or null"
    },
    "values_summary": "2-3 sentence summary of their values",
    "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 | unsure",
    "treatment_details": "Details about 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",
      "relationship": "relationship",
      "phone": "phone or null"
    },
    "financial_conversation_summary": "Summary of who handles finances",
    "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 401k, pension, retirement accounts",
    "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",
    "has_specific_items": true,
    "specific_items": [
      {"item": "Fishing gear", "recipient": "Son Charlie"},
      {"item": "Guitar", "recipient": "Daughter Betsy"}
    ]
  },
  "funeral_plans": {
    "service_type": "MUST BE ONE OF: funeral | memorial | celebration_of_life | other",
    "service_type_other": "If other, describe",
    "body_preference": "MUST BE ONE OF: burial | cremation | donation | undecided",
    "body_details": "Location details like 'Ashes spread at favorite fishing lake'",
    "conversation_summary": "Summary of funeral wishes",
    "preferred_location": "Where service should be held",
    "service_leader": "Who should lead",
    "music_readings": "Music and reading preferences",
    "appearance_clothing": "Dress code preferences",
    "charity_donations": "Charity for donations",
    "cost_planning": "MUST BE ONE OF: prepaid | family_aware | needs_discussion",
    "additional_notes": "Notes about life insurance, costs, etc."
  },
  "values_reflections": {
    "what_matters_most": "How they want to be remembered",
    "meaning_and_joy": "What gives their life meaning",
    "want_remembered_for": "What they hope people remember"
  },
  "facilitator_summary": "Facilitator's closing summary and recommendations"
}
```
DECISION GUIDE FOR COMMON SCENARIOS:
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 want treatment/CPR/ventilation IF there's hope of recovery -> "full_treatment_if_recovery"
- If they only want comfort/palliative care, no machines -> "comfort_care_only"
- If they're unsure or need more information -> "unsure"
beneficiary_status:
- If they're not sure who's listed or need to check -> "unsure"
- If they know some need updating -> "need_to_update"
- If everything is current -> "all_current"
has_info_list:
- If they have files/info but haven't shared 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 written down -> "yes_not_written"
- If they've discussed AND written it down -> "yes_written"
- If they haven't discussed wishes yet -> "not_yet"
service_type:
- Celebration of life, casual gathering, party -> "celebration_of_life"
- Traditional funeral -> "funeral"
- Memorial service -> "memorial"
body_preference:
- Cremation, ashes spread somewhere -> "cremation"
- Burial in cemetery/ground -> "burial"
- Donate body to science -> "donation"
cost_planning:
- 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"
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."""
    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-2.0-flash-exp')
    
    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'],
        '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'],
        'comfort_care_only': ['comfort', 'comfort_care', 'palliative', 'no machines'],
        'unsure': ['unsure', 'not sure', 'uncertain', 'undecided', 'need more info'],
        'need_to_update': ['need_to_update', 'needs update', 'update', 'outdated'],
        'all_current': ['all_current', 'current', 'up to date', 'all updated'],
        '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', 'none', "haven't created", 'not yet'],
        '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', 'no', "haven't", 'not discussed'],
        'celebration_of_life': ['celebration', 'celebration_of_life', 'party', 'gathering', 'casual'],
        'funeral': ['funeral', 'traditional'],
        'memorial': ['memorial', 'memorial_service'],
        'other': ['other'],
        'cremation': ['cremation', 'cremate', 'ashes', 'cremated'],
        'burial': ['burial', 'bury', 'buried', 'cemetery', 'ground'],
        'donation': ['donation', 'donate', 'science', 'donate body'],
        'undecided': ['undecided', 'unsure', 'not sure'],
        'family_aware': ['family_aware', 'family aware', 'life insurance', 'insurance', 'covered'],
        'prepaid': ['prepaid', 'pre-paid', 'pre paid', 'paid'],
        'needs_discussion': ['needs_discussion', 'need to discuss', 'not discussed']
    }
    
    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
    
    # 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', 'unsure'],
            'full_treatment_if_recovery'
        )
    
    # 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'
        )
    
    # 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'],
            'celebration_of_life'
        )
        fun['body_preference'] = normalize_value(
            fun.get('body_preference'),
            ['burial', 'cremation', 'donation', 'undecided'],
            'cremation'
        )
        fun['cost_planning'] = normalize_value(
            fun.get('cost_planning'),
            ['prepaid', 'family_aware', 'needs_discussion'],
            'family_aware'
        )
    
    return data

# ============================================================================
# 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)
    
    # ===== 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 =====
    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')} – {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')} – {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 if dying)")
    add_checkbox_paragraph(doc, treatment == 'full_treatment_if_recovery', "Full medical treatment if recovery possible")
    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 =====
    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')} – {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')} – {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 – my trusted person knows where it is")
    add_checkbox_paragraph(doc, info_status == 'yes_not_shared', "Yes – 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 – discussed and written down")
    add_checkbox_paragraph(doc, shared == 'yes_not_written', "Yes – 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?")
    has_items = is_true(financial.get('has_specific_items'))
    add_checkbox_paragraph(doc, has_items, "Yes (list below)")
    add_checkbox_paragraph(doc, not has_items, "Not yet decided")
    
    items = financial.get('specific_items', [])
    if items and has_items:
        add_field_label(doc, "Items / recipients:")
        for item in items:
            add_field_value(doc, f"– {item.get('item', 'Item')} → {item.get('recipient', 'Recipient')}")
    
    # ===== FUNERAL & MEMORIAL =====
    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'))
    
    # ===== 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 =====
    add_section_header(doc, "Next Steps & Resources")
    add_field_label(doc, "From today's conversation, the next steps we identified:")
    
    add_checkbox_paragraph(doc, True, "Update or create a Health Care Power of Attorney (POA) or Living Will")
    add_checkbox_paragraph(doc, True, "Provide copies of my Health Care POA and Living Will to my health care team")
    add_checkbox_paragraph(doc, True, "Complete or update Financial Power of Attorney / Will / Trust")
    add_checkbox_paragraph(doc, True, "Review and update beneficiaries on insurance, retirement, and bank accounts")
    add_checkbox_paragraph(doc, True, "Create or update a list of key financial information and tell my trusted person where it's stored")
    add_checkbox_paragraph(doc, True, "Talk with my loved ones about my wishes for health, finances, and funeral planning")
    add_checkbox_paragraph(doc, True, "Store all important documents safely in a clearly labeled folder or binder at home")
    add_checkbox_paragraph(doc, True, "Review all plans annually or after major life events")
    
    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
# ============================================================================

# 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__":
    demo.launch(theme=custom_theme)