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
  8. Resilience: retry logic (up to 3 attempts), robust JSON parser, raw response
     surfaced on failure for debugging
  9. Hallucination guard: prompt instructs model to return not_a_planning_conversation
     flag when audio does not contain advance care planning content; validation rejects
     such responses before generating a report.
"""

import os
import re
import json
import time
import logging
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

logger = logging.getLogger(__name__)

# ============================================================================
# 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.

IMPORTANT - NON-RELEVANT AUDIO DETECTION:
Before extracting any planning data, first determine whether this audio actually
contains an advance care planning conversation. The audio MUST contain a real
discussion about at least one of these topics: health care wishes, financial
planning, funeral/memorial preferences, or values and legacy.

If the audio is:
- Silence, noise, music, or unintelligible speech
- A conversation about unrelated topics (e.g. casual chat, a lecture, a podcast
  not about advance care planning)
- Too short or too unclear to extract meaningful planning information
- Random test audio or microphone checks

Then return ONLY this JSON and nothing else:
```json
{
  "not_a_planning_conversation": true,
  "reason": "Brief explanation of what the audio actually contains"
}
```

Do NOT invent, fabricate, or hallucinate planning data. If you cannot clearly hear
a real advance care planning discussion, you MUST return the above JSON.

Only proceed with the full extraction below if you are confident the audio contains
a genuine advance care planning conversation.

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.

    Uses a multi-strategy approach so that trailing prose, markdown fences,
    or minor formatting differences do not cause a parse failure.
    """
    if not response_text:
        return None

    # Strip markdown code fences
    cleaned = re.sub(r'```json\s*', '', response_text)
    cleaned = re.sub(r'```\s*', '', cleaned)
    cleaned = cleaned.strip()

    # Strategy 1: try parsing the entire cleaned text directly
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        pass

    # Strategy 2: locate the first '{' and try progressively shorter
    # substrings ending at each '}' from the end backward.  This handles
    # cases where Gemini appends explanatory text after the JSON object.
    start = cleaned.find('{')
    if start == -1:
        return None

    for end in range(len(cleaned) - 1, start, -1):
        if cleaned[end] == '}':
            try:
                return json.loads(cleaned[start:end + 1])
            except json.JSONDecodeError:
                continue

    return None


def _is_valid_planning_data(data: dict) -> tuple[bool, str]:
    """Check whether parsed data represents a genuine planning conversation.

    Returns (is_valid, reason).  When is_valid is False, reason contains a
    user-facing message explaining why the audio was rejected.
    """
    if not data:
        return False, "No data could be extracted from the audio."

    # Explicit flag set by the model when the audio is not relevant
    if data.get("not_a_planning_conversation"):
        reason = data.get("reason", "The audio does not appear to contain an advance care planning conversation.")
        return False, (
            "This audio does not contain an advance care planning conversation. "
            f"({reason}) Please record or upload a conversation that discusses "
            "health care wishes, financial planning, or funeral preferences."
        )

    # Secondary heuristic: if none of the core sections are present, the
    # model may have returned something unexpected.
    has_any_section = any(
        key in data for key in (
            "health_care_wishes", "financial_planning", "funeral_plans",
            "participant", "topics_discussed",
        )
    )
    if not has_any_section:
        return False, (
            "The audio could not be matched to a planning conversation. "
            "Please make sure the recording contains a discussion about "
            "health care wishes, financial planning, or funeral preferences."
        )

    return True, ""


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
# ============================================================================

MAX_RETRIES = 3
RETRY_DELAY_SECONDS = 3


def process_audio(audio_file):
    """Main function to process audio and generate Word document.

    Retries up to MAX_RETRIES times when the Gemini response cannot be
    parsed, since transient malformed responses are the most common
    failure mode.
    """
    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:
        # Retry loop: the Gemini API occasionally returns responses that
        # cannot be parsed (extra prose, truncated JSON, etc.).  A simple
        # retry with a short delay resolves this the vast majority of the
        # time without any user intervention.
        raw_response = None
        data = None

        for attempt in range(1, MAX_RETRIES + 1):
            raw_response = analyze_audio(audio_file, api_key)
            data = parse_json_response(raw_response)
            if data is not None:
                break
            logger.warning(
                "Attempt %d/%d: failed to parse Gemini response (length=%d)",
                attempt, MAX_RETRIES, len(raw_response) if raw_response else 0,
            )
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY_SECONDS)

        if not data:
            # Surface the raw response so the user (or developer) can
            # inspect what Gemini actually returned.
            return (
                None,
                f"Failed to parse the AI response after {MAX_RETRIES} attempts. "
                "Open the JSON panel below to see the raw API output.",
                raw_response,
            )

        # ---- Hallucination guard ----
        is_valid, rejection_reason = _is_valid_planning_data(data)
        if not is_valid:
            return (
                None,
                rejection_reason,
                json.dumps(data, indent=2),
            )

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