aionboarding/ai-patient-onboarding/app/api/voice_onboarding.py

import base64
import json
from uuid import uuid4
from typing import Optional

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session

from app.database.database import get_db
from app.database.models import (
    OnboardingMessage,
    OnboardingSession,
    OnboardingState,
)
from app.schemas.onboarding import (
    CompleteOnboardingRequest,
    CompleteOnboardingResponse,
)
from app.schemas.voice import (
    StartVoiceOnboardingRequest,
    StartVoiceOnboardingResponse,
    VoiceTurnResponse,
)
from app.services.ai_onboarding_service import AIOnboardingService
from app.services.completion_service import CompletionService
from app.services.onboarding_service import OnboardingService
from app.services.patient_submission_service import PatientSubmissionService
from app.services.result_extraction_service import ResultExtractionService
from app.services.voice_service import VoiceService

router = APIRouter(
    prefix="/v1/voice-onboarding",
    tags=["Voice Onboarding"],
)


@router.post(
    "/start",
    response_model=StartVoiceOnboardingResponse,
)
async def start_voice_onboarding(
    request: StartVoiceOnboardingRequest,
    db: Session = Depends(get_db),
):
    onboarding_service = OnboardingService()
    voice_service = VoiceService()

    session_id = f"VOICE-{uuid4().hex[:8].upper()}"

    session = OnboardingSession(
        session_id=session_id,
        patient_id=request.patient.id,
        patient_name=request.patient.name,
        age=request.patient.age,
        gender=request.patient.gender,
        category_type=request.category.type,
        category_code=request.category.code,
        category_name=request.category.name,
        language=request.language,
        status="in_progress",
        stage="greeting",
    )

    db.add(session)
    db.commit()
    db.refresh(session)

    # Initialize empty state
    onboarding_service.get_state(db, session_id)

    # Generate localized voice greeting
    greeting_text = voice_service.get_initial_greeting(
        request.patient.name,
        request.language,
    )

    # Synthesize audio greeting
    audio_id, greeting_bytes = await voice_service.text_to_speech(
        greeting_text,
        request.language,
    )
    greeting_b64 = f"data:audio/wav;base64,{base64.b64encode(greeting_bytes).decode('utf-8')}" if greeting_bytes else None

    # Save initial assistant message
    onboarding_service.save_message(
        db=db,
        session_id=session_id,
        role="assistant",
        message=greeting_text,
        stage="greeting",
    )

    return StartVoiceOnboardingResponse(
        success=True,
        session_id=session.session_id,
        patient_id=session.patient_id,
        category=request.category,
        language=session.language,
        status=session.status,
        stage=session.stage,
        greeting_text=greeting_text,
        audio_url=f"/v1/voice-onboarding/audio/{audio_id}",
        audio_base64=greeting_b64,
    )


@router.post(
    "/{session_id}/turn",
    response_model=VoiceTurnResponse,
)
async def voice_turn(
    session_id: str,
    audio: Optional[UploadFile] = File(None),
    text_fallback: Optional[str] = Form(None),
    db: Session = Depends(get_db),
):
    onboarding_service = OnboardingService()
    voice_service = VoiceService()
    ai_service = AIOnboardingService()

    # 1. Fetch active session
    session = onboarding_service.get_session(db, session_id)
    if not session:
        raise HTTPException(
            status_code=404,
            detail="Voice onboarding session not found",
        )

    if session.status != "in_progress":
        raise HTTPException(
            status_code=400,
            detail="Voice onboarding session is not active",
        )

    # 2. Fast-path transcript: use real-time Web Speech text_fallback first for zero-latency processing
    patient_transcript = ""
    if text_fallback and text_fallback.strip():
        from app.services.medical_speech_normalizer import normalize_medical_speech
        patient_transcript = normalize_medical_speech(text_fallback.strip(), language_code=session.language)

    # If text_fallback was not available, run server-side multi-model STT
    if not patient_transcript and audio:
        audio_bytes = await audio.read()
        if len(audio_bytes) > 0:
            patient_transcript = await voice_service.speech_to_text(
                audio_bytes=audio_bytes,
                filename=audio.filename or "recording.webm",
                language_code=session.language,
                specialty_code=session.category_code or "general",
                patient_context={
                    "name": session.patient_name,
                    "age": session.age,
                    "gender": session.gender,
                },
            )

    # If STILL empty, user spoke nothing or silence was recorded.
    # Do NOT substitute fake default text or auto-fill clinical fields!
    if not patient_transcript or not patient_transcript.strip():
        lang = session.language.lower()
        if lang.startswith("ta"):
            no_speech_msg = "மன்னிக்கவும், உங்கள் குரல் கேட்கவில்லை. தயவுசெய்து உங்கள் உடல்நலப் பிரச்சினை பற்றி பேசவும்."
        elif lang.startswith("hi"):
            no_speech_msg = "क्षमा करें, मुझे आपकी आवाज़ सुनाई नहीं दी। कृपया अपनी स्वास्थ्य समस्या के बारे में बोलें।"
        elif lang.startswith("te"):
            no_speech_msg = "క్షమించండి, మీ ధ్వని వినిపించలేదు. దయచేసి మీ ఆరోగ్య సమస్య గురించి మాట్లాడండి."
        elif lang.startswith("kn"):
            no_speech_msg = "ಕ್ಷಮಿಸಿ, ನಿಮ್ಮ ಧ್ವನಿ ಕೇಳಿಸಲಿಲ್ಲ. ದಯವಿಟ್ಟು ನಿಮ್ಮ ಆರೋಗ್ಯದ ತೊಂದರೆ ಬಗ್ಗೆ ಮಾತನಾಡಿ."
        else:
            no_speech_msg = "I didn't hear anything. Please speak into your microphone and describe your health symptoms."

        audio_id, _ = await voice_service.text_to_speech(
            text=no_speech_msg,
            language_code=session.language,
        )

        return VoiceTurnResponse(
            success=True,
            session_id=session_id,
            patient_transcript="",
            ai_text=no_speech_msg,
            stage=session.stage,
            completed=False,
            audio_url=f"/v1/voice-onboarding/audio/{audio_id}",
        )



    # 3. Save patient message
    onboarding_service.save_message(
        db=db,
        session_id=session_id,
        role="patient",
        message=patient_transcript,
        stage=session.stage,
    )

    # 4. Fetch state & history
    state = onboarding_service.get_state(db, session_id)
    messages = (
        db.query(OnboardingMessage)
        .filter(OnboardingMessage.session_id == session_id)
        .order_by(OnboardingMessage.created_at.asc())
        .all()
    )
    conversation_history = [
        {
            "role": m.role,
            "message": m.message,
            "stage": m.stage,
        }
        for m in messages
    ]

    # 5. AI Reasoning & Extraction
    ai_response = await ai_service.process_message(
        session=session,
        state=state,
        conversation_history=conversation_history,
        user_message=patient_transcript,
    )

    # 6. Update Structured Clinical State
    current_state = json.loads(state.state_json)
    updates = ai_response.get("state_updates", {})
    for k, v in updates.items():
        if v is not None:
            if isinstance(v, list) and len(v) == 0:
                continue
            current_state[k] = v

    onboarding_service.update_state(
        db=db,
        state=state,
        data=current_state,
    )

    # 7. Update session stage & status
    ai_reply_text = ai_response["reply"]
    session.stage = ai_response.get("next_stage", session.stage)
    is_completed = bool(ai_response.get("completed", False))
    if is_completed:
        session.status = "completed"

    db.commit()

    # 8. Save Assistant Message
    onboarding_service.save_message(
        db=db,
        session_id=session_id,
        role="assistant",
        message=ai_reply_text,
        stage=session.stage,
    )

    # 9. Synthesize AI Speech Audio (TTS)
    audio_id, audio_bytes = await voice_service.text_to_speech(
        text=ai_reply_text,
        language_code=session.language,
    )

    audio_b64 = f"data:audio/wav;base64,{base64.b64encode(audio_bytes).decode('utf-8')}" if audio_bytes else None

    return VoiceTurnResponse(
        success=True,
        session_id=session_id,
        patient_transcript=patient_transcript,
        ai_text=ai_reply_text,
        stage=session.stage,
        completed=is_completed,
        audio_url=f"/v1/voice-onboarding/audio/{audio_id}",
        audio_base64=audio_b64,
    )


@router.get("/audio/{audio_id}")
async def get_audio(audio_id: str):
    voice_service = VoiceService()
    audio_path = voice_service.get_audio_path(audio_id)

    if not audio_path or not audio_path.exists():
        raise HTTPException(status_code=404, detail="Audio file not found")

    media_type = "audio/mpeg" if audio_path.suffix == ".mp3" else "audio/wav"
    return FileResponse(
        path=str(audio_path),
        media_type=media_type,
        filename=audio_path.name,
    )


@router.post(
    "/{session_id}/complete",
    response_model=CompleteOnboardingResponse,
)
async def complete_voice_onboarding(
    session_id: str,
    request: CompleteOnboardingRequest,
    db: Session = Depends(get_db),
):
    onboarding_service = OnboardingService()
    completion_service = CompletionService()
    result_service = ResultExtractionService()
    submission_service = PatientSubmissionService()

    session = onboarding_service.get_session(db, session_id)
    if not session:
        raise HTTPException(
            status_code=404,
            detail="Voice onboarding session not found",
        )

    state = onboarding_service.get_state(db, session_id)

    # Check if there is any conversation to summarize
    has_details = completion_service.is_complete(session, state)


    messages = (
        db.query(OnboardingMessage)
        .filter(OnboardingMessage.session_id == session_id)
        .order_by(OnboardingMessage.created_at.asc())
        .all()
    )
    conversation_history = [
        {"role": m.role, "message": m.message, "stage": m.stage}
        for m in messages
    ]

    result = await result_service.extract_result(
        session=session,
        state=state,
        conversation_history=conversation_history,
    )

    if not request.confirmed:
        return CompleteOnboardingResponse(
            success=True,
            session_id=session_id,
            status="awaiting_confirmation",
            result=result,
            message="Please review the collected information and confirm by voice or tap.",
        )

    submission = await submission_service.submit(result)

    if submission.get("success"):
        db.query(OnboardingMessage).filter(OnboardingMessage.session_id == session_id).delete()
        db.query(OnboardingState).filter(OnboardingState.session_id == session_id).delete()
        db.delete(session)
        db.commit()

        return CompleteOnboardingResponse(
            success=True,
            session_id=session_id,
            status="submitted_and_cleaned",
            result=result,
            message="Voice onboarding successfully saved to patient record.",
        )

    session.status = "submission_pending_retry"
    db.commit()

    return CompleteOnboardingResponse(
        success=False,
        session_id=session_id,
        status="submission_pending_retry",
        result=result,
        message="Patient API submission failed. Session preserved for retry.",
    )