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

import json
from uuid import uuid4

from fastapi import APIRouter, Depends, HTTPException
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,
    SendMessageRequest,
    SendMessageResponse,
    StartOnboardingRequest,
    StartOnboardingResponse,
)
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,
)


router = APIRouter(
    prefix="/v1/onboarding",
    tags=["Patient Onboarding"],
)


@router.post(
    "/start",
    response_model=StartOnboardingResponse,
)
async def start_onboarding(
    request: StartOnboardingRequest,
    db: Session = Depends(get_db),
):
    onboarding_service = OnboardingService()
    session_id = f"ONB-{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)

    greeting_msg = (
        f"Hello {request.patient.name or 'there'}, "
        "welcome to CureSelect Healthcare. "
        "I'll ask you a few questions to understand "
        "your current health concern. Shall we begin?"
    )

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

    return StartOnboardingResponse(
        success=True,
        session_id=session.session_id,
        patient_id=session.patient_id,
        category=request.category,
        language=session.language,
        status=session.status,
        stage=session.stage,
        message=greeting_msg,
    )


@router.post(
    "/{session_id}/message",
    response_model=SendMessageResponse,
)
async def send_message(
    session_id: str,
    request: SendMessageRequest,
    db: Session = Depends(get_db),
):
    onboarding_service = OnboardingService()
    ai_service = AIOnboardingService()

    # 1. Get session
    session = onboarding_service.get_session(
        db,
        session_id,
    )

    if not session:
        raise HTTPException(
            status_code=404,
            detail="Onboarding session not found",
        )

    # 2. Check session
    if session.status != "in_progress":
        raise HTTPException(
            status_code=400,
            detail="Onboarding session is not active",
        )

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

    # 4. Get state
    state = onboarding_service.get_state(
        db,
        session_id,
    )

    # 5. Get conversation history
    messages = (
        db.query(OnboardingMessage)
        .filter(
            OnboardingMessage.session_id == session_id
        )
        .order_by(
            OnboardingMessage.created_at.asc()
        )
        .all()
    )

    conversation_history = [
        {
            "role": message.role,
            "message": message.message,
            "stage": message.stage,
        }
        for message in messages
    ]

    # 6. Send to AI
    ai_response = await ai_service.process_message(
        session=session,
        state=state,
        conversation_history=conversation_history,
        user_message=request.message,
    )

    # 7. Update state
    current_state = json.loads(
        state.state_json
    )

    updates = ai_response.get(
        "state_updates",
        {},
    )

    for key, value in updates.items():
        if value is not None:
            if isinstance(value, list) and len(value) == 0:
                continue
            current_state[key] = value

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

    # 8. Update stage
    session.stage = ai_response.get("next_stage", session.stage)

    # 9. Check completion
    if ai_response.get("completed"):
        session.status = "completed"

    db.commit()

    # 10. Save AI message
    onboarding_service.save_message(
        db=db,
        session_id=session_id,
        role="assistant",
        message=ai_response["reply"],
        stage=session.stage,
    )

    return SendMessageResponse(
        success=True,
        session_id=session_id,
        stage=session.stage,
        message=ai_response["reply"],
        status=session.status,
    )


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

    # 1. Find session
    session = onboarding_service.get_session(
        db,
        session_id,
    )

    if not session:
        raise HTTPException(
            status_code=404,
            detail="Onboarding session not found",
        )

    # 2. Load state
    state = onboarding_service.get_state(
        db,
        session_id,
    )

    # Check if complete
    has_details = completion_service.is_complete(session, state)

    # 4. Get conversation history
    messages = (
        db.query(OnboardingMessage)
        .filter(OnboardingMessage.session_id == session_id)
        .order_by(OnboardingMessage.created_at.asc())
        .all()
    )

    conversation_history = [
        {
            "role": message.role,
            "message": message.message,
            "stage": message.stage,
        }
        for message in messages
    ]

    # 5. Generate final structured result
    result = await result_service.extract_result(
        session=session,
        state=state,
        conversation_history=conversation_history,
    )

    # 6. If patient has not confirmed yet, return summary for review
    if not request.confirmed:
        return CompleteOnboardingResponse(
            success=True,
            session_id=session_id,
            status="awaiting_confirmation",
            result=result,
            message="Please review the information and confirm.",
        )

    # 7. Patient confirmed -> submit to existing patient / hospital API
    submission = await submission_service.submit(result)

    if submission.get("success"):
        # 8. Successful submission -> clean up temporary SQLite session
        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="Onboarding successfully completed, saved to patient record, and temporary session cleaned.",
        )

    # If submission failed, preserve SQLite session for retry
    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.",
    )