resolutionflow/backend/app/api/endpoints/ai_sessions.py

"""FlowPilot AI session endpoints.

CRUD and interaction endpoints for AI-powered troubleshooting sessions:
  POST   /ai-sessions                    — Start a new session
  POST   /ai-sessions/{id}/respond       — Submit step response, get next step
  POST   /ai-sessions/{id}/resolve       — Resolve the session
  POST   /ai-sessions/{id}/escalate      — Escalate the session
  GET    /ai-sessions                    — List user's sessions (paginated)
  GET    /ai-sessions/{id}              — Get session detail with all steps
  GET    /ai-sessions/{id}/documentation — Get auto-generated documentation
  POST   /ai-sessions/{id}/rate         — Submit post-session rating
"""
import logging
from typing import Annotated, Optional
from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.core.rate_limit import limiter
from app.api.deps import get_current_active_user, get_db, require_engineer_or_admin
from app.core.config import settings
from app.core.ai_quota_service import check_ai_quota, record_ai_usage, get_user_plan
from app.models.user import User
from app.models.ai_session import AISession
from app.schemas.ai_session import (
    AISessionCreateRequest,
    AISessionCreateResponse,
    StepResponseRequest,
    StepResponseResponse,
    ResolveSessionRequest,
    EscalateSessionRequest,
    SessionCloseResponse,
    SessionDocumentation,
    RateSessionRequest,
    AISessionSummary,
    AISessionDetail,
    AISessionStepResponse,
    StepOptionSchema,
)
from app.services import flowpilot_engine

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/ai-sessions", tags=["ai-sessions"])


def _require_ai_enabled() -> None:
    if not settings.ai_enabled:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="AI is not configured. Set GOOGLE_AI_API_KEY or ANTHROPIC_API_KEY.",
        )


async def _check_quota(user: User, db: AsyncSession) -> None:
    """Check AI quota and raise 429 if exceeded."""
    allowed, quota_status = await check_ai_quota(
        user_id=user.id,
        account_id=user.account_id,
        db=db,
        billing_anchor=user.ai_billing_cycle_anchor_at,
        is_super_admin=user.is_super_admin,
    )
    if not allowed:
        reset_key = "daily_reset_at" if quota_status.get("deny_reason") == "daily" else "monthly_reset_at"
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail={
                "message": f"AI limit exceeded ({quota_status['deny_reason']})",
                "reset_at": quota_status.get(reset_key),
                "quota": quota_status,
            },
        )


async def _record_usage(
    user: User,
    db: AsyncSession,
    generation_type: str,
    input_tokens: int,
    output_tokens: int,
    succeeded: bool,
    session_id: Optional[UUID] = None,
    error_code: Optional[str] = None,
) -> None:
    """Record AI usage after an LLM call."""
    plan = await get_user_plan(user.account_id, db)
    estimated_cost = (
        input_tokens * 3.0 / 1_000_000
        + output_tokens * 15.0 / 1_000_000
    )
    await record_ai_usage(
        user_id=user.id,
        account_id=user.account_id,
        conversation_id=None,
        generation_type=generation_type,
        tier=plan,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        estimated_cost=estimated_cost,
        succeeded=succeeded,
        counts_toward_quota=True,
        error_code=error_code,
        extra_data={"ai_session_id": str(session_id)} if session_id else None,
        db=db,
    )


# ── Create session ──

@router.post("", response_model=AISessionCreateResponse, status_code=201)
@limiter.limit("5/minute")
async def create_session(
    request: Request,
    data: AISessionCreateRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Start a new FlowPilot troubleshooting session."""
    _require_ai_enabled()
    await _check_quota(current_user, db)

    try:
        result = await flowpilot_engine.start_session(
            request=data,
            user_id=current_user.id,
            account_id=current_user.account_id,
            team_id=current_user.team_id,
            db=db,
        )
    except Exception as e:
        logger.exception("FlowPilot session start failed: %s", e)
        await _record_usage(
            current_user, db,
            generation_type="flowpilot_start",
            input_tokens=0, output_tokens=0,
            succeeded=False, error_code=type(e).__name__,
        )
        await db.commit()
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail=f"AI provider error ({type(e).__name__}). Please try again.",
        )

    await _record_usage(
        current_user, db,
        generation_type="flowpilot_start",
        input_tokens=result.first_step.confidence_score and 0,  # Tracked on session
        output_tokens=0,
        succeeded=True,
        session_id=result.session_id,
    )
    await db.commit()

    return result


# ── Respond to step ──

@router.post("/{session_id}/respond", response_model=StepResponseResponse)
@limiter.limit("15/minute")
async def respond_to_step(
    request: Request,
    session_id: UUID,
    data: StepResponseRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Submit an engineer's response to a FlowPilot step and get the next step."""
    _require_ai_enabled()
    await _check_quota(current_user, db)

    try:
        result = await flowpilot_engine.process_response(
            session_id=session_id,
            request=data,
            user_id=current_user.id,
            db=db,
        )
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=str(e))
    except Exception as e:
        logger.exception("FlowPilot response failed: %s", e)
        await _record_usage(
            current_user, db,
            generation_type="flowpilot_respond",
            input_tokens=0, output_tokens=0,
            succeeded=False,
            session_id=session_id,
            error_code=type(e).__name__,
        )
        await db.commit()
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail=f"AI provider error ({type(e).__name__}). Please try again.",
        )

    await _record_usage(
        current_user, db,
        generation_type="flowpilot_respond",
        input_tokens=0, output_tokens=0,
        succeeded=True,
        session_id=session_id,
    )
    await db.commit()

    return result


# ── Resolve ──

@router.post("/{session_id}/resolve", response_model=SessionCloseResponse)
@limiter.limit("15/minute")
async def resolve_session(
    request: Request,
    session_id: UUID,
    data: ResolveSessionRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Resolve a FlowPilot session and generate documentation."""
    try:
        result = await flowpilot_engine.resolve_session(
            session_id=session_id,
            request=data,
            user_id=current_user.id,
            db=db,
        )
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=str(e))

    await db.commit()
    return result


# ── Escalate ──

@router.post("/{session_id}/escalate", response_model=SessionCloseResponse)
@limiter.limit("15/minute")
async def escalate_session(
    request: Request,
    session_id: UUID,
    data: EscalateSessionRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Escalate a FlowPilot session to another engineer."""
    try:
        result = await flowpilot_engine.escalate_session(
            session_id=session_id,
            request=data,
            user_id=current_user.id,
            db=db,
        )
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=str(e))

    await db.commit()
    return result


# ── List sessions ──

@router.get("", response_model=list[AISessionSummary])
@limiter.limit("30/minute")
async def list_sessions(
    request: Request,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    session_status: Optional[str] = Query(None, alias="status"),
    skip: int = Query(0, ge=0),
    limit: int = Query(20, ge=1, le=100),
):
    """List the current user's AI sessions."""
    query = (
        select(AISession)
        .where(AISession.user_id == current_user.id)
        .order_by(AISession.created_at.desc())
        .offset(skip)
        .limit(limit)
    )

    if session_status:
        query = query.where(AISession.status == session_status)

    result = await db.execute(query)
    sessions = result.scalars().all()

    return [AISessionSummary.model_validate(s) for s in sessions]


# ── Get session detail ──

@router.get("/{session_id}", response_model=AISessionDetail)
@limiter.limit("30/minute")
async def get_session(
    request: Request,
    session_id: UUID,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
):
    """Get full session detail with all steps."""
    result = await db.execute(
        select(AISession)
        .options(selectinload(AISession.steps))
        .where(AISession.id == session_id)
    )
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Session not found")

    # Allow access if user is owner or escalation target
    if session.user_id != current_user.id and session.escalated_to_id != current_user.id:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Not authorized")

    # Build step responses
    step_responses = []
    for step in session.steps:
        options = []
        if step.options_presented:
            options = [
                StepOptionSchema(
                    label=opt.get("label", ""),
                    value=opt.get("value", ""),
                    followup_hint=opt.get("followup_hint"),
                )
                for opt in step.options_presented
            ]
        content = step.content or {}
        step_responses.append(AISessionStepResponse(
            step_id=step.id,
            step_order=step.step_order,
            step_type=step.step_type,
            content=content,
            context_message=step.context_message,
            options=options,
            allow_free_text=content.get("allow_free_text", True),
            allow_skip=content.get("allow_skip", True),
            confidence_tier=session.confidence_tier,
            confidence_score=step.confidence_at_step,
        ))

    detail = AISessionDetail.model_validate(session)
    detail.steps = step_responses
    return detail


# ── Documentation ──

@router.get("/{session_id}/documentation", response_model=SessionDocumentation)
@limiter.limit("30/minute")
async def get_documentation(
    request: Request,
    session_id: UUID,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
):
    """Get auto-generated documentation for a session."""
    try:
        return await flowpilot_engine.get_session_documentation(
            session_id=session_id,
            user_id=current_user.id,
            db=db,
        )
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=str(e))


# ── Rate ──

@router.post("/{session_id}/rate", status_code=204)
@limiter.limit("15/minute")
async def rate_session(
    request: Request,
    session_id: UUID,
    data: RateSessionRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Submit a post-session rating."""
    try:
        await flowpilot_engine.rate_session(
            session_id=session_id,
            rating=data.rating,
            feedback=data.feedback,
            user_id=current_user.id,
            db=db,
        )
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=str(e))

    await db.commit()