import asyncio
import base64
import logging
import re
import time
import uuid
from datetime import datetime, timedelta, timezone
from typing import Optional

import httpx
import websockets.exceptions
from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile, File, Form, WebSocket
from fastapi.responses import JSONResponse, Response
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, delete
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.exc import IntegrityError
from pydantic import BaseModel, field_validator

from app.core.deps import get_db, get_current_staff, decode_staff_token
from app.core.config import settings
from app.models.patient import Patient
from app.models.intake_session import IntakeSession, SessionContributor
from app.models.question_response import QuestionResponse
from app.models.patient_summary import PatientSummary
from app.models.user import User
from app.models.mstr_role import MstrRole
from app.models.mstr_department import MstrDepartment
from app.models.mstr_question import MstrQuestion
from app.models.question_dept_map import QuestionDeptMap
from app.models.active_session import ActiveSession
from app.models.patient_lock import PatientLock
from app.models.patient_record import PatientRecord
from app.models.conversation_utterance import ConversationUtterance
from app.models.conversation_custom_question import ConversationCustomQuestion
from app.services.gemini import (
    get_first_question, get_next_question, generate_summary,
    extract_from_image, transliterate_to_roman, extract_patient_answer,
    classify_and_map_utterances, segment_and_map_utterance,
)
from app.services.sarvam import (
    transcribe_audio as sarvam_transcribe, transcribe_audio_diarized as sarvam_transcribe_diarized,
    translate_to_english as sarvam_translate, translate_to_language as sarvam_translate_to_lang,
    text_to_speech as sarvam_tts, SUPPORTED_LANGUAGES,
    SarvamStreamSession, pcm_to_wav_bytes, STREAM_SAMPLE_RATE, STREAM_RECONNECT_ATTEMPTS, STREAM_RECONNECT_BACKOFF_S,
)
from app.constants import error_codes
from app.constants.session_status import ACTIVE, FIRST_ROUND_COMPLETE, COMPLETED
from app.constants.role_keys import SR_CONSULTANT, ADMIN, SUPER_ADMIN

REVIEWER_ROLE_KEYS = (SUPER_ADMIN, ADMIN, SR_CONSULTANT)

logger = logging.getLogger(__name__)

router = APIRouter()

MAX_ANSWER_TEXT_LENGTH = 2000
MAX_NOTE_LENGTH = 2000
MAX_TRANSCRIPTION_ATTEMPTS = 2
CONVERSATION_MAPPING_CONFIDENCE_THRESHOLD = 0.6
CONVERSATION_RECENT_CONTEXT_SIZE = 6

# Gemini classify calls run strictly one at a time per session (each depends on the
# previous one's committed context), so retries here must stay quick — a stuck attempt
# blocks every classify call queued behind it, not just its own. The attempt count itself
# is env-configurable (settings.CLASSIFY_MAX_ATTEMPTS) — backoff is computed from a formula
# rather than a fixed lookup table so it stays correct regardless of how high that's set.
CLASSIFY_RETRY_BASE_BACKOFF_S = 0.5
CLASSIFY_RETRY_MAX_BACKOFF_S = 4.0


def _classify_retry_backoff_s(attempt: int) -> float:
    return min(CLASSIFY_RETRY_BASE_BACKOFF_S * (2 ** (attempt - 1)), CLASSIFY_RETRY_MAX_BACKOFF_S)

# WebSocket close codes used by /conversation-stream — outside the reserved 0-2999 range,
# following RFC 6455 §7.4.2's guidance for application-specific codes.
WS_CODE_AUTH_FAILED = 4401
WS_CODE_NOT_FOUND = 4404
WS_CODE_CONFLICT = 4409
WS_CODE_LOCK_LOST = 4423
WS_CODE_FALLBACK_REQUIRED = 4503

FLAGGED_PHRASES = {
    "no", "na", "n/a", "nil", "none", "unknown",
    "not sure", "don't know", "i don't know",
    "doesn't know", "not known", "not applicable",
}

_IMPLICIT_MARKER = "[IMPLICITLY_ANSWERED]"


# ── Request schemas ────────────────────────────────────────────────────────────

class StartSessionRequest(BaseModel):
    patientId: str

    @field_validator("patientId")
    @classmethod
    def validate_patient_id(cls, v):
        v = v.strip()
        try:
            uuid.UUID(v)
        except ValueError:
            raise ValueError("patientId must be a valid UUID")
        return v


class CompleteSessionRequest(BaseModel):
    additionalNote: Optional[str] = None

    @field_validator("additionalNote")
    @classmethod
    def strip_note(cls, v):
        if v is not None:
            v = v.strip()
            if len(v) > MAX_NOTE_LENGTH:
                raise ValueError(f"additionalNote max {MAX_NOTE_LENGTH} chars")
            return v if v else None
        return v


class RespondRequest(BaseModel):
    questionId: Optional[str] = None
    questionText: Optional[str] = None
    text: str
    inputType: str = "TEXT"
    inputLanguage: str = "te-IN"

    @field_validator("text")
    @classmethod
    def text_not_empty(cls, v):
        v = v.strip()
        if not v:
            raise ValueError("text cannot be empty")
        if len(v) > MAX_ANSWER_TEXT_LENGTH:
            raise ValueError(f"text max {MAX_ANSWER_TEXT_LENGTH} chars")
        return v

    @field_validator("inputType")
    @classmethod
    def validate_input_type(cls, v):
        if v not in ("VOICE", "TEXT"):
            raise ValueError("inputType must be VOICE or TEXT")
        return v

    @field_validator("inputLanguage")
    @classmethod
    def validate_input_language(cls, v):
        return v.strip() if v else "ENGLISH"


# ── Helper: load department prompts from DB ───────────────────────────────────

async def _get_dept_context(db: AsyncSession, dept_id: uuid.UUID) -> tuple[str, str, list[dict]]:
    """Returns (llm_system_prompt, llm_summary_prompt, questions) from DB."""
    result = await db.execute(select(MstrDepartment).where(MstrDepartment.id == dept_id))
    dept = result.scalar_one_or_none()
    if not dept:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.DEPARTMENT_NOT_FOUND, "message": "Department not found"},
        )
    q_result = await db.execute(
        select(MstrQuestion)
        .join(QuestionDeptMap, QuestionDeptMap.question_id == MstrQuestion.id)
        .where(QuestionDeptMap.department_id == dept_id)
        .where(MstrQuestion.is_active == True)
        .order_by(QuestionDeptMap.sequence_number)
    )
    questions = [
        {"id": str(q.id), "text": q.text}
        for q in q_result.scalars().all()
    ]
    return dept.llm_system_prompt, dept.llm_summary_prompt, questions


# ── Helper: which standard bank questions remain unasked in this session ───────

def _remaining_questions_for(responses, all_questions: list[dict], extra_excluded: set[str] | None = None) -> list:
    asked_ids = {str(r.question_id) for r in responses if r.question_id}
    if extra_excluded:
        asked_ids |= extra_excluded
    return [q for q in all_questions if q["id"] not in asked_ids]


# ── Helper: build intakeHistory from session_contributors + staff ──────────────

# ── Helper: build intakeHistory from session_contributors + user ──────────────

async def _build_intake_history(db: AsyncSession, session_id: uuid.UUID) -> list[dict]:
    result = await db.execute(
        select(SessionContributor, User, MstrRole)
        .join(User, SessionContributor.staff_id == User.id)
        .join(MstrRole, User.role_id == MstrRole.id)
        .where(SessionContributor.session_id == session_id)
        .order_by(SessionContributor.sequence_number)
    )
    rows = result.all()
    return [
        {
            "staffName": user.name,
            "role": role.name,
            "joinedAt": contributor.joined_at.isoformat(),
            "sequenceNumber": contributor.sequence_number,
        }
        for contributor, user, role in rows
    ]


# ── Helper: reject if another staff member's lock on this patient is still valid ──

def _ensure_lock_not_held_by_other_staff(lock: PatientLock | None, staff_id: uuid.UUID, now: datetime) -> None:
    if lock and lock.expires_at > now and lock.staff_id != staff_id:
        raise HTTPException(
            status_code=409,
            detail={"error": error_codes.SESSION_IN_USE, "message": "This patient is being seen by another staff member"}
        )


# ── GET /sessions ─────────────────────────────────────────────────────────────

@router.get("")
async def list_sessions(
    sort_by: str = Query("createdAt", pattern="^(name|id|age|gender|createdAt)$"),
    order: str = Query("desc", pattern="^(asc|desc)$"),
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    dept_id = uuid.UUID(staff["departmentId"])

    # Subquery: response count per session so we can prefer non-empty sessions in dedup
    response_counts = (
        select(
            QuestionResponse.session_id,
            func.count(QuestionResponse.id).label("cnt"),
        )
        .group_by(QuestionResponse.session_id)
        .subquery()
    )

    # Subquery: name of the staff member who recorded the first response for each session
    first_responder = (
        select(
            QuestionResponse.session_id,
            User.name.label("staff_name"),
        )
        .join(User, QuestionResponse.responded_by_staff_id == User.id)
        .where(QuestionResponse.sequence_number == 1)
        .subquery()
    )

    result = await db.execute(
        select(IntakeSession, Patient, MstrDepartment,
               func.coalesce(response_counts.c.cnt, 0).label("response_count"),
               first_responder.c.staff_name)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .join(MstrDepartment, IntakeSession.department_id == MstrDepartment.id)
        .outerjoin(response_counts, IntakeSession.id == response_counts.c.session_id)
        .outerjoin(first_responder, IntakeSession.id == first_responder.c.session_id)
        .where(IntakeSession.department_id == dept_id)
        # Sessions with responses come first; within same patient ties by most recent
        .order_by(
            func.coalesce(response_counts.c.cnt, 0).desc(),
            IntakeSession.started_at.desc(),
        )
    )
    rows = result.all()

    # One row per patient (by daily_id).
    # Because rows are ordered by response_count DESC then started_at DESC,
    # the first hit per daily_id is the session with the most data (not an empty orphan).
    seen: set[str] = set()
    unique_rows = []
    for session, patient, department, response_count, staff_name in rows:
        if patient.daily_id not in seen:
            seen.add(patient.daily_id)
            unique_rows.append((session, patient, department, response_count, staff_name))

    reverse = order == "desc"
    sort_key_map = {
        "name":      lambda r: (r[1].name or "").lower(),
        "id":        lambda r: (r[1].daily_id or "").lower(),
        "age":       lambda r: r[1].age if r[1].age is not None else 0,
        "gender":    lambda r: (r[1].gender or "").lower(),
        "createdAt": lambda r: r[0].started_at,
    }
    unique_rows.sort(key=sort_key_map[sort_by], reverse=reverse)

    return {
        "sessions": [
            {
                "id": str(session.id),
                "patientName": patient.name,
                "dailyPatientId": patient.daily_id,
                "age": patient.age,
                "gender": patient.gender,
                "department": department.name,
                "createdAt": session.started_at.isoformat(),
                "status": session.status.lower(),
                # An ACTIVE session with zero recorded answers hasn't actually been opened
                # yet — the frontend uses this to show "Start Intake" instead of "In Progress".
                "hasResponses": response_count > 0,
                "staffName": staff_name,
            }
            for session, patient, department, response_count, staff_name in unique_rows
        ],
        "total": len(unique_rows),
    }


# ── POST /sessions ─────────────────────────────────────────────────────────────

@router.post("")
async def create_session(
    body: StartSessionRequest,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    dept_id = uuid.UUID(staff["departmentId"])
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)
    lock_expires = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)
    patient_id = uuid.UUID(body.patientId)

    # Step 1 — Verify patient exists
    patient_result = await db.execute(
        select(Patient).where(Patient.id == patient_id)
    )
    patient = patient_result.scalar_one_or_none()
    if not patient:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.PATIENT_NOT_FOUND, "message": "Patient not found"}
        )

    # Step 2 — A patient can have at most one intake_session row, ever (see the unique
    # constraint on IntakeSession.patient_id) — look it up by patient alone rather than
    # assuming "no active session in this department" means no row exists at all.
    existing = await db.execute(
        select(IntakeSession).where(IntakeSession.patient_id == patient_id)
    )
    session = existing.scalar_one_or_none()

    if session and session.status == ACTIVE and session.department_id != dept_id:
        raise HTTPException(
            status_code=409,
            detail={
                "error": error_codes.SESSION_IN_USE,
                "message": "This patient already has an active session in another department",
            },
        )

    if session and session.status != ACTIVE:
        raise HTTPException(
            status_code=409,
            detail={"error": error_codes.SESSION_COMPLETED, "message": "Intake is already completed for this patient"}
        )

    # ── RESUME existing session ────────────────────────────────────────────────
    if session:
        # Check patient lock
        lock_result = await db.execute(
            select(PatientLock).where(PatientLock.patient_id == patient.id)
        )
        lock = lock_result.scalar_one_or_none()

        _ensure_lock_not_held_by_other_staff(lock, staff_id, now)

        # UPSERT patient lock for this staff
        await db.execute(
            pg_insert(PatientLock)
            .values(id=uuid.uuid4(), patient_id=patient.id, staff_id=staff_id, expires_at=lock_expires)
            .on_conflict_do_update(
                index_elements=['patient_id'],
                set_={'staff_id': staff_id, 'expires_at': lock_expires},
            )
        )

        # Update active_session to link to this session
        active_result = await db.execute(
            select(ActiveSession).where(ActiveSession.staff_id == staff_id)
        )
        active_slot = active_result.scalar_one_or_none()
        if active_slot:
            active_slot.session_id = session.id
            active_slot.expires_at = lock_expires

        # Check if this staff is already a contributor; if not, add them
        existing_contributor = await db.execute(
            select(SessionContributor).where(
                SessionContributor.session_id == session.id,
                SessionContributor.staff_id == staff_id,
            )
        )
        if not existing_contributor.scalar_one_or_none():
            count_result = await db.execute(
                select(func.count()).where(SessionContributor.session_id == session.id)
            )
            next_seq = count_result.scalar() + 1
            db.add(SessionContributor(
                id=uuid.uuid4(),
                session_id=session.id,
                staff_id=staff_id,
                sequence_number=next_seq,
            ))
        await db.commit()

        # Load intake history
        intake_history = await _build_intake_history(db, session.id)

        # Load all responses so far
        responses_result = await db.execute(
            select(QuestionResponse, User.name.label("staff_name"))
            .join(User, QuestionResponse.responded_by_staff_id == User.id)
            .where(QuestionResponse.session_id == session.id)
            .order_by(QuestionResponse.sequence_number)
        )
        all_response_rows = responses_result.all()
        all_responses = [r for r, _ in all_response_rows]

        responses = [
            {
                "questionId": r.question_id,
                "questionText": r.question_text,
                "answer": r.transcribed_text,
                "sequenceNumber": r.sequence_number,
                "isFlagged": r.is_flagged,
                "respondedByStaffName": staff_name,
            }
            for r, staff_name in all_response_rows
            if r.transcribed_text != _IMPLICIT_MARKER
        ]

        system_prompt, _, questions = await _get_dept_context(db, session.department_id)

        # Load patient records so the LLM can skip questions already answered by uploaded docs
        rec_result = await db.execute(
            select(PatientRecord).where(PatientRecord.session_id == session.id)
        )
        session_records = [r.extracted_text for r in rec_result.scalars().all() if r.extracted_text]

        # Re-evaluate the next question using the updated prompt so implicit answers
        # (e.g. "cold for two days" answering both chief complaint AND onset/duration)
        # are correctly recognised and that question is skipped.
        next_question = None
        skipped_ids: set[str] = set()
        if all_responses:
            conversation = [
                {"question": r.question_text, "answer": r.transcribed_text}
                for r in all_responses
                if r.transcribed_text != _IMPLICIT_MARKER
            ]
            llm_result = await get_next_question(
                patient_info={"name": patient.name, "age": patient.age, "gender": patient.gender, "daily_id": patient.daily_id},
                conversation=conversation,
                latest_answer=next((r.transcribed_text for r in reversed(all_responses) if r.transcribed_text != _IMPLICIT_MARKER), ""),
                system_prompt=system_prompt,
                questions=questions,
                patient_records=session_records or None,
            )
            skipped_ids = set(llm_result.get("skippedIds", []))
            existing_ids = {str(r.question_id) for r in all_responses if r.question_id}
            questions_map = {q["id"]: q["text"] for q in questions}
            max_seq = max((r.sequence_number for r in all_responses), default=0)
            for i, skipped_id in enumerate(sorted(skipped_ids - existing_ids)):
                q_text = questions_map.get(skipped_id)
                if q_text:
                    db.add(QuestionResponse(
                        id=uuid.uuid4(),
                        session_id=session.id,
                        question_id=skipped_id,
                        question_text=q_text,
                        transcribed_text=_IMPLICIT_MARKER,
                        input_type="TEXT",
                        input_language="ENGLISH",
                        is_flagged=False,
                        sequence_number=max_seq + i + 1,
                        responded_by_staff_id=staff_id,
                    ))
            if not llm_result["intakeComplete"]:
                next_question = llm_result["nextQuestion"]
            session.pending_question = next_question["text"] if next_question else None
            await db.commit()
        elif session.pending_question:
            # No prior answers — use the stored first question as-is (nothing to validate against)
            next_question = {"id": None, "text": session.pending_question}

        return JSONResponse(
            status_code=200,
            content={
                "sessionId": str(session.id),
                "status": "RESUMED",
                "patient": {
                    "dailyId": patient.daily_id,
                    "name": patient.name,
                    "age": patient.age,
                    "gender": patient.gender,
                },
                "intakeHistory": intake_history,
                "responses": responses,
                "resumeFromQuestion": None,
                "nextQuestion": next_question,
                "remainingQuestions": _remaining_questions_for(all_responses, questions, extra_excluded=skipped_ids),
                "questions": questions,
                "streamingEnabled": settings.CONVERSATION_STREAMING_ENABLED,
            }
        )

    # ── CREATE new session ─────────────────────────────────────────────────────
    system_prompt, _, questions = await _get_dept_context(db, dept_id)


    # Step 3 — Create session for the existing patient
    session = IntakeSession(
        id=uuid.uuid4(),
        patient_id=patient.id,
        department_id=dept_id,
        status=ACTIVE,
    )
    db.add(session)
    try:
        await db.flush()
    except IntegrityError:
        # Another request created this patient's session between our check above and this
        # insert — the DB's patient_id-unique constraint is the real source of truth here.
        await db.rollback()
        raise HTTPException(
            status_code=409,
            detail={
                "error": error_codes.SESSION_IN_USE,
                "message": "Session creation for this patient conflicted with a concurrent request — please retry",
            },
        )

    # Step 4 — Add staff as first contributor
    db.add(SessionContributor(
        id=uuid.uuid4(),
        session_id=session.id,
        staff_id=staff_id,
        sequence_number=1,
    ))

    # Step 5 — Acquire patient lock (upsert — safe when two staff race on the same patient)
    await db.execute(
        pg_insert(PatientLock)
        .values(id=uuid.uuid4(), patient_id=patient.id, staff_id=staff_id, expires_at=lock_expires)
        .on_conflict_do_update(
            index_elements=['patient_id'],
            set_={'staff_id': staff_id, 'expires_at': lock_expires},
        )
    )

    # Step 6 — Link active_session to this session
    active_result = await db.execute(
        select(ActiveSession).where(ActiveSession.staff_id == staff_id)
    )
    active_slot = active_result.scalar_one_or_none()
    if active_slot:
        active_slot.session_id = session.id
        active_slot.expires_at = lock_expires

    await db.commit()

    # Step 7 — Get first question and intake history
    first = await get_first_question(
        patient_info={
            "name": patient.name,
            "age": patient.age,
            "gender": patient.gender,
            "daily_id": patient.daily_id,
        },
        system_prompt=system_prompt,
        questions=questions,
    )
    session.pending_question = first["firstQuestion"]["text"]
    await db.commit()
    intake_history = await _build_intake_history(db, session.id)

    logger.info("Session created for patient %s: %s", patient.id, session.id)
    return JSONResponse(
        status_code=201,
        content={
            "sessionId": str(session.id),
            "status": "CREATED",
            "patient": {
                "dailyId": patient.daily_id,
                "name": patient.name,
                "age": patient.age,
                "gender": patient.gender,
                "expiresAt": patient.expires_at.isoformat(),
            },
            "intakeHistory": intake_history,
            "firstQuestion": first["firstQuestion"],
            "createdAt": session.created_at.isoformat(),
            "remainingQuestions": _remaining_questions_for([], questions),
            "questions": questions,
            "streamingEnabled": settings.CONVERSATION_STREAMING_ENABLED,
        }
    )


# ── GET /sessions/{id} ────────────────────────────────────────────────────────

@router.get("/{session_id}")
async def get_session(
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)
    lock_expires = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)

    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()

    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    session, patient = row

    intake_history = await _build_intake_history(db, session.id)

    responses_result = await db.execute(
        select(QuestionResponse, User.name.label("staff_name"))
        .join(User, QuestionResponse.responded_by_staff_id == User.id)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_response_rows = responses_result.all()
    all_responses = [r for r, _ in all_response_rows]

    responses = [
        {
            "questionId": r.question_id,
            "questionText": r.question_text,
            "answer": r.transcribed_text,
            "sequenceNumber": r.sequence_number,
            "inputType": r.input_type,
            "inputLanguage": r.input_language,
            "isFlagged": r.is_flagged,
            "createdAt": r.created_at.isoformat(),
            "respondedByStaffId": str(r.responded_by_staff_id),
            "respondedByStaffName": staff_name,
        }
        for r, staff_name in all_response_rows
        if r.transcribed_text != _IMPLICIT_MARKER
    ]

    system_prompt, _, questions = await _get_dept_context(db, session.department_id)
    summary_text = None
    next_question = None
    join_skipped_ids: set[str] = set()

    if session.status in (FIRST_ROUND_COMPLETE, COMPLETED):
        summary_result = await db.execute(
            select(PatientSummary).where(PatientSummary.session_id == session.id)
        )
        summary_row = summary_result.scalar_one_or_none()
        if summary_row:
            summary_text = summary_row.summary_text
    else:
        # Active session — acquire lock, register contributor, determine next question
        lock_result = await db.execute(
            select(PatientLock).where(PatientLock.patient_id == patient.id)
        )
        lock = lock_result.scalar_one_or_none()

        _ensure_lock_not_held_by_other_staff(lock, staff_id, now)

        if lock:
            lock.staff_id = staff_id
            lock.expires_at = lock_expires
        else:
            db.add(PatientLock(
                id=uuid.uuid4(),
                patient_id=patient.id,
                staff_id=staff_id,
                expires_at=lock_expires,
            ))

        active_result = await db.execute(
            select(ActiveSession).where(ActiveSession.staff_id == staff_id)
        )
        active_slot = active_result.scalar_one_or_none()
        if active_slot:
            active_slot.session_id = session.id
            active_slot.expires_at = lock_expires

        existing_contributor = await db.execute(
            select(SessionContributor).where(
                SessionContributor.session_id == session.id,
                SessionContributor.staff_id == staff_id,
            )
        )
        if not existing_contributor.scalar_one_or_none():
            count_result = await db.execute(
                select(func.count()).where(SessionContributor.session_id == session.id)
            )
            db.add(SessionContributor(
                id=uuid.uuid4(),
                session_id=session.id,
                staff_id=staff_id,
                sequence_number=count_result.scalar() + 1,
            ))

        await db.commit()
        patient_info = {"name": patient.name, "age": patient.age, "gender": patient.gender, "daily_id": patient.daily_id}
        if session.pending_question:
            next_question = {"id": None, "text": session.pending_question}
        else:
            try:
                rec_result = await db.execute(
                    select(PatientRecord)
                    .where(PatientRecord.session_id == session.id)
                    .order_by(PatientRecord.created_at)
                )
                session_records = [r.extracted_text for r in rec_result.scalars().all() if r.extracted_text]

                if all_responses:
                    conversation_so_far = [
                        {"question": r.question_text, "answer": r.transcribed_text}
                        for r in all_responses
                        if r.transcribed_text != _IMPLICIT_MARKER
                    ]
                    llm_result = await get_next_question(
                        patient_info=patient_info,
                        conversation=conversation_so_far,
                        latest_answer=next((r.transcribed_text for r in reversed(all_responses) if r.transcribed_text != _IMPLICIT_MARKER), ""),
                        system_prompt=system_prompt,
                        questions=questions,
                        patient_records=session_records or None,
                    )
                    join_skipped_ids = set(llm_result.get("skippedIds", []))
                    existing_ids = {str(r.question_id) for r in all_responses if r.question_id}
                    join_questions_map = {q["id"]: q["text"] for q in questions}
                    max_seq = max((r.sequence_number for r in all_responses), default=0)
                    for i, skipped_id in enumerate(sorted(join_skipped_ids - existing_ids)):
                        q_text = join_questions_map.get(skipped_id)
                        if q_text:
                            db.add(QuestionResponse(
                                id=uuid.uuid4(),
                                session_id=session.id,
                                question_id=skipped_id,
                                question_text=q_text,
                                transcribed_text=_IMPLICIT_MARKER,
                                input_type="TEXT",
                                input_language="ENGLISH",
                                is_flagged=False,
                                sequence_number=max_seq + i + 1,
                                responded_by_staff_id=staff_id,
                            ))
                    await db.commit()
                    if not llm_result["intakeComplete"]:
                        next_question = llm_result["nextQuestion"]
                else:
                    first = await get_first_question(
                        patient_info=patient_info,
                        system_prompt=system_prompt,
                        questions=questions,
                        patient_records=session_records or None,
                    )
                    next_question = first["firstQuestion"]
            except Exception as exc:
                exc_type = type(exc).__name__
                if "RateLimitError" in exc_type or "429" in str(exc):
                    raise HTTPException(
                        status_code=503,
                        detail={"error": error_codes.INTERNAL_ERROR, "message": "Service is not available. Please try again."},
                    )
                logger.error("LLM call failed in get_session: %s", str(exc))
                raise HTTPException(
                    status_code=500,
                    detail={"error": error_codes.INTERNAL_ERROR, "message": "Failed to determine next question."},
                )

    records_result = await db.execute(
        select(PatientRecord)
        .where(PatientRecord.session_id == session.id)
        .order_by(PatientRecord.created_at)
    )
    patient_records = [
        {
            "imageData": r.image_data,
            "mimeType": r.mime_type,
            "extractedText": r.extracted_text,
            "createdAt": r.created_at.isoformat(),
        }
        for r in records_result.scalars().all()
    ]

    return {
        "session": {
            "id": str(session.id),
            "status": session.status,
            "startedAt": session.started_at.isoformat(),
            "completedAt": session.completed_at.isoformat() if session.completed_at else None,
        },
        "patient": {
            "dailyId": patient.daily_id,
            "name": patient.name,
            "age": patient.age,
            "gender": patient.gender,
        },
        "intakeHistory": intake_history,
        "responses": responses,
        "summaryText": summary_text,
        "nextQuestion": next_question,
        "patientRecords": patient_records,
        "remainingQuestions": _remaining_questions_for(all_responses, questions, extra_excluded=join_skipped_ids),
        "streamingEnabled": settings.CONVERSATION_STREAMING_ENABLED,
    }


# ── GET /sessions/{id}/summary ───────────────────────────────────────────────

@router.get("/{session_id}/summary")
async def get_summary(
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    session_result = await db.execute(
        select(IntakeSession).where(IntakeSession.id == session_id)
    )
    session = session_result.scalar_one_or_none()
    if not session:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    summary_result = await db.execute(
        select(PatientSummary).where(PatientSummary.session_id == session.id)
    )
    summary = summary_result.scalar_one_or_none()
    if not summary:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SUMMARY_NOT_FOUND, "message": "Summary not found for this session"}
        )

    return {
        "sessionId": str(session.id),
        "summaryText": summary.summary_text,
        "additionalNote": summary.additional_note,
        "generationLatencyMs": summary.generation_latency_ms,
        "generatedAt": summary.generated_at.isoformat(),
    }


# ── POST /sessions/{id}/respond ───────────────────────────────────────────────

@router.post("/{session_id}/respond")
async def respond(
    session_id: uuid.UUID,
    body: RespondRequest,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)

    # Step 1 — Find session
    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()

    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    session, patient = row

    if session.status != ACTIVE:
        raise HTTPException(
            status_code=409,
            detail={"error": error_codes.SESSION_COMPLETED, "message": "Session is already completed"}
        )

    # Step 2 — Verify patient lock ownership
    lock_result = await db.execute(
        select(PatientLock).where(PatientLock.patient_id == patient.id)
    )
    lock = lock_result.scalar_one_or_none()
    if not lock or lock.expires_at <= now or lock.staff_id != staff_id:
        raise HTTPException(
            status_code=423,
            detail={"error": error_codes.SESSION_LOCK_LOST, "message": "Patient lock lost. Please re-acquire before responding."}
        )

    # Step 3 — Get current response count for sequence number
    count_result = await db.execute(
        select(func.count()).where(QuestionResponse.session_id == session.id)
    )
    response_count = count_result.scalar()

    # Step 4 — Resolve question text and ID (needed before answer extraction)
    system_prompt, _, questions = await _get_dept_context(db, session.department_id)
    _qmap = {q["id"]: q["text"] for q in questions}
    _text_to_id = {q["text"].strip().lower(): q["id"] for q in questions}
    question_text = body.questionText or _qmap.get(body.questionId or "") or body.questionId or ""
    question_id = body.questionId or _text_to_id.get(question_text.strip().lower())

    # Step 5 — For voice input, strip background noise before storing or advancing the LLM.
    # extract_patient_answer returns "" when the transcript contains no valid patient response.
    answer_text = body.text
    if body.inputType == "VOICE" and question_text:
        try:
            extracted = await extract_patient_answer(question_text, body.text)
        except Exception:
            extracted = body.text
        if not extracted:
            logger.info("No valid patient response in voice transcript; re-asking question.")
            no_answer_result = await db.execute(
                select(QuestionResponse)
                .where(QuestionResponse.session_id == session.id)
                .order_by(QuestionResponse.sequence_number)
            )
            return {
                "nextQuestion": {
                    "id": str(question_id) if question_id else None,
                    "text": session.pending_question or question_text,
                },
                "remainingQuestions": _remaining_questions_for(no_answer_result.scalars().all(), questions),
                "intakeComplete": False,
                "totalResponses": response_count,
                "isFlagged": False,
            }
        answer_text = extracted

    # Step 6 — Check if answer is flagged
    is_flagged = answer_text.strip().lower() in FLAGGED_PHRASES

    # Step 7 — Save Q&A to DB (idempotent)
    # If the LLM call (Step 9) times out and the client retries, the DB commit from the
    # first attempt already exists. Detect this by checking whether the latest saved
    # response matches this exact question + answer, and skip the INSERT if so.
    latest_resp_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number.desc())
        .limit(1)
    )
    latest_resp = latest_resp_result.scalar_one_or_none()
    already_saved = (
        latest_resp is not None
        and latest_resp.question_text.strip() == question_text.strip()
        and latest_resp.transcribed_text.strip() == answer_text.strip()
    )

    if not already_saved:
        db.add(QuestionResponse(
            id=uuid.uuid4(),
            session_id=session.id,
            question_id=question_id,
            question_text=question_text,
            transcribed_text=answer_text,
            input_type=body.inputType,
            input_language=body.inputLanguage,
            is_flagged=is_flagged,
            sequence_number=response_count + 1,
            responded_by_staff_id=staff_id,
        ))

    total_responses = response_count + (0 if already_saved else 1)

    # Step 7b — Refresh lock and session TTL
    new_expires = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)
    lock.expires_at = new_expires
    active_result = await db.execute(
        select(ActiveSession).where(ActiveSession.staff_id == staff_id)
    )
    active_slot = active_result.scalar_one_or_none()
    if active_slot:
        active_slot.expires_at = new_expires

    await db.commit()

    # Step 8 — Build conversation history for LLM
    history_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_responses = history_result.scalars().all()

    conversation = [
        {"question": r.question_text, "answer": r.transcribed_text}
        for r in all_responses
        if r.transcribed_text != _IMPLICIT_MARKER
    ]

    patient_info = {
        "name": patient.name,
        "age": patient.age,
        "gender": patient.gender,
        "daily_id": patient.daily_id,
    }

    # Step 9 — Call LLM for next question
    records_result = await db.execute(
        select(PatientRecord)
        .where(PatientRecord.session_id == session.id)
        .order_by(PatientRecord.created_at)
    )
    patient_records = [r.extracted_text for r in records_result.scalars().all() if r.extracted_text]

    llm_result = await get_next_question(
        patient_info=patient_info,
        conversation=conversation,
        latest_answer=answer_text,
        system_prompt=system_prompt,
        questions=questions,
        patient_records=patient_records or None,
        language=body.inputLanguage,
    )

    # Step 9b — Persist implicitly-answered question records so they are excluded
    # from remainingQuestions in this response and in complete_session later.
    skipped_ids = set(llm_result.get("skippedIds", []))
    for i, skipped_id in enumerate(sorted(skipped_ids)):
        q_text = _qmap.get(skipped_id)
        if q_text:
            db.add(QuestionResponse(
                id=uuid.uuid4(),
                session_id=session.id,
                question_id=skipped_id,
                question_text=q_text,
                transcribed_text=_IMPLICIT_MARKER,
                input_type="TEXT",
                input_language=body.inputLanguage,
                is_flagged=False,
                sequence_number=total_responses + i + 1,
                responded_by_staff_id=staff_id,
            ))

    # Step 9c — Fallback: if the LLM returned an empty question text (e.g. it output only a
    # SKIPPED line with no question after it), use the first remaining unanswered question so
    # the session doesn't stall silently on the frontend.
    if not llm_result["intakeComplete"] and not (llm_result["nextQuestion"] or {}).get("text"):
        remaining = _remaining_questions_for(all_responses, questions, extra_excluded=skipped_ids)
        if remaining:
            llm_result["nextQuestion"] = {"id": remaining[0]["id"], "text": remaining[0]["text"]}
        else:
            llm_result["intakeComplete"] = True
            llm_result["nextQuestion"] = None

    # Step 10 — Cache next question / handle intake complete
    if llm_result["intakeComplete"]:
        session.pending_question = None
        await db.commit()

        flagged_result = await db.execute(
            select(QuestionResponse)
            .where(
                QuestionResponse.session_id == session.id,
                QuestionResponse.is_flagged == True,
            )
            .order_by(QuestionResponse.sequence_number)
        )
        flagged = flagged_result.scalars().all()

        return {
            "nextQuestion": None,
            "intakeComplete": True,
            "totalResponses": total_responses,
            "flaggedQuestions": [
                {
                    "questionText": r.question_text,
                    "answerText": r.transcribed_text,
                    "sequenceNumber": r.sequence_number,
                }
                for r in flagged
                if r.transcribed_text != _IMPLICIT_MARKER
            ],
            "remainingQuestions": _remaining_questions_for(all_responses, questions, extra_excluded=skipped_ids),
        }

    session.pending_question = llm_result["nextQuestion"]["text"]
    await db.commit()

    return {
        "nextQuestion": llm_result["nextQuestion"],
        "remainingQuestions": _remaining_questions_for(all_responses, questions, extra_excluded=skipped_ids),
        "intakeComplete": False,
        "totalResponses": total_responses,
        "isFlagged": is_flagged,
    }


# ── POST /sessions/{id}/complete ──────────────────────────────────────────────

@router.post("/{session_id}/complete")
async def complete_session(
    session_id: uuid.UUID,
    body: CompleteSessionRequest = CompleteSessionRequest(),
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])

    # Step 1 — Find session
    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()

    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    session, patient = row
    role = staff.get("role")

    if session.status == COMPLETED:
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.SESSION_COMPLETED, "message": "Session is already completed"}
        )

    if session.status == FIRST_ROUND_COMPLETE and role not in REVIEWER_ROLE_KEYS:
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.SESSION_NOT_ACTIVE, "message": "Session is awaiting senior consultant review"}
        )

    if session.status not in (ACTIVE, FIRST_ROUND_COMPLETE):
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.SESSION_NOT_ACTIVE, "message": "Session is not active"}
        )

    # Step 2 — Load all responses
    history_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_responses = history_result.scalars().all()

    if not all_responses:
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.NO_RESPONSES, "message": "Cannot complete a session with no responses"}
        )

    conversation = [
        {"question": r.question_text, "answer": r.transcribed_text}
        for r in all_responses
        if r.transcribed_text != _IMPLICIT_MARKER
    ]

    records_result = await db.execute(
        select(PatientRecord)
        .where(PatientRecord.session_id == session.id)
        .order_by(PatientRecord.created_at)
    )
    patient_records = [r.extracted_text for r in records_result.scalars().all() if r.extracted_text]

    patient_info = {
        "name": patient.name,
        "age": patient.age,
        "gender": patient.gender,
        "daily_id": patient.daily_id,
    }

    # Step 3 — Mark COMPLETED (sr.consultant+) or FIRST_ROUND_COMPLETE (consultant)
    now = datetime.now(timezone.utc)
    new_status = COMPLETED if role in REVIEWER_ROLE_KEYS else FIRST_ROUND_COMPLETE
    session.status = new_status
    session.completed_at = now
    await db.flush()

    # Step 4 — Generate summary
    _, summary_prompt, questions = await _get_dept_context(db, session.department_id)
    start_time = datetime.now(timezone.utc)
    summary_text = await generate_summary(
        patient_info=patient_info,
        conversation=conversation,
        intake_date=str(patient.intake_date),
        summary_prompt=summary_prompt,
        patient_records=patient_records or None,
    )
    latency_ms = int((datetime.now(timezone.utc) - start_time).total_seconds() * 1000)

    # Step 5 — Save summary (upsert in case session was re-completed)
    existing_summary = await db.execute(
        select(PatientSummary).where(PatientSummary.session_id == session.id)
    )
    existing = existing_summary.scalar_one_or_none()
    if existing:
        existing.summary_text = summary_text
        existing.generation_latency_ms = latency_ms
        existing.additional_note = body.additionalNote
    else:
        db.add(PatientSummary(
            id=uuid.uuid4(),
            session_id=session.id,
            summary_text=summary_text,
            additional_note=body.additionalNote,
            generation_latency_ms=latency_ms,
        ))

    # Step 6 — Release patient lock and active session slot
    lock_result = await db.execute(
        select(PatientLock).where(PatientLock.patient_id == patient.id)
    )
    lock = lock_result.scalar_one_or_none()
    if lock:
        await db.delete(lock)

    active_result = await db.execute(
        select(ActiveSession).where(ActiveSession.staff_id == staff_id)
    )
    active_slot = active_result.scalar_one_or_none()
    if active_slot:
        await db.delete(active_slot)

    await db.commit()

    logger.info("Session completed: %s", session_id)

    return {
        "summaryId": str(session.id),
        "status": new_status.lower(),
        "summaryText": summary_text,
        "additionalNote": body.additionalNote,
        "generatedAt": now.isoformat(),
        "latencyMs": latency_ms,
        "remainingQuestions": _remaining_questions_for(all_responses, questions),
    }


# ── PATCH /sessions/{id}/note ─────────────────────────────────────────────────

class UpdateNoteRequest(BaseModel):
    additionalNote: Optional[str] = None

    @field_validator("additionalNote")
    @classmethod
    def strip_note(cls, v: Optional[str]) -> Optional[str]:
        if v is not None:
            v = v.strip()
            if len(v) > MAX_NOTE_LENGTH:
                raise ValueError(f"additionalNote max {MAX_NOTE_LENGTH} chars")
            return v if v else None
        return v


@router.patch("/{session_id}/note")
async def update_session_note(
    session_id: uuid.UUID,
    body: UpdateNoteRequest,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    result = await db.execute(
        select(PatientSummary).where(PatientSummary.session_id == session_id)
    )
    summary = result.scalar_one_or_none()
    if not summary:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SUMMARY_NOT_FOUND, "message": "Session summary not found"}
        )

    summary.additional_note = body.additionalNote
    await db.commit()
    return {"additionalNote": body.additionalNote}


# ── POST /sessions/{id}/ask-more ──────────────────────────────────────────────

@router.post("/{session_id}/ask-more")
async def ask_more(
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)
    lock_expires = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)

    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    session, patient = row

    if session.status not in (FIRST_ROUND_COMPLETE, COMPLETED):
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.SESSION_NOT_COMPLETED, "message": "Session is not completed"}
        )

    # Check concurrent session limit before re-opening
    count_result = await db.execute(
        select(func.count()).select_from(ActiveSession).where(
            ActiveSession.expires_at > now
        )
    )
    active_count = count_result.scalar()
    if active_count >= settings.MAX_CONCURRENT_SESSIONS:
        raise HTTPException(
            status_code=503,
            detail={"error": error_codes.MAX_SESSIONS_REACHED, "message": "Maximum concurrent sessions reached. Try again later."}
        )

    # Re-open session
    session.status = ACTIVE
    session.completed_at = None
    await db.flush()

    # Add contributor if new
    existing_contributor = await db.execute(
        select(SessionContributor).where(
            SessionContributor.session_id == session.id,
            SessionContributor.staff_id == staff_id,
        )
    )
    if not existing_contributor.scalar_one_or_none():
        count_result = await db.execute(
            select(func.count()).where(SessionContributor.session_id == session.id)
        )
        db.add(SessionContributor(
            id=uuid.uuid4(),
            session_id=session.id,
            staff_id=staff_id,
            sequence_number=count_result.scalar() + 1,
        ))

    # UPSERT patient lock
    lock_result = await db.execute(
        select(PatientLock).where(PatientLock.patient_id == patient.id)
    )
    lock = lock_result.scalar_one_or_none()
    if lock:
        lock.staff_id = staff_id
        lock.expires_at = lock_expires
    else:
        db.add(PatientLock(
            id=uuid.uuid4(),
            patient_id=patient.id,
            staff_id=staff_id,
            expires_at=lock_expires,
        ))

    # Update active_session TTL and link to re-opened session
    active_result = await db.execute(
        select(ActiveSession).where(ActiveSession.staff_id == staff_id)
    )
    active_slot = active_result.scalar_one_or_none()
    if active_slot:
        active_slot.session_id = session.id
        active_slot.expires_at = lock_expires

    await db.commit()

    history_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_responses = history_result.scalars().all()

    conversation_history = [
        {"questionId": r.question_id, "questionText": r.question_text, "answer": r.transcribed_text,
         "sequenceNumber": r.sequence_number, "isFlagged": r.is_flagged}
        for r in all_responses
        if r.transcribed_text != _IMPLICIT_MARKER
    ]
    _, _, questions = await _get_dept_context(db, session.department_id)

    return {
        "conversationHistory": conversation_history,
        "remainingQuestions": _remaining_questions_for(all_responses, questions),
    }


# ── POST /sessions/{id}/ask-more/next-question ────────────────────────────────

@router.post("/{session_id}/ask-more/next-question")
async def ask_more_next_question(
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    session, patient = row

    if session.status != ACTIVE:
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.SESSION_NOT_ACTIVE, "message": "Session is not active"}
        )

    history_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_responses = history_result.scalars().all()

    system_prompt, _, questions = await _get_dept_context(db, session.department_id)
    conversation_for_llm = [
        {"question": r.question_text, "answer": r.transcribed_text}
        for r in all_responses
        if r.transcribed_text != _IMPLICIT_MARKER
    ]

    records_result = await db.execute(
        select(PatientRecord)
        .where(PatientRecord.session_id == session.id)
        .order_by(PatientRecord.created_at)
    )
    patient_records = [r.extracted_text for r in records_result.scalars().all() if r.extracted_text]

    llm_result = await get_next_question(
        patient_info={"name": patient.name, "age": patient.age, "gender": patient.gender, "daily_id": patient.daily_id},
        conversation=conversation_for_llm,
        latest_answer=all_responses[-1].transcribed_text if all_responses else "",
        system_prompt=system_prompt,
        questions=questions,
        patient_records=patient_records or None,
    )

    next_question = None if llm_result["intakeComplete"] else llm_result["nextQuestion"]
    return {"nextQuestion": next_question}


# ── POST /sessions/{id}/current-question ─────────────────────────────────────

class CurrentQuestionRequest(BaseModel):
    language: str = "te-IN"

    @field_validator("language")
    @classmethod
    def validate_language(cls, v):
        if not re.match(r'^[a-z]{2}-[A-Z]{2}$', v):
            raise ValueError("language must be a valid language code (e.g. te-IN, hi-IN, kn-IN, ta-IN)")
        return v


@router.post("/{session_id}/current-question")
async def get_current_question(
    session_id: uuid.UUID,
    body: CurrentQuestionRequest,
    db: AsyncSession = Depends(get_db),
    current_staff: dict = Depends(get_current_staff),
):
    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        raise HTTPException(status_code=404, detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"})
    session, patient = row

    history_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session.id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_responses = history_result.scalars().all()

    patient_info = {"name": patient.name, "age": patient.age, "gender": patient.gender, "daily_id": patient.daily_id}
    system_prompt, _, questions = await _get_dept_context(db, session.department_id)

    cur_q_records_result = await db.execute(
        select(PatientRecord)
        .where(PatientRecord.session_id == session.id)
        .order_by(PatientRecord.created_at)
    )
    cur_q_records = [r.extracted_text for r in cur_q_records_result.scalars().all() if r.extracted_text]

    try:
        if all_responses:
            conversation = [
                {"question": r.question_text, "answer": r.transcribed_text}
                for r in all_responses
                if r.transcribed_text != _IMPLICIT_MARKER
            ]
            llm_result = await get_next_question(
                patient_info=patient_info,
                conversation=conversation,
                latest_answer=next((r.transcribed_text for r in reversed(all_responses) if r.transcribed_text != _IMPLICIT_MARKER), ""),
                system_prompt=system_prompt,
                questions=questions,
                patient_records=cur_q_records or None,
                language=body.language,
            )
            question_text = llm_result.get("nextQuestion") or ""
        else:
            first = await get_first_question(
                patient_info=patient_info,
                system_prompt=system_prompt,
                questions=questions,
                patient_records=cur_q_records or None,
                language=body.language,
            )
            question_text = first["firstQuestion"]["text"]
    except Exception as exc:
        logger.error("LLM call failed in current-question: %s", str(exc))
        raise HTTPException(status_code=500, detail={"error": error_codes.INTERNAL_ERROR, "message": "Failed to generate question"})

    return {"question": question_text}


# ── POST /sessions/{id}/extract-record ────────────────────────────────────────

@router.post("/{session_id}/extract-record")
async def extract_patient_record(
    session_id: uuid.UUID,
    file: UploadFile = File(...),
    db: AsyncSession = Depends(get_db),
    current_staff: dict = Depends(get_current_staff),
):
    result = await db.execute(
        select(IntakeSession).where(IntakeSession.id == session_id)
    )
    session = result.scalar_one_or_none()
    if not session:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )
    if session.status in (FIRST_ROUND_COMPLETE, COMPLETED):
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.SESSION_COMPLETED, "message": "Cannot upload records to a completed session"}
        )

    contents = await file.read()
    base64_image = base64.b64encode(contents).decode("utf-8")
    mime_type = file.content_type or "image/jpeg"

    try:
        extracted_text = await extract_from_image(base64_image, mime_type)
    except Exception as exc:
        exc_type = type(exc).__name__
        if "RateLimitError" in exc_type or "429" in str(exc):
            raise HTTPException(
                status_code=503,
                detail={"error": error_codes.INTERNAL_ERROR, "message": "Service is not available. Please try again."},
            )
        logger.error("Image extraction failed: %s", str(exc))
        raise HTTPException(
            status_code=500,
            detail={"error": error_codes.INTERNAL_ERROR, "message": "Failed to extract data from image"}
        )

    try:
        db.add(PatientRecord(
            id=uuid.uuid4(),
            session_id=session_id,
            image_data=base64_image,
            mime_type=mime_type,
            extracted_text=extracted_text,
        ))
        await db.commit()
    except Exception as exc:
        logger.error("Failed to save patient record: %s", str(exc))
        raise HTTPException(
            status_code=500,
            detail={"error": error_codes.INTERNAL_ERROR, "message": "Failed to save patient record"}
        )

    next_question = None
    try:
        pt_result = await db.execute(select(Patient).where(Patient.id == session.patient_id))
        pt = pt_result.scalar_one_or_none()

        resp_result = await db.execute(
            select(QuestionResponse)
            .where(QuestionResponse.session_id == session.id)
            .order_by(QuestionResponse.sequence_number)
        )
        responses = resp_result.scalars().all()

        rec_result = await db.execute(
            select(PatientRecord)
            .where(PatientRecord.session_id == session.id)
            .order_by(PatientRecord.created_at)
        )
        all_records = [r.extracted_text for r in rec_result.scalars().all() if r.extracted_text]

        system_prompt, _, questions = await _get_dept_context(db, session.department_id)
        pt_info = {
            "name": pt.name if pt else "",
            "age": pt.age if pt else "",
            "gender": pt.gender if pt else "",
            "daily_id": pt.daily_id if pt else "",
        }

        if responses:
            conversation = [
                {"question": r.question_text, "answer": r.transcribed_text}
                for r in responses
                if r.transcribed_text != _IMPLICIT_MARKER
            ]
            llm_result = await get_next_question(
                patient_info=pt_info,
                conversation=conversation,
                latest_answer=next((r.transcribed_text for r in reversed(responses) if r.transcribed_text != _IMPLICIT_MARKER), ""),
                system_prompt=system_prompt,
                questions=questions,
                patient_records=all_records or None,
            )
            if not llm_result["intakeComplete"]:
                next_question = llm_result["nextQuestion"]
        else:
            first = await get_first_question(
                patient_info=pt_info,
                system_prompt=system_prompt,
                questions=questions,
                patient_records=all_records or None,
            )
            next_question = first["firstQuestion"]
    except Exception as exc:
        logger.error("Failed to get next question after record extraction: %s", str(exc))

    logger.info("Patient record extracted for session: %s", session_id)
    return {"extractedText": extracted_text, "nextQuestion": next_question}


# ── POST /sessions/{id}/heartbeat ────────────────────────────────────────────

@router.post("/{session_id}/heartbeat")
async def heartbeat_session(
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)
    lock_expires = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)

    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        return {"ok": False}
    session, patient = row
    if session.status != ACTIVE:
        return {"ok": False}

    lock_result = await db.execute(select(PatientLock).where(PatientLock.patient_id == patient.id))
    lock = lock_result.scalar_one_or_none()
    if lock:
        lock.staff_id = staff_id
        lock.expires_at = lock_expires
    else:
        db.add(PatientLock(id=uuid.uuid4(), patient_id=patient.id, staff_id=staff_id, expires_at=lock_expires))

    active_result = await db.execute(select(ActiveSession).where(ActiveSession.staff_id == staff_id))
    active_slot = active_result.scalar_one_or_none()
    if active_slot:
        active_slot.expires_at = lock_expires

    await db.commit()
    return {"ok": True}


# ── POST /sessions/{id}/transcribe ────────────────────────────────────────────

@router.post("/{session_id}/transcribe")
async def transcribe_session_audio(
    session_id: uuid.UUID,
    file: UploadFile = File(...),
    language: str = Form("te-IN"),
    question: str = Form(""),
    db: AsyncSession = Depends(get_db),
    current_staff: dict = Depends(get_current_staff),
):
    audio_bytes = await file.read()
    t_start = datetime.now(timezone.utc)
    logger.warning("Transcribe request — language: %s, file size: %d bytes", language, len(audio_bytes))
    last_exc: Exception | None = None
    transcript = ""
    detected_lang = language
    for attempt in range(1, MAX_TRANSCRIPTION_ATTEMPTS + 1):
        try:
            t_stt_start = datetime.now(timezone.utc)
            transcript, detected_lang = await sarvam_transcribe(
                audio_bytes,
                file.filename or "audio.wav",
                file.content_type or "audio/wav",
                language,
            )
            t_stt_ms = int((datetime.now(timezone.utc) - t_stt_start).total_seconds() * 1000)
            logger.warning("[Latency] Sarvam STT attempt %d: %dms | transcript: %r", attempt, t_stt_ms, transcript)
            last_exc = None
            break
        except Exception as exc:
            body = getattr(getattr(exc, "response", None), "text", None)
            logger.error("Sarvam STT attempt %d failed: %s | response body: %s", attempt, str(exc), body, exc_info=True)
            last_exc = exc

    if last_exc is not None:
        raise HTTPException(
            status_code=500,
            detail={"error": error_codes.INTERNAL_ERROR, "message": "Transcription failed"}
        )

    try:
        if transcript and transcript.count('\n') >= 2:
            logger.warning("Multi-speaker filtered on raw transcript: %r", transcript)
            return {"transcript": ""}

        if transcript and detected_lang and detected_lang != "en-IN":
            t_translate_start = datetime.now(timezone.utc)
            transcript = await sarvam_translate(transcript, detected_lang)
            t_translate_ms = int((datetime.now(timezone.utc) - t_translate_start).total_seconds() * 1000)
            logger.warning("[Latency] Sarvam Translate: %dms | result: %r", t_translate_ms, transcript)
    except Exception as exc:
        body = getattr(getattr(exc, "response", None), "text", None)
        logger.error("Sarvam translation failed: %s | response body: %s", str(exc), body, exc_info=True)
        raise HTTPException(
            status_code=500,
            detail={"error": error_codes.INTERNAL_ERROR, "message": "Transcription failed"}
        )
    t_total_ms = int((datetime.now(timezone.utc) - t_start).total_seconds() * 1000)
    logger.warning("[Latency] Total transcribe pipeline: %dms", t_total_ms)
    return {"transcript": transcript}


# ── Request schema: /conversation-segment/classify ────────────────────────────

class ClassifyUtteranceRef(BaseModel):
    utteranceId: str
    speakerId: str

    @field_validator("utteranceId")
    @classmethod
    def validate_utterance_id(cls, v):
        v = v.strip()
        try:
            uuid.UUID(v)
        except ValueError:
            raise ValueError("utteranceId must be a valid UUID")
        return v


class ClassifyConversationSegmentRequest(BaseModel):
    utterances: list[ClassifyUtteranceRef]

    @field_validator("utterances")
    @classmethod
    def utterances_not_empty(cls, v):
        if not v:
            raise ValueError("utterances cannot be empty")
        return v


# ── POST /sessions/{id}/conversation-segment/transcribe ──────────────────────
# Free-conversation intake mode: the doctor and patient speak naturally (no
# turn-taking Q/A loop). The frontend VAD-segments the shared mic feed and posts
# each segment here first — the fast leg of the pipeline (Sarvam STT + translation
# only) — so the raw transcript can be shown in the UI immediately. Each diarized
# entry is persisted as a PENDING placeholder utterance (claiming its sequence
# number right away) and handed back to the frontend, which then calls
# /conversation-segment/classify below to run the slower Gemini classification/
# mapping step without blocking on it first.

async def _persist_conversation_entries(
    db: AsyncSession, session: IntakeSession, lock: PatientLock, now: datetime,
    entries: list[dict], detected_lang: str,
) -> dict:
    """Persists STT entries (list of {"speakerId", "text"}) as PENDING placeholder
    utterances, claiming their sequence numbers. Shared by the batch
    /conversation-segment/transcribe endpoint below and the real-time
    conversation-stream WebSocket — the two differ only in how `entries` were produced
    (a Sarvam batch job vs. a live streaming connection), not in how they're persisted."""
    lock.expires_at = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)

    if not entries:
        _, _, questions = await _get_dept_context(db, session.department_id)
        history_result = await db.execute(
            select(QuestionResponse).where(QuestionResponse.session_id == session.id)
        )
        current_responses = history_result.scalars().all()
        await db.commit()
        return {
            "utterances": [],
            "remainingQuestions": _remaining_questions_for(current_responses, questions),
        }

    # Serialize sequence-number assignment against any other transcribe/classify call for
    # this session — see the matching lock + comment in _classify_and_persist_utterances.
    # Acquired only now, after the slow STT step, so two segments' STT calls can still run
    # fully in parallel; this only serializes the brief count-then-insert section below.
    await db.execute(
        select(IntakeSession.id).where(IntakeSession.id == session.id).with_for_update()
    )

    utterance_count_result = await db.execute(
        select(func.count()).where(ConversationUtterance.session_id == session.id)
    )
    utterance_count = utterance_count_result.scalar()

    # Claim a sequence number and persist the raw transcript for every diarized entry right
    # away, as PENDING placeholders — this is what lets the frontend show the transcript the
    # instant Sarvam returns, well before the slower Gemini classify step (see
    # _classify_and_persist_utterances below) fills in speaker role and question mapping
    # for these same rows.
    placeholder_rows: list[tuple[ConversationUtterance, str]] = []
    for i, entry in enumerate(entries):
        placeholder = ConversationUtterance(
            id=uuid.uuid4(),
            session_id=session.id,
            sequence_number=utterance_count + i + 1,
            speaker_role="UNKNOWN",
            raw_transcript=entry["text"],
            language=detected_lang,
            mapping_status="PENDING",
        )
        db.add(placeholder)
        placeholder_rows.append((placeholder, entry["speakerId"]))

    await db.commit()

    return {
        "utterances": [
            {
                "utteranceId": str(placeholder.id),
                "sequenceNumber": placeholder.sequence_number,
                "speakerId": speaker_id,
                "text": placeholder.raw_transcript,
            }
            for placeholder, speaker_id in placeholder_rows
        ],
    }


@router.post("/{session_id}/conversation-segment/transcribe")
async def transcribe_conversation_segment(
    session_id: uuid.UUID,
    file: UploadFile = File(...),
    language: str = Form("te-IN"),
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)

    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )
    session, patient = row

    if session.status != ACTIVE:
        raise HTTPException(
            status_code=409,
            detail={"error": error_codes.SESSION_COMPLETED, "message": "Session is already completed"}
        )

    lock_result = await db.execute(
        select(PatientLock).where(PatientLock.patient_id == patient.id)
    )
    lock = lock_result.scalar_one_or_none()
    if not lock or lock.expires_at <= now or lock.staff_id != staff_id:
        raise HTTPException(
            status_code=423,
            detail={"error": error_codes.SESSION_LOCK_LOST, "message": "Patient lock lost. Please re-acquire before responding."}
        )

    audio_bytes = await file.read()
    try:
        if settings.CONVERSATION_DIARIZATION_ENABLED:
            entries, detected_lang = await sarvam_transcribe_diarized(
                audio_bytes, file.filename or "segment.wav", file.content_type or "audio/wav", language,
            )
        else:
            transcript, detected_lang = await sarvam_transcribe(
                audio_bytes, file.filename or "segment.wav", file.content_type or "audio/wav", language,
            )
            entries = [{"speakerId": "0", "text": transcript}] if transcript.strip() else []
        if detected_lang and detected_lang != "en-IN":
            for entry in entries:
                entry["text"] = await sarvam_translate(entry["text"], detected_lang)
    except Exception as exc:
        logger.error("Conversation segment STT failed: %s", str(exc), exc_info=True)
        raise HTTPException(
            status_code=500,
            detail={"error": error_codes.INTERNAL_ERROR, "message": "Transcription failed"}
        )

    return await _persist_conversation_entries(db, session, lock, now, entries, detected_lang)


# ── POST /sessions/{id}/conversation-segment/classify ─────────────────────────
# Slow leg of free-conversation capture: classifies + maps the PENDING placeholder
# utterances handed back by /transcribe above via Gemini, then — when a patient
# answer maps confidently to a standard question — saves it as a QuestionResponse
# row exactly like the structured /respond flow, so remainingQuestions / summary /
# edit endpoints all keep working unchanged regardless of which mode produced the
# answer.

@router.post("/{session_id}/conversation-segment/classify")
async def classify_conversation_segment(
    session_id: uuid.UUID,
    body: ClassifyConversationSegmentRequest,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)

    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )
    session, patient = row

    if session.status != ACTIVE:
        raise HTTPException(
            status_code=409,
            detail={"error": error_codes.SESSION_COMPLETED, "message": "Session is already completed"}
        )

    lock_result = await db.execute(
        select(PatientLock).where(PatientLock.patient_id == patient.id)
    )
    lock = lock_result.scalar_one_or_none()
    if not lock or lock.expires_at <= now or lock.staff_id != staff_id:
        raise HTTPException(
            status_code=423,
            detail={"error": error_codes.SESSION_LOCK_LOST, "message": "Patient lock lost. Please re-acquire before responding."}
        )

    utterance_refs = [{"utteranceId": u.utteranceId, "speakerId": u.speakerId} for u in body.utterances]
    return await _classify_and_persist_utterances(db, session, staff_id, lock, now, utterance_refs)


async def _classify_and_persist_utterances(
    db: AsyncSession, session: IntakeSession, staff_id: uuid.UUID, lock: PatientLock,
    now: datetime, utterance_refs: list[dict],
) -> dict:
    """utterance_refs: list of {"utteranceId": str, "speakerId": str}. Shared by the batch
    /conversation-segment/classify endpoint above and the conversation-stream WebSocket —
    once an utterance is a PENDING placeholder row, classification doesn't care whether the
    audio behind it came from a batch job or a live stream."""
    utterance_ids = [uuid.UUID(u["utteranceId"]) for u in utterance_refs]
    speaker_ids_by_id = {u["utteranceId"]: u["speakerId"] for u in utterance_refs}

    pending_result = await db.execute(
        select(ConversationUtterance).where(
            ConversationUtterance.id.in_(utterance_ids),
            ConversationUtterance.session_id == session.id,
            ConversationUtterance.mapping_status == "PENDING",
        )
    )
    pending_by_id = {row.id: row for row in pending_result.scalars().all()}
    if len(pending_by_id) != len(utterance_ids):
        raise HTTPException(
            status_code=400,
            detail={"error": error_codes.VALIDATION_ERROR, "message": "One or more utterances were not found or already classified"}
        )
    # Preserve the order the client sent them in — this is the original Sarvam entry order,
    # which classify_and_map_utterances/segment_and_map_utterance depend on for alignment.
    pending_rows = [pending_by_id[uid] for uid in utterance_ids]
    entries = [{"speakerId": speaker_ids_by_id[str(row.id)], "text": row.raw_transcript} for row in pending_rows]
    detected_lang = pending_rows[0].language

    _, _, questions = await _get_dept_context(db, session.department_id)

    min_sequence = pending_rows[0].sequence_number
    recent_result = await db.execute(
        select(ConversationUtterance)
        .where(
            ConversationUtterance.session_id == session.id,
            ConversationUtterance.sequence_number < min_sequence,
            ConversationUtterance.mapping_status != "PENDING",
        )
        .order_by(ConversationUtterance.sequence_number.desc())
        .limit(CONVERSATION_RECENT_CONTEXT_SIZE)
    )
    recent_rows = recent_result.scalars().all()
    recent_context = [
        {"speakerRole": r.speaker_role, "text": r.raw_transcript}
        for r in reversed(recent_rows)
    ]

    # VAD cuts a segment on any silence gap, so one continuous speaking turn — a patient
    # describing several aspects of the same symptom in separate breaths, e.g. "It is a
    # pressure like pain." / "It sometimes spreads to my left arm." / "It lasts about 10
    # minutes." — routinely arrives as several independent segments even though nothing
    # else was said in between. Whether text alone LOOKS "incomplete" is not a reliable
    # signal here (each fragment above reads as a grammatically complete sentence on its
    # own) — so the merge trigger is structural: same speaker still talking, no other
    # speaker's turn in between. merge_target stays the live ORM row so it can be updated
    # in place (grown and re-classified) instead of inserted as a disconnected new row.
    merge_target = recent_rows[0] if recent_rows else None
    previous_utterance = (
        {"speakerRole": merge_target.speaker_role, "text": merge_target.raw_transcript}
        if merge_target else None
    )

    qmap = {q["id"]: q["text"] for q in questions}

    # Custom-question and QuestionResponse rows are added after every ConversationUtterance
    # row below is flushed, since they carry a foreign key to the utterance and have no ORM
    # relationship() to order the flush for us.
    pending_custom_questions = []
    pending_responses = []
    utterances_out = []

    # Both branches converge on the same shape — a list of {"speakerRole", "text", "mappings",
    # "customQuestion"} turns, in chronological order — so one shared loop below saves the DB
    # rows and builds the response regardless of which path produced the turns:
    #   - diarized: Sarvam already split the audio into per-speaker entries; classify each.
    #   - non-diarized: Sarvam returns one flat transcript; the LLM splits it into turns by
    #     text alone (segment_and_map_utterance) instead of relying on audio diarization.
    segments = None
    max_attempts = settings.CLASSIFY_MAX_ATTEMPTS
    for attempt in range(1, max_attempts + 1):
        try:
            if settings.CONVERSATION_DIARIZATION_ENABLED:
                classifications = await classify_and_map_utterances(entries, recent_context, questions, previous_utterance)
                segments = [
                    {
                        "speakerRole": c["speakerRole"], "text": e["text"], "mappings": c["mappings"],
                        "customQuestion": c["customQuestion"], "continuesFromPrevious": c["continuesFromPrevious"],
                    }
                    for e, c in zip(entries, classifications)
                ]
            else:
                segments = await segment_and_map_utterance(entries[0]["text"], recent_context, questions, previous_utterance)
            break
        except Exception as exc:
            exc_type = type(exc).__name__
            if attempt < max_attempts:
                logger.warning(
                    "Conversation classify attempt %d/%d failed, retrying: %s",
                    attempt, max_attempts, str(exc),
                )
                await asyncio.sleep(_classify_retry_backoff_s(attempt))
                continue
            if "RateLimitError" in exc_type or "429" in str(exc):
                raise HTTPException(
                    status_code=503,
                    detail={"error": error_codes.INTERNAL_ERROR, "message": "Service is not available. Please try again."},
                )
            logger.error(
                "Conversation utterance classification failed after %d attempts: %s",
                max_attempts, str(exc), exc_info=True,
            )
            raise HTTPException(
                status_code=500,
                detail={"error": error_codes.INTERNAL_ERROR, "message": "Failed to process conversation segment"}
            )

    # Serialize sequence-number assignment against any other transcribe/classify call for
    # this session — a plain count() query race here (e.g. this classify's "extra" turns vs.
    # the next segment's /transcribe placeholders, both reading the same stale count) can
    # hand out duplicate sequence numbers. Acquired only now, after the slow Gemini call
    # above, so two segments' classify calls don't serialize on network time — just on this
    # short DB section. Released automatically at commit.
    await db.execute(
        select(IntakeSession.id).where(IntakeSession.id == session.id).with_for_update()
    )

    response_count_result = await db.execute(
        select(func.count()).where(QuestionResponse.session_id == session.id)
    )
    response_count = response_count_result.scalar()

    # Running counter for any turn beyond what the transcribe-time placeholders cover — only
    # possible in non-diarized mode, where the LLM can split one flat transcript into more
    # turns than the single placeholder it arrived as (see the loop below).
    extra_count_result = await db.execute(
        select(func.count()).where(ConversationUtterance.session_id == session.id)
    )
    next_extra_sequence = extra_count_result.scalar() + 1

    # Merge trigger is structural — same speaker still talking, nobody else has taken a turn
    # since — not the classifier's own "continuesFromPrevious" guess, which text alone can't
    # answer reliably (see the comment above merge_target). Update that existing row in place
    # instead of inserting a disconnected new one — the live transcript then shows one clean,
    # growing bubble rather than a fragment per VAD cut. Reusing its sequence_number in the
    # response signals the frontend to replace that entry rather than append a duplicate.
    should_merge = (
        merge_target is not None
        and bool(segments)
        and segments[0]["speakerRole"] == merge_target.speaker_role
    )

    if should_merge:
        seg = segments[0]
        confident_mappings = [
            m for m in seg["mappings"]
            if m["confidence"] >= CONVERSATION_MAPPING_CONFIDENCE_THRESHOLD
        ]
        primary_mapping = confident_mappings[0] if confident_mappings else None
        mapping_status = "MAPPED" if confident_mappings else ("SMALL_TALK" if seg["speakerRole"] == "SMALL_TALK" else "UNMAPPED")

        # The growing utterance may have already produced a QuestionResponse in an earlier
        # round (e.g. a vague first fragment mapped weakly) — re-classifying the fuller,
        # merged text can shift or drop that mapping. Clear it before writing this round's
        # mappings so a fuller answer replaces the stale one instead of duplicating it.
        old_question_id = merge_target.mapped_question_id
        if old_question_id:
            await db.execute(
                delete(QuestionResponse).where(
                    QuestionResponse.session_id == session.id,
                    QuestionResponse.question_id == old_question_id,
                )
            )

        merged_text = f"{merge_target.raw_transcript.rstrip()} {seg['text'].lstrip()}".strip()
        merge_target.raw_transcript = merged_text
        merge_target.mapped_question_id = primary_mapping["questionId"] if primary_mapping else None
        merge_target.mapping_confidence = primary_mapping["confidence"] if primary_mapping else None
        merge_target.mapping_status = mapping_status

        if seg["customQuestion"]:
            pending_custom_questions.append((merge_target.id, seg["customQuestion"]))

        mapped_questions = []
        for m in confident_mappings:
            question_text = qmap.get(m["questionId"], "")
            mapped_questions.append({
                "questionId": m["questionId"],
                "questionText": question_text,
                "answer": m["answer"],
                "confidence": m["confidence"],
            })
            pending_responses.append({
                "questionId": m["questionId"],
                "questionText": question_text,
                "answer": m["answer"],
            })

        utterances_out.append({
            "speakerRole": seg["speakerRole"],
            "transcript": merged_text,
            "sequenceNumber": merge_target.sequence_number,
            "mappedQuestions": mapped_questions,
        })

        # The placeholder that produced this merged segment is superseded by merge_target —
        # its content now lives there instead, so remove it rather than leaving a disconnected,
        # empty duplicate row behind.
        await db.delete(pending_rows[0])
        remaining_segments = segments[1:]
        remaining_placeholders = pending_rows[1:]
    else:
        remaining_segments = segments
        remaining_placeholders = pending_rows

    for i, segment in enumerate(remaining_segments):
        confident_mappings = [
            m for m in segment["mappings"]
            if m["confidence"] >= CONVERSATION_MAPPING_CONFIDENCE_THRESHOLD
        ]
        primary_mapping = confident_mappings[0] if confident_mappings else None
        mapping_status = "MAPPED" if confident_mappings else ("SMALL_TALK" if segment["speakerRole"] == "SMALL_TALK" else "UNMAPPED")

        if i < len(remaining_placeholders):
            # Fill in the placeholder created at transcribe-time — same row, same sequence number.
            placeholder = remaining_placeholders[i]
            placeholder.speaker_role = segment["speakerRole"]
            placeholder.mapped_question_id = primary_mapping["questionId"] if primary_mapping else None
            placeholder.mapping_confidence = primary_mapping["confidence"] if primary_mapping else None
            placeholder.mapping_status = mapping_status
            utterance_id = placeholder.id
            sequence_number = placeholder.sequence_number
        else:
            # Non-diarized mode can split one flat transcript into more turns than the single
            # placeholder it arrived as — anything beyond that placeholder is a brand-new row.
            utterance_id = uuid.uuid4()
            sequence_number = next_extra_sequence
            next_extra_sequence += 1
            db.add(ConversationUtterance(
                id=utterance_id,
                session_id=session.id,
                sequence_number=sequence_number,
                speaker_role=segment["speakerRole"],
                raw_transcript=segment["text"],
                language=detected_lang,
                mapped_question_id=primary_mapping["questionId"] if primary_mapping else None,
                mapping_confidence=primary_mapping["confidence"] if primary_mapping else None,
                mapping_status=mapping_status,
            ))

        # Doctor questions the LLM couldn't match to the standard bank — logged for later
        # review and potential model fine-tuning, not surfaced anywhere in the live session.
        if segment["customQuestion"]:
            pending_custom_questions.append((utterance_id, segment["customQuestion"]))

        mapped_questions = []
        for m in confident_mappings:
            question_text = qmap.get(m["questionId"], "")
            mapped_questions.append({
                "questionId": m["questionId"],
                "questionText": question_text,
                "answer": m["answer"],
                "confidence": m["confidence"],
            })
            pending_responses.append({
                "questionId": m["questionId"],
                "questionText": question_text,
                "answer": m["answer"],
            })

        utterances_out.append({
            "speakerRole": segment["speakerRole"],
            "transcript": segment["text"],
            "sequenceNumber": sequence_number,
            "mappedQuestions": mapped_questions,
        })

    # Rare: the LLM returned fewer turns than there were placeholders. Mark the leftovers
    # UNMAPPED rather than leaving them stuck in PENDING (which would hide them from
    # get_conversation_history and silently drop their sequence number forever).
    for leftover in remaining_placeholders[len(remaining_segments):]:
        leftover.mapping_status = "UNMAPPED"

    await db.flush()

    for utterance_id, custom_question_text in pending_custom_questions:
        db.add(ConversationCustomQuestion(
            id=uuid.uuid4(),
            session_id=session.id,
            department_id=session.department_id,
            utterance_id=utterance_id,
            question_text=custom_question_text,
        ))

    # One utterance can answer several standard questions at once (e.g. "I have high
    # blood pressure and high cholesterol, but not diabetes" answers three) — save a
    # QuestionResponse row for every confident mapping across every entry in this segment.
    for i, r in enumerate(pending_responses):
        db.add(QuestionResponse(
            id=uuid.uuid4(),
            session_id=session.id,
            question_id=r["questionId"],
            question_text=r["questionText"],
            transcribed_text=r["answer"],
            input_type="VOICE",
            input_language=detected_lang,
            is_flagged=r["answer"].strip().lower() in FLAGGED_PHRASES,
            sequence_number=response_count + i + 1,
            responded_by_staff_id=staff_id,
        ))

    lock.expires_at = now + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)
    await db.commit()

    refreshed_result = await db.execute(
        select(QuestionResponse).where(QuestionResponse.session_id == session.id)
    )
    refreshed_responses = refreshed_result.scalars().all()

    return {
        "utterances": utterances_out,
        "remainingQuestions": _remaining_questions_for(refreshed_responses, questions),
    }


# ── WS /sessions/{id}/conversation-stream ─────────────────────────────────────
# Real-time counterpart to the batch /conversation-segment/transcribe + /classify pair
# above, gated by settings.CONVERSATION_STREAMING_ENABLED. Sarvam's streaming STT has no
# diarization, so every utterance here is speakerId "0" — same as the non-diarized branch
# of the batch path. A browser WebSocket handshake can't carry an Authorization header, so
# the first message the client sends after connecting must be {"type": "auth", "token"}.

_LOCK_REFRESH_INTERVAL_S = 300  # well under SESSION_LOCK_TTL_MINUTES (30 min), plenty of margin

# Sarvam's own documented WS close codes (distinct from generic RFC 6455 codes — Sarvam
# repurposes/extends them for its own API semantics):
#   1003 — rate limit, quota exceeded, or invalid subscription key   -> permanent, don't retry
#   4000 — invalid model/language_code/parameter, account not enabled -> permanent, don't retry
#   1008 — inactivity timeout (avoided by periodic pings)              -> not our concern here
#   1011 — internal server error                                      -> transient, retry with backoff
_FATAL_SARVAM_WS_CLOSE_CODES = (1003, 4000)


def _is_fatal_sarvam_error(exc: Exception) -> bool:
    if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code in (402, 429):
        return True
    if isinstance(exc, websockets.exceptions.ConnectionClosed) and exc.rcvd is not None:
        return exc.rcvd.code in _FATAL_SARVAM_WS_CLOSE_CODES
    return False


class _ConversationStreamState:
    """Shared between the two concurrent relay tasks below (browser->Sarvam and
    Sarvam->browser) plus the periodic lock-refresh task — all touch the same Sarvam
    session and the same "are we mid-outage" bookkeeping."""

    def __init__(self, language: str, sample_rate: int = STREAM_SAMPLE_RATE):
        # Browsers don't always honor a requested AudioContext sample rate exactly — the
        # frontend reports back whatever it actually got, and every PCM-rate-sensitive call
        # below (the live Sarvam connection and the outage-gap WAV wrapper) uses that value
        # instead of assuming the 16kHz default.
        self.sample_rate = sample_rate
        self.stream = SarvamStreamSession(language=language, sample_rate=sample_rate)
        self.outage_buffer = bytearray()
        self.in_outage = False
        self.give_up = False
        # Set once Sarvam reports (or we infer from a close/HTTP error) that the account is
        # out of quota/credits — a terminal condition no amount of reconnecting fixes.
        self.fatal_error: str | None = None
        # Set the moment audio starts flowing with no data event yet answered, cleared the
        # moment any data event arrives — lets relay_sarvam_to_browser force a flush on a
        # continuous talker whose utterance Sarvam's own VAD never finds a gap in.
        self.utterance_started_at: float | None = None
        self.reconnect_lock = asyncio.Lock()
        # AsyncSession is not safe for concurrent use across asyncio tasks — every DB-touching
        # section below (finalize, outage catch-up, lock refresh) acquires this first.
        self.db_lock = asyncio.Lock()


async def _submit_outage_buffer(
    state: _ConversationStreamState, db: AsyncSession, session: IntakeSession,
    staff_id: uuid.UUID, lock: PatientLock, pcm_bytes: bytes, language: str,
) -> None:
    """Audio captured while the backend<->Sarvam leg was down, submitted once through the
    batch STT path instead of replayed into the freshly reconnected live stream — bursting
    stale audio into a live connection would confuse Sarvam's own real-time VAD timing."""
    try:
        wav_bytes = pcm_to_wav_bytes(pcm_bytes, state.sample_rate)
        transcript, detected_lang = await sarvam_transcribe(wav_bytes, "outage-gap.wav", "audio/wav", language)
        transcript = transcript.strip()
        if not transcript:
            return
        if detected_lang and detected_lang != "en-IN":
            transcript = await sarvam_translate(transcript, detected_lang)
        async with state.db_lock:
            now = datetime.now(timezone.utc)
            persisted = await _persist_conversation_entries(
                db, session, lock, now, [{"speakerId": "0", "text": transcript}], detected_lang,
            )
            if persisted["utterances"]:
                await _classify_and_persist_utterances(
                    db, session, staff_id, lock, datetime.now(timezone.utc),
                    [{"utteranceId": u["utteranceId"], "speakerId": u["speakerId"]} for u in persisted["utterances"]],
                )
    except Exception as exc:
        if _is_fatal_sarvam_error(exc):
            # The batch fallback hits the same Sarvam account as the realtime socket — if
            # it's out of quota, every future call to *either* API will fail identically.
            # Stop both, not just this one, so the caller doesn't go on to fire off another
            # (doomed) realtime request thinking recovery succeeded.
            state.fatal_error = str(exc)
            state.give_up = True
            logger.error("Sarvam batch fallback hit a fatal account error, giving up: %s", str(exc))
            return
        # Best-effort catch-up for one outage gap — losing this specific segment's classify
        # step is a much smaller problem than crashing the whole stream over it.
        logger.error("Failed to submit Sarvam-outage audio gap via batch fallback: %s", str(exc), exc_info=True)


async def _recover_sarvam_stream(
    state: _ConversationStreamState, db: AsyncSession, session: IntakeSession,
    staff_id: uuid.UUID, lock: PatientLock, language: str,
) -> bool:
    """Called by whichever relay task first notices the Sarvam connection died. Only one
    caller actually retries — guarded by reconnect_lock — since both directions share the
    one underlying Sarvam connection and would otherwise both try to reconnect at once."""
    async with state.reconnect_lock:
        if not state.in_outage:
            return True  # another caller already recovered it while we waited for the lock
        if state.fatal_error:
            # Already known to be a dead account — connect() would just succeed at the
            # transport level and get closed again the moment audio flows, looping forever.
            state.give_up = True
            return False
        for attempt, backoff in enumerate(STREAM_RECONNECT_BACKOFF_S, start=1):
            try:
                await state.stream.connect()
                state.in_outage = False
                if state.outage_buffer:
                    buffered = bytes(state.outage_buffer)
                    state.outage_buffer = bytearray()
                    await _submit_outage_buffer(state, db, session, staff_id, lock, buffered, language)
                    if state.give_up:
                        # The batch fallback just discovered the account is out of quota —
                        # the realtime socket we "successfully" reconnected to is doomed
                        # too, so report failure instead of inviting one more live request.
                        return False
                return True
            except Exception as exc:
                if _is_fatal_sarvam_error(exc):
                    state.fatal_error = str(exc)
                    state.give_up = True
                    logger.error("Sarvam stream reconnect hit a fatal account error, giving up: %s", str(exc))
                    return False
                logger.warning(
                    "Sarvam stream reconnect attempt %d/%d failed: %s",
                    attempt, STREAM_RECONNECT_ATTEMPTS, str(exc),
                )
                await asyncio.sleep(backoff)
        state.give_up = True
        return False


async def _finalize_stream_utterance(
    db: AsyncSession, session: IntakeSession, staff_id: uuid.UUID, lock: PatientLock,
    text: str, language: str, websocket: WebSocket, db_lock: asyncio.Lock,
) -> None:
    """Runs once per Sarvam END_SPEECH — persists the finished utterance as a PENDING
    placeholder, then classifies it, same two-step shape as the batch path. Awaited inline
    (never spawned as a background task) so classify calls for one session stay strictly
    ordered, matching the batch /classify endpoint's own sequencing guarantee."""
    t_start = time.monotonic()
    logger.info("Stream utterance finalize start: %r", text)

    # Unlike the batch path (which connects with mode="transcribe" and needs a separate
    # translate_to_english call), this stream's SarvamStreamSession.connect() always sets
    # mode="translate" — Sarvam itself returns transcript.final already translated to English,
    # so text here is already English. Re-running it through translate_to_english would treat
    # already-English text as if it were still source-language, corrupting the result.
    async with db_lock:
        now = datetime.now(timezone.utc)
        persisted = await _persist_conversation_entries(
            db, session, lock, now, [{"speakerId": "0", "text": text}], language,
        )
        if not persisted["utterances"]:
            logger.info("Stream utterance had nothing to persist — nothing sent to browser")
            return
        utterance_refs = [
            {"utteranceId": u["utteranceId"], "speakerId": u["speakerId"]} for u in persisted["utterances"]
        ]
        t_before_classify = time.monotonic()
        try:
            classified = await _classify_and_persist_utterances(
                db, session, staff_id, lock, datetime.now(timezone.utc), utterance_refs,
            )
        except HTTPException as exc:
            logger.warning(
                "Stream utterance classify failed after %.2fs: %s", time.monotonic() - t_before_classify, exc.detail,
            )
            await websocket.send_json({"type": "classify_failed", "detail": exc.detail})
            return
        logger.info("Stream utterance classified in %.2fs", time.monotonic() - t_before_classify)

    await websocket.send_json({"type": "utterance", **classified})
    logger.info("Stream utterance sent to browser — total finalize time %.2fs", time.monotonic() - t_start)


@router.websocket("/{session_id}/conversation-stream")
async def conversation_stream(
    websocket: WebSocket,
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
):
    await websocket.accept()

    try:
        auth_msg = await asyncio.wait_for(websocket.receive_json(), timeout=10.0)
    except Exception:
        await websocket.close(code=WS_CODE_AUTH_FAILED)
        return

    if not isinstance(auth_msg, dict) or auth_msg.get("type") != "auth" or not auth_msg.get("token"):
        await websocket.close(code=WS_CODE_AUTH_FAILED)
        return

    try:
        staff = await decode_staff_token(auth_msg["token"], db)
    except HTTPException:
        await websocket.send_json({"type": "auth_failed", "reason": "INVALID_TOKEN"})
        await websocket.close(code=WS_CODE_AUTH_FAILED)
        return

    staff_id = uuid.UUID(staff["sub"])
    now = datetime.now(timezone.utc)

    result = await db.execute(
        select(IntakeSession, Patient)
        .join(Patient, IntakeSession.patient_id == Patient.id)
        .where(IntakeSession.id == session_id)
    )
    row = result.first()
    if not row:
        await websocket.send_json({"type": "auth_failed", "reason": "SESSION_NOT_FOUND"})
        await websocket.close(code=WS_CODE_NOT_FOUND)
        return
    session, patient = row

    if session.status != ACTIVE:
        await websocket.send_json({"type": "auth_failed", "reason": "SESSION_COMPLETED"})
        await websocket.close(code=WS_CODE_CONFLICT)
        return

    lock_result = await db.execute(select(PatientLock).where(PatientLock.patient_id == patient.id))
    lock = lock_result.scalar_one_or_none()
    if not lock or lock.expires_at <= now or lock.staff_id != staff_id:
        await websocket.send_json({"type": "auth_failed", "reason": "SESSION_LOCK_LOST"})
        await websocket.close(code=WS_CODE_LOCK_LOST)
        return

    language = auth_msg.get("language") or "te-IN"
    sample_rate = int(auth_msg.get("sampleRate") or STREAM_SAMPLE_RATE)
    state = _ConversationStreamState(language, sample_rate)

    try:
        await state.stream.connect()
    except Exception as exc:
        logger.error("Sarvam stream connect failed: %s", str(exc), exc_info=True)
        reason = "SARVAM_QUOTA_EXCEEDED" if _is_fatal_sarvam_error(exc) else "SARVAM_UNAVAILABLE"
        await websocket.send_json({"type": "fallback_required", "reason": reason})
        await websocket.close(code=WS_CODE_FALLBACK_REQUIRED)
        return

    await websocket.send_json({"type": "ready"})

    async def refresh_lock_periodically():
        while True:
            await asyncio.sleep(_LOCK_REFRESH_INTERVAL_S)
            async with state.db_lock:
                lock.expires_at = datetime.now(timezone.utc) + timedelta(minutes=settings.SESSION_LOCK_TTL_MINUTES)
                await db.commit()

    async def relay_browser_to_sarvam():
        while True:
            message = await websocket.receive()
            if message["type"] == "websocket.disconnect":
                return
            audio_bytes = message.get("bytes")
            if not audio_bytes:
                continue
            if state.in_outage:
                state.outage_buffer.extend(audio_bytes)
                continue
            if state.utterance_started_at is None:
                state.utterance_started_at = time.monotonic()
                logger.info("Stream utterance audio started")
            try:
                await state.stream.send_audio(audio_bytes)
            except Exception as exc:
                state.in_outage = True
                state.outage_buffer.extend(audio_bytes)
                if _is_fatal_sarvam_error(exc):
                    state.fatal_error = str(exc)
                    logger.error("Sarvam stream hit a fatal account error, giving up: %s", str(exc))
                    state.give_up = True
                    return
                if not await _recover_sarvam_stream(state, db, session, staff_id, lock, language):
                    return

    async def force_flush_stalled_utterance():
        """Watchdog for the case Sarvam's own VAD never finds a gap at all (a continuous
        talker): if audio has been flowing for too long with no transcript.final event yet,
        logs it and restarts the clock so this fires at most once per STREAM_MAX_UTTERANCE_SECONDS
        window rather than every tick.

        NOTE: the speech-to-text-realtime endpoint's flush() only takes effect under
        endpointing="manual" — SarvamStreamSession runs endpointing="vad" (see its
        docstring), so the call below is a documented no-op today. It's left in place as
        the hook to act on if this watchdog needs real teeth again — e.g. by switching to
        endpointing="manual" and driving speech_start/speech_end explicitly, or by tuning
        silence_duration_ms/threshold on connect() so Sarvam's VAD itself cuts the boundary
        sooner for continuous talkers."""
        while True:
            await asyncio.sleep(1.0)
            if state.utterance_started_at is None:
                continue
            elapsed = time.monotonic() - state.utterance_started_at
            if elapsed >= settings.STREAM_MAX_UTTERANCE_SECONDS:
                logger.warning("Stream utterance stalled %.2fs with no transcript.final event", elapsed)
                state.utterance_started_at = time.monotonic()  # restart the clock, not spin-flush every tick
                try:
                    await state.stream.flush()
                except Exception as exc:
                    # A dead connection surfaces again on the next send_audio/events call in
                    # the other two tasks, which already run _recover_sarvam_stream — this
                    # nudge is best-effort, so just log rather than duplicating that recovery.
                    logger.warning("Stall-flush failed, leaving recovery to the relay tasks: %s", str(exc))

    async def relay_sarvam_to_browser():
        while True:
            if state.give_up:
                return
            try:
                async for event in state.stream.events():
                    event_type = event.get("event")
                    if event_type == "transcript.partial":
                        transcript = event.get("text") or ""
                        if transcript:
                            await websocket.send_json({"type": "partial", "text": transcript})
                    elif event_type == "transcript.final":
                        transcript = (event.get("text") or "").strip()
                        recognize_elapsed = (
                            time.monotonic() - state.utterance_started_at
                            if state.utterance_started_at else 0.0
                        )
                        logger.info(
                            "Stream transcript.final after %.2fs of audio: %r", recognize_elapsed, transcript,
                        )
                        state.utterance_started_at = None
                        if transcript:
                            await _finalize_stream_utterance(
                                db, session, staff_id, lock, transcript, language, websocket, state.db_lock,
                            )
                    elif event_type == "vad.speech_end":
                        logger.info("Stream vad.speech_end received (informational only)")
                    elif event_type == "error":
                        logger.warning("Sarvam stream error event: %s", event)
                        if event.get("is_fatal"):
                            # Sarvam itself is telling us this account can't stream right
                            # now (out of credits/rate-limited past recovery) — reconnecting
                            # would just repeat the same rejection, so stop here.
                            state.fatal_error = event.get("message") or event.get("code") or "fatal Sarvam stream error"
                            state.give_up = True
                            return
            except Exception as exc:
                if state.give_up:
                    return
                state.in_outage = True
                if _is_fatal_sarvam_error(exc):
                    state.fatal_error = str(exc)
                    logger.error("Sarvam stream hit a fatal account error, giving up: %s", str(exc))
                    state.give_up = True
                    return
                if not await _recover_sarvam_stream(state, db, session, staff_id, lock, language):
                    return

    tasks = [
        asyncio.create_task(relay_browser_to_sarvam()),
        asyncio.create_task(relay_sarvam_to_browser()),
        asyncio.create_task(refresh_lock_periodically()),
        asyncio.create_task(force_flush_stalled_utterance()),
    ]
    try:
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for t in pending:
            t.cancel()
        for t in done:
            exc = t.exception()
            if exc:
                logger.error("Conversation stream task ended with error: %s", str(exc), exc_info=exc)
    finally:
        await state.stream.close()
        try:
            if state.give_up:
                reason = "SARVAM_UNAVAILABLE"
                if state.fatal_error:
                    logger.error("Conversation stream giving up — fatal Sarvam account error: %s", state.fatal_error)
                    reason = "SARVAM_QUOTA_EXCEEDED"
                await websocket.send_json({"type": "fallback_required", "reason": reason})
                await websocket.close(code=WS_CODE_FALLBACK_REQUIRED)
            else:
                await websocket.close()
        except Exception:
            pass


# ── GET /sessions/{id}/conversation-history ───────────────────────────────────
# Returns every diarized utterance captured so far in free-conversation mode,
# ordered by sequence — used to hydrate the live transcript whenever
# conversation mode is (re)opened, including after navigating away and
# resuming the session later.

@router.get("/{session_id}/conversation-history")
async def get_conversation_history(
    session_id: uuid.UUID,
    db: AsyncSession = Depends(get_db),
    staff=Depends(get_current_staff),
):
    session_result = await db.execute(
        select(IntakeSession).where(IntakeSession.id == session_id)
    )
    if not session_result.scalar_one_or_none():
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"}
        )

    utterances_result = await db.execute(
        select(ConversationUtterance)
        .where(
            ConversationUtterance.session_id == session_id,
            # A PENDING row means classification never completed (e.g. the page was closed
            # mid-flight) — its audio is gone by now, so there's nothing to recover; exclude
            # it rather than showing a half-processed, roleless bubble on reload.
            ConversationUtterance.mapping_status != "PENDING",
        )
        .order_by(ConversationUtterance.sequence_number)
    )
    return {
        "utterances": [
            {
                "speakerRole": u.speaker_role,
                "transcript": u.raw_transcript,
                "sequenceNumber": u.sequence_number,
            }
            for u in utterances_result.scalars().all()
        ]
    }


# ── POST /sessions/{id}/speak ─────────────────────────────────────────────────

# ── PATCH /sessions/{id}/responses/{seq} ─────────────────────────────────────

class UpdateResponseRequest(BaseModel):
    answer: str

    @field_validator("answer")
    @classmethod
    def validate_answer(cls, v):
        v = v.strip()
        if not v:
            raise ValueError("answer cannot be empty")
        return v


@router.patch("/{session_id}/responses/{sequence_number}")
async def update_response(
    session_id: uuid.UUID,
    sequence_number: int,
    body: UpdateResponseRequest,
    db: AsyncSession = Depends(get_db),
    current_staff: dict = Depends(get_current_staff),
):
    # Fetch session and the target response together
    session_result = await db.execute(
        select(IntakeSession).where(IntakeSession.id == session_id)
    )
    session = session_result.scalar_one_or_none()
    if not session:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Session not found"},
        )

    response_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session_id)
        .where(QuestionResponse.sequence_number == sequence_number)
    )
    response = response_result.scalar_one_or_none()
    if not response:
        raise HTTPException(
            status_code=404,
            detail={"error": error_codes.SESSION_NOT_FOUND, "message": "Response not found"},
        )

    response.transcribed_text = body.answer
    await db.commit()

    # Regenerate summary only once a round has been submitted
    if session.status not in (FIRST_ROUND_COMPLETE, COMPLETED):
        return {"sequenceNumber": sequence_number, "answer": body.answer, "summaryText": None}

    patient_result = await db.execute(
        select(Patient).where(Patient.id == session.patient_id)
    )
    patient = patient_result.scalar_one_or_none()

    all_responses_result = await db.execute(
        select(QuestionResponse)
        .where(QuestionResponse.session_id == session_id)
        .order_by(QuestionResponse.sequence_number)
    )
    all_responses = all_responses_result.scalars().all()

    conversation = [
        {"question": r.question_text, "answer": r.transcribed_text}
        for r in all_responses
        if r.transcribed_text != _IMPLICIT_MARKER
    ]

    patient_info = {
        "name": patient.name,
        "age": patient.age,
        "gender": patient.gender,
        "daily_id": patient.daily_id,
    }

    records_result = await db.execute(
        select(PatientRecord)
        .where(PatientRecord.session_id == session_id)
        .order_by(PatientRecord.created_at)
    )
    patient_records = [r.extracted_text for r in records_result.scalars().all() if r.extracted_text]

    _, summary_prompt, _ = await _get_dept_context(db, session.department_id)

    try:
        start_time = datetime.now(timezone.utc)
        summary_text = await generate_summary(
            patient_info=patient_info,
            conversation=conversation,
            intake_date=str(patient.intake_date),
            summary_prompt=summary_prompt,
            patient_records=patient_records or None,
        )
        latency_ms = int((datetime.now(timezone.utc) - start_time).total_seconds() * 1000)
    except Exception as exc:
        logger.error("Summary regeneration failed after response edit: %s", str(exc))
        return {"sequenceNumber": sequence_number, "answer": body.answer, "summaryText": None}

    existing_result = await db.execute(
        select(PatientSummary).where(PatientSummary.session_id == session_id)
    )
    existing = existing_result.scalar_one_or_none()
    if existing:
        existing.summary_text = summary_text
        existing.generation_latency_ms = latency_ms
    else:
        db.add(PatientSummary(
            id=uuid.uuid4(),
            session_id=session_id,
            summary_text=summary_text,
            generation_latency_ms=latency_ms,
        ))

    await db.commit()
    return {"sequenceNumber": sequence_number, "answer": body.answer, "summaryText": summary_text}


class SpeakRequest(BaseModel):
    text: str
    language: str = "te-IN"

    @field_validator("language")
    @classmethod
    def validate_language(cls, v):
        if v not in SUPPORTED_LANGUAGES:
            raise ValueError(f"language must be one of: {', '.join(SUPPORTED_LANGUAGES)}")
        return v


@router.post("/{session_id}/speak")
async def speak_text(
    session_id: uuid.UUID,
    body: SpeakRequest,
    db: AsyncSession = Depends(get_db),
    current_staff: dict = Depends(get_current_staff),
):
    try:
        translated = await sarvam_translate_to_lang(body.text, body.language)
        audio_bytes = await sarvam_tts(translated, body.language)
    except Exception as exc:
        logger.error("TTS failed: %s", str(exc))
        raise HTTPException(
            status_code=500,
            detail={"error": error_codes.INTERNAL_ERROR, "message": "Unable to read this question aloud in the selected language. Please try again or choose a different language."},
        )
    return Response(content=audio_bytes, media_type="audio/wav")
