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

import json
import time
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,
)
from app.services.voice_service import VoiceService


from fastapi import APIRouter, Depends, HTTPException, Request
from sqlalchemy.orm import Session
from app.core.security import get_auth_context, AuthContext, check_rate_limit

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


@router.post(
    "/start",
    response_model=StartOnboardingResponse,
)
async def start_onboarding(
    request: StartOnboardingRequest,
    req_obj: Request = None,
    auth: AuthContext = Depends(get_auth_context),
    db: Session = Depends(get_db),
):
    if req_obj and req_obj.client:
        check_rate_limit(req_obj.client.host, limit=30, window_sec=60)

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

    # Use authenticated patient identity
    patient_id = request.patient.id or auth.patient_id

    session = OnboardingSession(
        session_id=session_id,
        patient_id=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)

    voice_service = VoiceService()
    greeting_msg = voice_service.get_initial_greeting(
        request.patient.name,
        request.language,
    )

    # 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,
    req_obj: Request = None,
    auth: AuthContext = Depends(get_auth_context),
    db: Session = Depends(get_db),
):
    if req_obj and req_obj.client:
        check_rate_limit(req_obj.client.host, limit=30, window_sec=60)

    t0 = time.perf_counter()
    onboarding_service = OnboardingService()
    ai_service = AIOnboardingService()

    # 1. Get session
    t_sess_start = time.perf_counter()
    session = onboarding_service.get_session(
        db,
        session_id,
    )
    t_sess_end = time.perf_counter()

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

    # 1b. Verify Session Ownership (IDOR Prevention)
    if session.patient_id != auth.patient_id and auth.role != "admin":
        raise HTTPException(
            status_code=403,
            detail="Forbidden: You do not have authorization to access or modify this patient onboarding session.",
        )

    # 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
    t_state_start = time.perf_counter()
    state = onboarding_service.get_state(
        db,
        session_id,
    )
    t_state_end = time.perf_counter()

    # 5. Get conversation history
    t_hist_start = time.perf_counter()
    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
    ]
    t_hist_end = time.perf_counter()

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

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

    t_db_start = time.perf_counter()
    db.commit()
    t_db_end = time.perf_counter()

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

    t_total_end = time.perf_counter()
    total_ms = round((t_total_end - t0) * 1000, 2)
    total_sec = round(t_total_end - t0, 3)

    print(f"[PERF PROFILE - send_message] SESSION_LOOKUP: {(t_sess_end - t_sess_start)*1000:.2f}ms | STATE_LOAD: {(t_state_end - t_state_start)*1000:.2f}ms | HIST_LOAD: {(t_hist_end - t_hist_start)*1000:.2f}ms | AI_PROCESS: {(t_ai_end - t_ai_start)*1000:.2f}ms | DB_COMMIT: {(t_db_end - t_db_start)*1000:.2f}ms | TOTAL: {total_ms}ms")

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


@router.post(
    "/{session_id}/complete",
    response_model=CompleteOnboardingResponse,
)
async def complete_onboarding(
    session_id: str,
    request: CompleteOnboardingRequest,
    auth: AuthContext = Depends(get_auth_context),
    db: Session = Depends(get_db),
):
    start_time = time.time()
    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",
        )

    # 1b. Verify Session Ownership (IDOR Prevention)
    if session.patient_id != auth.patient_id and auth.role != "admin":
        raise HTTPException(
            status_code=403,
            detail="Forbidden: You do not have authorization to complete this onboarding session.",
        )

    # 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 isinstance(result, dict):
        result["onboarding_completed"] = True

    if session.status == "completed":
        return CompleteOnboardingResponse(
            success=True,
            session_id=session_id,
            status="completed",
            result=result,
            message="Onboarding session is already completed.",
            response_time_ms=round((time.time() - start_time) * 1000, 2),
            response_time_sec=round(time.time() - start_time, 3),
        )

    if submission.get("success"):
        # Mark session completed in DB instead of hard-deleting
        session.status = "completed"
        session.stage = "confirmation"
        db.commit()

        elapsed = time.time() - start_time
        return CompleteOnboardingResponse(
            success=True,
            session_id=session_id,
            status="completed",
            result=result,
            message="Onboarding successfully completed and saved to patient record.",
            response_time_ms=round(elapsed * 1000, 2),
            response_time_sec=round(elapsed, 3),
        )

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