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 datetime import datetime
from typing import Annotated, Optional
from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from sqlalchemy import or_, select, func, text
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,
    PickupSessionRequest,
    LinkTicketRequest,
    AISessionSummary,
    AISessionDetail,
    AISessionStepResponse,
    AISessionSearchResult,
    StepOptionSchema,
)
from app.services import flowpilot_engine
from app.services.psa_documentation_service import retry_failed_push

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


# ── Pause ──

@router.post("/{session_id}/pause", status_code=204)
@limiter.limit("15/minute")
async def pause_session(
    request: Request,
    session_id: UUID,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Pause an active FlowPilot session for later resume."""
    try:
        await flowpilot_engine.pause_session(
            session_id=session_id,
            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()


# ── Resume ──

@router.post("/{session_id}/resume", status_code=204)
@limiter.limit("15/minute")
async def resume_session(
    request: Request,
    session_id: UUID,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Resume a paused FlowPilot session."""
    try:
        await flowpilot_engine.resume_session(
            session_id=session_id,
            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()


# ── Escalation Queue ──

@router.get("/escalation-queue", response_model=list[AISessionSummary])
@limiter.limit("30/minute")
async def get_escalation_queue(
    request: Request,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """List sessions requesting escalation for the current user's team."""
    if not current_user.team_id:
        return []

    result = await db.execute(
        select(AISession)
        .where(
            AISession.team_id == current_user.team_id,
            AISession.status == "requesting_escalation",
            AISession.user_id != current_user.id,  # Don't show own escalated sessions
        )
        .order_by(AISession.created_at.desc())
    )
    sessions = result.scalars().all()
    return [AISessionSummary.model_validate(s) for s in sessions]


# ── Pickup Escalated Session ──

@router.post("/{session_id}/pickup", response_model=StepResponseResponse)
@limiter.limit("5/minute")
async def pickup_session(
    request: Request,
    session_id: UUID,
    data: PickupSessionRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Pick up an escalated session as a new engineer."""
    _require_ai_enabled()
    await _check_quota(current_user, db)

    try:
        result = await flowpilot_engine.pickup_session(
            session_id=session_id,
            resume_mode=data.resume_mode,
            additional_context=data.additional_context,
            user_id=current_user.id,
            team_id=current_user.team_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 pickup failed: %s", e)
        await _record_usage(
            current_user, db,
            generation_type="flowpilot_pickup",
            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_pickup",
        input_tokens=0, output_tokens=0,
        succeeded=True,
        session_id=session_id,
    )
    await db.commit()

    return result


# ── Link Ticket ──

@router.post("/{session_id}/link-ticket", response_model=AISessionDetail)
@limiter.limit("10/minute")
async def link_ticket_to_session(
    request: Request,
    session_id: UUID,
    data: LinkTicketRequest,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Link a PSA ticket to an in-progress session retroactively."""
    try:
        await flowpilot_engine.link_ticket(
            session_id=session_id,
            psa_ticket_id=data.psa_ticket_id,
            psa_connection_id=data.psa_connection_id,
            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 updated session detail
    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")

    detail = AISessionDetail.model_validate(session)
    return detail


# ── Search sessions (Command Palette) ──

@router.get("/search", response_model=list[AISessionSearchResult])
@limiter.limit("30/minute")
async def search_sessions(
    request: Request,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    q: str = Query(..., min_length=2, max_length=200),
    limit: int = Query(5, ge=1, le=20),
):
    """Search AI sessions by content using full-text search. Used by Command Palette."""
    result = await db.execute(
        select(AISession)
        .where(
            or_(
                AISession.user_id == current_user.id,
                AISession.account_id == current_user.account_id,
            ),
            text("ai_sessions.search_vector @@ plainto_tsquery('english', :q)"),
        )
        .params(q=q)
        .order_by(AISession.created_at.desc())
        .limit(limit)
    )
    sessions = result.scalars().all()
    return [
        AISessionSearchResult(
            id=s.id,
            problem_summary=s.problem_summary,
            problem_domain=s.problem_domain,
            status=s.status,
            created_at=s.created_at,
        )
        for s in sessions
    ]


# ── Similar Sessions ──

@router.get("/{session_id}/similar")
@limiter.limit("15/minute")
async def get_similar_sessions(
    request: Request,
    session_id: UUID,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    limit: int = Query(5, ge=1, le=20),
):
    """Find sessions semantically similar to this one using vector embeddings."""
    from app.services.session_embedding_service import find_similar_sessions

    if not current_user.account_id:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="No account")

    results = await find_similar_sessions(
        session_id=session_id,
        account_id=current_user.account_id,
        db=db,
        limit=limit,
    )
    return results


# ── 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),
    problem_domain: Optional[str] = Query(None),
    matched_flow_id: Optional[UUID] = Query(None),
    confidence_tier: Optional[str] = Query(None, pattern="^(guided|exploring|discovery)$"),
    ticket_id: Optional[str] = Query(None),
    date_from: Optional[datetime] = Query(None),
    date_to: Optional[datetime] = Query(None),
    q: Optional[str] = Query(None, min_length=2, max_length=200),
):
    """List the current user's AI sessions (owned or picked up)."""
    user_id_str = str(current_user.id)
    query = (
        select(AISession)
        .where(
            or_(
                AISession.user_id == current_user.id,
                AISession.escalation_package["picked_up_by"].as_string() == user_id_str,
            )
        )
        .order_by(AISession.created_at.desc())
        .offset(skip)
        .limit(limit)
    )

    if session_status:
        query = query.where(AISession.status == session_status)
    if problem_domain:
        query = query.where(AISession.problem_domain == problem_domain)
    if matched_flow_id:
        query = query.where(AISession.matched_flow_id == matched_flow_id)
    if confidence_tier:
        query = query.where(AISession.confidence_tier == confidence_tier)
    if ticket_id:
        query = query.where(AISession.psa_ticket_id == ticket_id)
    if date_from:
        query = query.where(AISession.created_at >= date_from)
    if date_to:
        query = query.where(AISession.created_at <= date_to)
    if q:
        query = query.where(
            text("ai_sessions.search_vector @@ plainto_tsquery('english', :q)")
        ).params(q=q)

    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, escalation target, or picked-up handler
    pkg = session.escalation_package or {}
    is_handler = pkg.get("picked_up_by") == str(current_user.id)
    if session.user_id != current_user.id and session.escalated_to_id != current_user.id and not is_handler:
        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()


# ── Retry PSA Push ──

@router.post("/{session_id}/retry-psa-push")
@limiter.limit("5/minute")
async def retry_psa_push_endpoint(
    request: Request,
    session_id: UUID,
    current_user: Annotated[User, Depends(get_current_active_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
    _: None = Depends(require_engineer_or_admin),
):
    """Manually retry a failed PSA documentation push."""
    from app.models.psa_post_log import PsaPostLog

    # Find the latest failed push log for this session
    result = await db.execute(
        select(PsaPostLog)
        .where(
            PsaPostLog.ai_session_id == session_id,
            PsaPostLog.status.in_(["failed", "pending_retry"]),
        )
        .order_by(PsaPostLog.posted_at.desc())
        .limit(1)
    )
    log_entry = result.scalar_one_or_none()

    if not log_entry:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="No failed PSA push found for this session",
        )

    # Reset to pending_retry and attempt immediately
    log_entry.status = "pending_retry"
    log_entry.retry_count = max(0, log_entry.retry_count - 1)  # Give one more attempt

    success = await retry_failed_push(log_entry, db)
    await db.commit()

    return {
        "psa_push_status": "sent" if success else log_entry.status,
        "psa_push_error": log_entry.error_message if not success else None,
    }