resolutionflow/backend/app/services/flowpilot_engine.py

"""FlowPilot Engine — core LLM orchestration for AI troubleshooting sessions.

Manages structured diagnostic conversations: intake analysis, step generation,
confidence tracking, and auto-documentation. All LLM responses are structured
JSON validated against known output shapes.
"""
import json
import logging
import uuid
from datetime import datetime, timezone
from collections.abc import AsyncIterator
from typing import Any, Optional
from uuid import UUID

from sqlalchemy import select, func, or_
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.core.ai_provider import get_ai_provider
from app.core.config import settings
from app.services.llm_utils import parse_llm_json
from app.services.notification_service import notify
from app.models.ai_session import AISession
from app.models.ai_session_step import AISessionStep
from app.models.tree import Tree
from app.schemas.ai_session import (
    AISessionCreateRequest,
    AISessionCreateResponse,
    AISessionStepResponse,
    StepOptionSchema,
    StepResponseRequest,
    StepResponseResponse,
    ResolveSessionRequest,
    EscalateSessionRequest,
    SessionCloseResponse,
    SessionDocumentation,
    DocumentationStep,
    StatusUpdateRequest,
    StatusUpdateResponse,
)

logger = logging.getLogger(__name__)

# Maximum steps per session as a safety limit
MAX_STEPS_PER_SESSION = 30

STRUCTURED_OUTPUT_SCHEMA = """\
Your response MUST be a valid JSON object with one of these shapes:

1. Diagnostic question:
{"type": "question", "content": "Brief description", "reasoning": "Internal why", "context_message": "Shown to engineer", "options": [{"label": "Human text", "value": "machine_value", "followup_hint": "or null"}], "allow_free_text": true, "allow_skip": true, "confidence": 0.65}

2. Suggested action:
{"type": "action", "content": "What to do", "reasoning": "Internal why", "context_message": "Here's what to try", "action_type": "instruction | script_generation | verification | info_request | open_script_builder", "expected_outcome": "What success looks like", "confidence": 0.78}

3. Resolution suggestion:
{"type": "resolution_suggestion", "content": "Summary of what we did", "reasoning": "Internal why", "resolution_summary": "Issue was caused by X, fixed by Y", "confidence": 0.92, "follow_up_recommendations": ["Monitor for 24 hours"]}

4. Diagnostic fork (explore multiple hypotheses in parallel):
{"type": "fork", "content": "Why we need to branch", "reasoning": "Internal why", "context_message": "Shown to engineer explaining the fork", "fork_reason": "Multiple possible root causes need independent investigation", "options": [{"label": "Branch name", "description": "What this branch will investigate"}], "confidence": 0.45}\
"""

FLOWPILOT_SYSTEM_PROMPT = """\
You are FlowPilot, an expert MSP troubleshooting assistant embedded in ResolutionFlow. You guide engineers through structured diagnosis of IT issues.

## YOUR ROLE
- Conduct systematic troubleshooting through targeted questions and actions
- Start broad, narrow down based on responses
- Never guess — ask clarifying questions when uncertain
- Suggest specific, actionable steps the engineer can verify
- When confidence is high, suggest resolution; when low, keep investigating

## RESPONSE FORMAT
You MUST respond with ONLY a valid JSON object. No markdown, no prose, no code fences.
Every response must have a "type" field: "question", "action", or "resolution_suggestion".

{structured_output_schema}

## RULES
- Maximum 5 options per question. Options should be the most likely scenarios.
- Always include relevant context in context_message — explain WHY you're asking
- confidence is a float 0.0-1.0 reflecting how certain you are about the diagnosis path
- When multiple symptoms point to one root cause with >90% confidence, suggest resolution
- When a task can be accomplished via script OR through a GUI (e.g., Active Directory Users & Computers, Microsoft 365 Admin Center), ALWAYS ask the engineer which approach they prefer BEFORE suggesting either. Present options like "Would you like me to guide you through the GUI, or would you prefer a script to automate this?" Never assume the engineer wants a script.
- Only suggest a script_generation action AFTER the engineer has confirmed they want a script-based approach
- When the engineer wants a custom script that doesn't match an existing template, suggest opening the Script Builder. Use action_type "open_script_builder" with a "script_prompt" field containing a clear description of what the script should do, and a "script_language" field (powershell, bash, or python).
- Never suggest restarting or rebooting as a first step — diagnose first
- Be specific: "Check Event Viewer > System > source NTFS" not "check the logs"

## DIAGNOSTIC FORKING
When you detect MULTIPLE equally plausible root causes that require DIFFERENT investigation paths, use a "fork" response to let the engineer explore them as parallel branches. Use forks when:
- Two or more hypotheses have similar probability and investigating one doesn't help eliminate the other
- The engineer has tried the obvious path and results are ambiguous (could be DNS OR firewall OR auth)
- Symptoms point to multiple subsystems (e.g., "slow login" could be AD replication, DNS, or group policy)
Do NOT fork when:
- One hypothesis is clearly more likely — just investigate that first
- You can ask a single question that would eliminate most possibilities
- The session has fewer than 3 steps (gather more info first)
Fork options should be 2-4 independent investigation paths. Each option label should be a clear, short hypothesis name (e.g., "DNS Resolution Issue", "AD Replication Lag").

{team_context}

{matched_flow_context}\
"""

INTAKE_CLASSIFICATION_PROMPT = """\
You are a triage classifier for IT support issues in an MSP environment.

Analyze the following intake and respond with ONLY a JSON object:
{
  "problem_summary": "One-line summary of the issue (max 120 chars)",
  "problem_domain": "One of: active_directory, networking, m365, hardware, endpoint, virtualization, security, backup, email, printing, cloud, other",
  "key_symptoms": ["symptom1", "symptom2"],
  "urgency": "low | medium | high | critical"
}\
"""


def _confidence_to_tier(confidence: float) -> str:
    """Map numeric confidence to tier label."""
    if confidence >= 0.8:
        return "guided"
    elif confidence >= 0.4:
        return "exploring"
    return "discovery"


def _parse_structured_output(raw_text: str) -> dict[str, Any]:
    """Parse and validate structured JSON from LLM response.

    Uses shared parse_llm_json for fence stripping and JSON parsing,
    then validates FlowPilot-specific output shape.
    """
    data = parse_llm_json(raw_text)

    if not isinstance(data, dict) or "type" not in data:
        raise ValueError("LLM response missing required 'type' field")

    valid_types = {"question", "action", "resolution_suggestion", "fork"}
    if data["type"] not in valid_types:
        raise ValueError(f"Unknown response type: {data['type']}")

    return data


def _build_step_response(step: AISessionStep, session: AISession) -> AISessionStepResponse:
    """Convert a model step + session state into an API response."""
    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 {}
    return 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=session.confidence_score,
    )


async def start_session(
    request: AISessionCreateRequest,
    user_id: UUID,
    account_id: UUID,
    team_id: Optional[UUID],
    db: AsyncSession,
) -> AISessionCreateResponse:
    """Start a new FlowPilot session: classify intake, match flows, get first step."""

    # 0. Process PSA ticket intake if applicable
    ticket_context_block = None
    ticket_data = None
    psa_context_status = None

    if request.intake_type == "psa_ticket" and request.psa_connection_id and request.psa_ticket_id:
        ticket_context_block, ticket_data, psa_context_status = await _process_ticket_intake(
            psa_connection_id=request.psa_connection_id,
            psa_ticket_id=request.psa_ticket_id,
            db=db,
        )
        # Enrich intake content with ticket context for classification
        if ticket_data:
            enriched_content = dict(request.intake_content)
            enriched_content["ticket_data"] = {
                "summary": ticket_data.get("ticket", {}).get("summary", ""),
                "company": ticket_data.get("company", {}).get("name", ""),
                "priority": ticket_data.get("ticket", {}).get("priority", ""),
            }
            request = request.model_copy(update={"intake_content": enriched_content})

    # 1. Classify intake via fast LLM call
    intake_text = _extract_intake_text(request.intake_content)
    # Include ticket context in classification text if available
    if ticket_context_block:
        intake_text = f"{ticket_context_block}\n\n{intake_text}"
    classification = await _classify_intake(intake_text)

    # 2. Try to match existing flows
    from app.services.flow_matching_engine import find_matches

    matches = await find_matches(
        intake_text=intake_text,
        problem_domain=classification.get("problem_domain"),
        account_id=account_id,
        db=db,
    )

    top_match = matches[0] if matches else None
    matched_flow_id = top_match["tree_id"] if top_match else None
    match_score = top_match["score"] if top_match else None
    matched_flow_name = top_match["tree_name"] if top_match else None

    # 3. Build system prompt
    matched_flow_context = ""
    if top_match and top_match.get("score", 0) > 0.5:
        matched_flow_context = (
            f"## MATCHED FLOW\n"
            f"A similar flow exists: \"{top_match['tree_name']}\" "
            f"(match score: {top_match['score']:.0%}). "
            f"Use it as a guide but adapt to the specific situation."
        )

    # Include ticket context in system prompt if available
    ticket_prompt_section = ""
    if ticket_context_block:
        ticket_prompt_section = f"\n## PSA TICKET CONTEXT\n{ticket_context_block}\n"

    # Include available script templates for in-session script generation
    script_context = await _build_script_context(team_id, db)
    if script_context:
        ticket_prompt_section += f"\n{script_context}\n"

    system_prompt = FLOWPILOT_SYSTEM_PROMPT.format(
        structured_output_schema=STRUCTURED_OUTPUT_SCHEMA,
        team_context=ticket_prompt_section,
        matched_flow_context=matched_flow_context,
    )

    # 4. Build first user message from intake
    user_message = _format_intake_message(request.intake_content, classification)

    messages = [{"role": "user", "content": user_message}]

    # 5. Call LLM for first diagnostic step
    provider = get_ai_provider(settings.get_model_for_action("open_chat"))
    raw_response, input_tokens, output_tokens = await provider.generate_json(
        system_prompt=system_prompt,
        messages=messages,
        max_tokens=2048,
    )

    # Parse with retry on failure
    try:
        parsed = _parse_structured_output(raw_response)
    except ValueError:
        # Retry once with nudge
        retry_messages = messages + [
            {"role": "assistant", "content": raw_response},
            {"role": "user", "content": "Please respond with ONLY valid JSON matching the required schema. No markdown or prose."},
        ]
        raw_response, retry_in, retry_out = await provider.generate_json(
            system_prompt=system_prompt,
            messages=retry_messages,
            max_tokens=2048,
        )
        input_tokens += retry_in
        output_tokens += retry_out
        parsed = _parse_structured_output(raw_response)

    confidence = parsed.get("confidence", 0.0)
    confidence_tier = _confidence_to_tier(confidence)

    # Initial confidence from match + classification
    if top_match and top_match.get("score", 0) > 0.8:
        confidence_tier = "guided"
        confidence = max(confidence, 0.8)

    # 6. Create session
    session = AISession(
        id=uuid.uuid4(),
        user_id=user_id,
        account_id=account_id,
        team_id=team_id,
        intake_type=request.intake_type,
        intake_content=request.intake_content,
        problem_summary=classification.get("problem_summary"),
        problem_domain=classification.get("problem_domain"),
        status="active",
        confidence_tier=confidence_tier,
        confidence_score=confidence,
        matched_flow_id=matched_flow_id,
        match_score=match_score,
        psa_ticket_id=request.psa_ticket_id,
        psa_connection_id=request.psa_connection_id,
        ticket_data=ticket_data,
        total_input_tokens=input_tokens,
        total_output_tokens=output_tokens,
        step_count=1,
        system_prompt_snapshot=system_prompt,
        conversation_messages=[
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": raw_response},
        ],
    )
    db.add(session)

    # 7a. Update matched flow usage tracking
    if matched_flow_id:
        try:
            flow_result = await db.get(Tree, matched_flow_id)
            if flow_result:
                flow_result.usage_count = (flow_result.usage_count or 0) + 1
                flow_result.last_matched_at = datetime.now(timezone.utc)
        except Exception as e:
            logger.warning("Failed to update flow usage stats for flow %s: %s", matched_flow_id, e)

    # 7. Create first step
    step = _create_step_from_parsed(
        session_id=session.id,
        step_order=0,
        parsed=parsed,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        branch_id=session.active_branch_id if session.is_branching else None,
    )
    db.add(step)

    await db.flush()

    # Generate session embedding for similar-session matching (fire-and-forget)
    try:
        from app.services.session_embedding_service import generate_session_embedding
        await generate_session_embedding(session.id, db)
    except Exception:
        logger.warning("Failed to generate session embedding on create", exc_info=True)

    return AISessionCreateResponse(
        session_id=session.id,
        status=session.status,
        confidence_tier=session.confidence_tier,
        problem_summary=session.problem_summary,
        problem_domain=session.problem_domain,
        matched_flow_id=matched_flow_id,
        matched_flow_name=matched_flow_name,
        match_score=match_score,
        first_step=_build_step_response(step, session),
        psa_context_status=psa_context_status,
    )


async def process_response(
    session_id: UUID,
    request: StepResponseRequest,
    user_id: UUID,
    db: AsyncSession,
) -> StepResponseResponse:
    """Process an engineer's response and generate the next FlowPilot step."""

    session = await _load_session(session_id, user_id, db)

    if session.status != "active":
        raise ValueError(f"Session is {session.status}, not active")

    if session.step_count >= MAX_STEPS_PER_SESSION:
        raise ValueError("Maximum steps reached for this session")

    # Update the current (latest) step with engineer's response
    latest_step = session.steps[-1] if session.steps else None
    if latest_step and latest_step.responded_at is None:
        latest_step.selected_option = request.selected_option
        latest_step.free_text_input = request.free_text_input
        latest_step.was_free_text = bool(request.free_text_input and not request.selected_option)
        latest_step.was_skipped = request.was_skipped
        latest_step.action_result = request.action_result
        latest_step.responded_at = datetime.now(timezone.utc)

    # Build the conversation message for the engineer's response
    response_text = _format_engineer_response(request)
    session.conversation_messages = session.conversation_messages + [
        {"role": "user", "content": response_text}
    ]

    # Call LLM with full conversation
    provider = get_ai_provider(settings.get_model_for_action("open_chat"))
    raw_response, input_tokens, output_tokens = await provider.generate_json(
        system_prompt=session.system_prompt_snapshot or "",
        messages=session.conversation_messages,
        max_tokens=2048,
    )

    try:
        parsed = _parse_structured_output(raw_response)
    except ValueError:
        retry_messages = session.conversation_messages + [
            {"role": "assistant", "content": raw_response},
            {"role": "user", "content": "Please respond with ONLY valid JSON matching the required schema."},
        ]
        raw_response, retry_in, retry_out = await provider.generate_json(
            system_prompt=session.system_prompt_snapshot or "",
            messages=retry_messages,
            max_tokens=2048,
        )
        input_tokens += retry_in
        output_tokens += retry_out
        parsed = _parse_structured_output(raw_response)

    # Append assistant response to conversation
    session.conversation_messages = session.conversation_messages + [
        {"role": "assistant", "content": raw_response}
    ]

    # Update session confidence
    confidence = parsed.get("confidence", session.confidence_score)
    session.confidence_score = confidence
    session.confidence_tier = _confidence_to_tier(confidence)
    session.total_input_tokens += input_tokens
    session.total_output_tokens += output_tokens
    session.step_count += 1

    # Create new step
    step = _create_step_from_parsed(
        session_id=session.id,
        step_order=session.step_count - 1,
        parsed=parsed,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        branch_id=session.active_branch_id if session.is_branching else None,
    )
    db.add(step)

    await db.flush()

    # Handle fork: create branches and enrich step content with branch IDs
    if parsed["type"] == "fork":
        from app.services.branch_manager import BranchManager
        mgr = BranchManager(db)

        # Create root branch if this is the first fork in the session
        if not session.is_branching:
            root = await mgr.create_root_branch(session.id)
            # Reassign the step to the root branch
            step.branch_id = root.id

        fork_options = parsed.get("options", [])
        fork_point, new_branches = await mgr.create_fork(
            session_id=session.id,
            parent_branch_id=session.active_branch_id,
            trigger_step_id=step.id,
            fork_reason=parsed.get("fork_reason", ""),
            options=[{"label": o["label"], "description": o.get("description", "")} for o in fork_options],
        )

        # Enrich the step content with fork_point_id and branch IDs for frontend
        enriched_content = dict(step.content or {})
        enriched_content["fork_point_id"] = str(fork_point.id)
        enriched_content["fork_branches"] = [
            {"branch_id": str(b.id), "label": b.label}
            for b in new_branches
        ]
        step.content = enriched_content
        step.is_fork_point = True
        step.fork_point_id = fork_point.id

        # Auto-switch to the first branch
        first_branch = new_branches[0]
        await mgr.switch_branch(session.id, first_branch.id)

        await db.flush()

    # Check if resolution was suggested
    resolution_suggested = parsed["type"] == "resolution_suggestion"
    resolution_summary = parsed.get("resolution_summary") if resolution_suggested else None

    return StepResponseResponse(
        session_id=session.id,
        status=session.status,
        confidence_tier=session.confidence_tier,
        confidence_score=session.confidence_score,
        next_step=_build_step_response(step, session),
        resolution_suggested=resolution_suggested,
        resolution_summary=resolution_summary,
    )


async def resolve_session(
    session_id: UUID,
    request: ResolveSessionRequest,
    user_id: UUID,
    db: AsyncSession,
) -> SessionCloseResponse:
    """Close a session as resolved and generate documentation."""
    session = await _load_session(session_id, user_id, db)

    if session.status not in ("active", "paused"):
        raise ValueError(f"Cannot resolve session in status: {session.status}")

    session.status = "resolved"
    session.resolved_at = datetime.now(timezone.utc)
    session.resolution_summary = request.resolution_summary
    session.resolution_action = request.resolution_action

    if request.session_rating is not None:
        session.session_rating = request.session_rating
    if request.session_feedback is not None:
        session.session_feedback = request.session_feedback

    documentation = _generate_documentation(session)

    # Queue for Knowledge Flywheel analysis
    session.analysis_status = "pending"

    # Recalculate success_rate for the matched flow
    if session.matched_flow_id:
        try:
            flow = await db.get(Tree, session.matched_flow_id)
            if flow:
                total_result = await db.execute(
                    select(func.count(AISession.id))
                    .where(AISession.matched_flow_id == flow.id)
                )
                resolved_result = await db.execute(
                    select(func.count(AISession.id))
                    .where(
                        AISession.matched_flow_id == flow.id,
                        AISession.status == "resolved",
                    )
                )
                total = total_result.scalar() or 0
                resolved_count = resolved_result.scalar() or 0
                flow.success_rate = round(resolved_count / total, 3) if total else None
        except Exception as e:
            logger.warning("Failed to recalculate success_rate for flow %s: %s", session.matched_flow_id, e)

    await db.flush()

    # Update session embedding with resolution data for similar-session matching
    try:
        from app.services.session_embedding_service import generate_session_embedding
        await generate_session_embedding(session.id, db)
    except Exception:
        logger.warning("Failed to update session embedding on resolve", exc_info=True)

    # Push documentation to PSA if ticket is linked
    psa_result = await _push_to_psa(session, user_id, db)

    return SessionCloseResponse(
        session_id=session.id,
        status=session.status,
        documentation=documentation,
        **psa_result,
    )


async def escalate_session(
    session_id: UUID,
    request: EscalateSessionRequest,
    user_id: UUID,
    db: AsyncSession,
) -> SessionCloseResponse:
    """Escalate a session — sets status to requesting_escalation for pickup."""
    session = await _load_session(session_id, user_id, db)

    if session.status not in ("active", "paused"):
        raise ValueError(f"Cannot escalate session in status: {session.status}")

    # Block self-escalation
    if request.escalated_to_id and request.escalated_to_id == user_id:
        raise ValueError("Cannot escalate a session to yourself. Use pause instead.")

    session.status = "requesting_escalation"
    # Don't set resolved_at — session isn't done yet
    session.escalation_reason = request.escalation_reason
    session.escalated_to_id = request.escalated_to_id

    # Build enhanced escalation package
    session.escalation_package = await _build_escalation_package_enhanced(session, user_id)

    documentation = _generate_documentation(session)

    await db.flush()

    # Notify about escalation
    await notify("session.escalated", session.account_id, {
        "session_id": str(session_id),
        "engineer_name": session.user.name if session.user else "Unknown",
        "escalation_reason": request.escalation_reason,
        "problem_summary": session.problem_summary or "N/A",
        "link": f"/pilot/{session_id}",
    }, db, target_user_ids=[request.escalated_to_id] if request.escalated_to_id else None)

    # Push documentation to PSA if ticket is linked
    psa_result = await _push_to_psa(session, user_id, db)

    return SessionCloseResponse(
        session_id=session.id,
        status=session.status,
        documentation=documentation,
        **psa_result,
    )


async def pickup_session(
    session_id: UUID,
    resume_mode: str,
    additional_context: Optional[str],
    user_id: UUID,
    team_id: Optional[UUID],
    db: AsyncSession,
) -> StepResponseResponse:
    """Pick up an escalated session as a new engineer.

    Generates a briefing step summarizing prior work, then either continues
    the conversation or starts fresh with the new engineer's context.
    """
    session = await _load_session(
        session_id, user_id, db,
        allow_team_access=True, team_id=team_id,
    )

    if session.status != "requesting_escalation":
        raise ValueError(f"Session is {session.status}, not requesting_escalation")

    # Can't pick up your own session
    if session.user_id == user_id:
        raise ValueError("Cannot pick up your own escalated session")

    # Record the pickup in the escalation package
    pkg = session.escalation_package or {}
    pkg["picked_up_by"] = str(user_id)
    pkg["picked_up_at"] = datetime.now(timezone.utc).isoformat()
    session.escalation_package = pkg

    # Reactivate the session
    session.status = "active"

    # Build a briefing message for the new engineer
    original_user_name = "the previous engineer"
    if session.user and session.user.name:
        original_user_name = session.user.name

    briefing_parts = [
        f"## Escalation Briefing",
        f"**Escalated by:** {original_user_name}",
        f"**Reason:** {session.escalation_reason or 'Not specified'}",
        "",
        f"**Problem:** {session.problem_summary or 'Unknown'}",
    ]

    steps_tried = pkg.get("steps_tried", [])
    if steps_tried:
        briefing_parts.append("")
        briefing_parts.append("**Steps already taken:**")
        for i, step in enumerate(steps_tried, 1):
            desc = step.get("description", "")
            resp = step.get("response", "")
            briefing_parts.append(f"{i}. {desc}")
            if resp:
                briefing_parts.append(f"   → {resp}")

    if hypotheses := pkg.get("remaining_hypotheses"):
        briefing_parts.append("")
        briefing_parts.append("**Remaining hypotheses:**")
        if isinstance(hypotheses, list):
            for h in hypotheses:
                briefing_parts.append(f"- {h}")
        else:
            briefing_parts.append(str(hypotheses))

    if suggestions := pkg.get("suggested_next_steps"):
        briefing_parts.append("")
        briefing_parts.append("**Suggested next steps:**")
        if isinstance(suggestions, list):
            for s in suggestions:
                briefing_parts.append(f"- {s}")
        else:
            briefing_parts.append(str(suggestions))

    briefing_text = "\n".join(briefing_parts)

    # Create a briefing step (special intake_analysis type)
    briefing_step = AISessionStep(
        id=uuid.uuid4(),
        session_id=session.id,
        branch_id=session.active_branch_id if session.is_branching else None,
        step_order=session.step_count,
        step_type="action",
        content={
            "text": briefing_text,
            "type": "briefing",
            "allow_free_text": False,
            "allow_skip": False,
        },
        context_message="Escalation briefing — here's what was tried before you.",
        confidence_at_step=session.confidence_score,
        ai_reasoning="Escalation handoff briefing for receiving engineer",
        input_tokens=0,
        output_tokens=0,
    )
    db.add(briefing_step)
    session.step_count += 1

    # Now generate the next step based on resume_mode
    if resume_mode == "fresh" and additional_context:
        # Engineer B provides their own input
        user_message = f"[Picking up escalated session] {additional_context}"
    else:
        # Continue where A left off
        user_message = (
            "[Picking up escalated session] I've reviewed the briefing above. "
            "Please continue the diagnosis based on everything tried so far."
        )

    # Append to conversation
    session.conversation_messages = session.conversation_messages + [
        {"role": "user", "content": user_message}
    ]

    # Call LLM for next step
    provider = get_ai_provider(settings.get_model_for_action("open_chat"))
    raw_response, input_tokens, output_tokens = await provider.generate_json(
        system_prompt=session.system_prompt_snapshot or "",
        messages=session.conversation_messages,
        max_tokens=2048,
    )

    try:
        parsed = _parse_structured_output(raw_response)
    except ValueError:
        retry_messages = session.conversation_messages + [
            {"role": "assistant", "content": raw_response},
            {"role": "user", "content": "Please respond with ONLY valid JSON matching the required schema."},
        ]
        raw_response, retry_in, retry_out = await provider.generate_json(
            system_prompt=session.system_prompt_snapshot or "",
            messages=retry_messages,
            max_tokens=2048,
        )
        input_tokens += retry_in
        output_tokens += retry_out
        parsed = _parse_structured_output(raw_response)

    session.conversation_messages = session.conversation_messages + [
        {"role": "assistant", "content": raw_response}
    ]

    confidence = parsed.get("confidence", session.confidence_score)
    session.confidence_score = confidence
    session.confidence_tier = _confidence_to_tier(confidence)
    session.total_input_tokens += input_tokens
    session.total_output_tokens += output_tokens
    session.step_count += 1

    next_step = _create_step_from_parsed(
        session_id=session.id,
        step_order=session.step_count - 1,
        parsed=parsed,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        branch_id=session.active_branch_id if session.is_branching else None,
    )
    db.add(next_step)

    await db.flush()

    return StepResponseResponse(
        session_id=session.id,
        status=session.status,
        confidence_tier=session.confidence_tier,
        confidence_score=session.confidence_score,
        next_step=_build_step_response(next_step, session),
        resolution_suggested=parsed["type"] == "resolution_suggestion",
        resolution_summary=parsed.get("resolution_summary") if parsed["type"] == "resolution_suggestion" else None,
    )


async def link_ticket(
    session_id: UUID,
    psa_ticket_id: str,
    psa_connection_id: UUID,
    user_id: UUID,
    db: AsyncSession,
) -> None:
    """Link a PSA ticket to an in-progress session and inject context."""
    session = await _load_session(session_id, user_id, db)

    if session.status not in ("active", "paused"):
        raise ValueError(f"Cannot link ticket to session in status: {session.status}")

    # Store the ticket link
    session.psa_ticket_id = psa_ticket_id
    session.psa_connection_id = psa_connection_id

    # Try to fetch ticket context
    ticket_context_block, ticket_data, _ = await _process_ticket_intake(
        psa_connection_id=psa_connection_id,
        psa_ticket_id=psa_ticket_id,
        db=db,
    )

    if ticket_data:
        session.ticket_data = ticket_data

    # Inject ticket context into the system prompt for subsequent steps
    if ticket_context_block and session.system_prompt_snapshot:
        ticket_section = f"\n\n## PSA TICKET CONTEXT (linked mid-session)\n{ticket_context_block}\n"
        session.system_prompt_snapshot = session.system_prompt_snapshot + ticket_section

    await db.flush()


async def pause_session(
    session_id: UUID,
    user_id: UUID,
    db: AsyncSession,
) -> None:
    """Pause an active session for the same engineer to resume later."""
    session = await _load_session(session_id, user_id, db)

    if session.status != "active":
        raise ValueError(f"Cannot pause session in status: {session.status}")

    session.status = "paused"
    await db.flush()


async def resume_session(
    session_id: UUID,
    user_id: UUID,
    db: AsyncSession,
) -> None:
    """Resume a paused session for the same engineer."""
    session = await _load_session(session_id, user_id, db)

    if session.status != "paused":
        raise ValueError(f"Cannot resume session in status: {session.status}")

    session.status = "active"
    await db.flush()


async def abandon_session(
    session_id: UUID,
    user_id: UUID,
    reason: Optional[str],
    db: AsyncSession,
) -> None:
    """Close a session without resolving or escalating.

    Used when the engineer no longer needs help, figured it out on their own,
    or the session is no longer relevant.
    """
    session = await _load_session(session_id, user_id, db)

    if session.status not in ("active", "paused"):
        raise ValueError(f"Cannot close session in status: {session.status}")

    session.status = "abandoned"
    session.resolved_at = datetime.now(timezone.utc)
    if reason:
        session.resolution_notes = reason
    await db.flush()


async def rate_session(
    session_id: UUID,
    rating: int,
    feedback: Optional[str],
    user_id: UUID,
    db: AsyncSession,
) -> None:
    """Submit post-session rating."""
    session = await _load_session(session_id, user_id, db)
    session.session_rating = rating
    session.session_feedback = feedback
    await db.flush()


async def get_session_documentation(
    session_id: UUID,
    user_id: UUID,
    db: AsyncSession,
) -> SessionDocumentation:
    """Get auto-generated documentation for a session."""
    session = await _load_session(session_id, user_id, db)
    return _generate_documentation(session)


async def generate_status_update(
    session_id: UUID,
    request: StatusUpdateRequest,
    user_id: UUID,
    db: AsyncSession,
) -> StatusUpdateResponse:
    """Generate a status update for ticket notes, client communication, or email draft."""
    session = await _load_session(session_id, user_id, db)

    # Build conversation summary from session steps
    steps_summary = []
    for step in sorted(session.steps, key=lambda s: s.step_order):
        content = step.content or {}
        if content.get("type") in ("resolution_suggestion", "briefing", "status_update"):
            continue
        text = content.get("text", "").strip()
        if not text:
            continue
        # Resolve option label instead of raw machine value
        response = None
        if step.was_skipped:
            response = "Skipped"
        elif step.selected_option and step.options_presented:
            for opt in step.options_presented:
                if opt.get("value") == step.selected_option:
                    response = opt.get("label", step.selected_option)
                    break
            else:
                response = step.selected_option
        elif step.selected_option:
            response = step.selected_option
        elif step.free_text_input:
            response = step.free_text_input
        outcome = None
        if step.action_result:
            outcome = "Succeeded" if step.action_result.get("success") else "Did not resolve"
        entry = f"{step.step_order + 1}. {text}"
        if response and response != "Skipped":
            entry += f" — {response}"
        elif response == "Skipped":
            entry += " (skipped)"
        if outcome:
            entry += f" [{outcome}]"
        steps_summary.append(entry)

    steps_text = "\n".join(steps_summary) if steps_summary else "No diagnostic steps yet."

    # Calculate time spent
    now = datetime.now(timezone.utc)
    ref_time = session.resolved_at or now
    delta = ref_time - session.created_at
    total_hrs = round(delta.total_seconds() / 3600, 2)
    time_display = f"{total_hrs} hrs"

    # Extract client name from intake or ticket data
    client_name = None
    intake = session.intake_content or {}
    if session.ticket_data:
        company = session.ticket_data.get("company", {})
        client_name = company.get("name") if isinstance(company, dict) else None
    if not client_name:
        client_name = intake.get("client_name") or intake.get("company_name")

    # Get engineer name for sign-off
    engineer_name = session.user.name if session.user and session.user.name else "Your support team"

    # Build system prompt based on audience and context
    system_prompt = _build_status_update_prompt(
        audience=request.audience,
        length=request.length,
        context=request.context,
        client_name=client_name,
        engineer_name=engineer_name,
    )

    # Build user message with full session context
    user_message = _build_status_update_context(
        session=session,
        steps_text=steps_text,
        time_display=time_display,
        context=request.context,
        client_name=client_name,
    )

    provider = get_ai_provider(settings.get_model_for_action("quick_action"))
    raw_response, input_tokens, output_tokens = await provider.generate_text(
        system_prompt=system_prompt,
        messages=[{"role": "user", "content": user_message}],
        max_tokens=1500,
    )

    # Track token usage
    session.total_input_tokens = (session.total_input_tokens or 0) + input_tokens
    session.total_output_tokens = (session.total_output_tokens or 0) + output_tokens

    # Store as a session step
    step = AISessionStep(
        id=uuid.uuid4(),
        session_id=session.id,
        branch_id=session.active_branch_id if session.is_branching else None,
        step_order=session.step_count,
        step_type="status_update",
        content={
            "audience": request.audience,
            "length": request.length,
            "context": request.context,
            "generated_content": raw_response.strip(),
            "client_name": client_name,
        },
        confidence_at_step=1.0,
    )
    db.add(step)
    session.step_count += 1
    await db.flush()

    return StatusUpdateResponse(
        content=raw_response.strip(),
        audience=request.audience,
        length=request.length,
        context=request.context,
        session_status=session.status,
        steps_completed=len(steps_summary),
        time_spent_display=time_display,
        client_name=client_name,
        generated_at=now,
    )


async def stream_ticket_notes(
    session_id: UUID,
    user_id: UUID,
    db: AsyncSession,
) -> AsyncIterator[str]:
    """Stream AI-generated structured ticket notes for a resolved session.

    Yields text chunks suitable for SSE streaming.
    """
    session = await _load_session(session_id, user_id, db)

    # Build conversation summary from messages (chat sessions)
    # or steps (guided sessions)
    messages = session.conversation_messages or []
    if messages:
        recent = messages[-20:]  # Last 20 messages for richer context
        convo_text = "\n".join(
            f"{'Engineer' if m['role'] == 'user' else 'AI Assistant'}: {m['content'][:500]}"
            for m in recent
            if isinstance(m, dict) and "role" in m and "content" in m
        )
    else:
        # Fall back to steps for guided sessions
        steps_summary = []
        for step in sorted(session.steps, key=lambda s: s.step_order):
            content = step.content or {}
            text = content.get("text", "")
            response = step.free_text_input or step.selected_option or ("Skipped" if step.was_skipped else None)
            entry = f"Step {step.step_order + 1}: {text}"
            if response:
                entry += f"\n  Engineer response: {response}"
            steps_summary.append(entry)
        convo_text = "\n".join(steps_summary) if steps_summary else "No session data."

    # Calculate time spent
    now = datetime.now(timezone.utc)
    ref_time = session.resolved_at or now
    delta = ref_time - session.created_at
    total_minutes = int(delta.total_seconds() / 60)
    time_display = f"{total_minutes} minutes" if total_minutes < 60 else f"{total_minutes // 60}h {total_minutes % 60}m"

    system_prompt = """You are generating internal ticket notes for an MSP engineer's PSA system.

Generate EXACTLY these four markdown sections, in this order:

## Problem Summary
Summarize what the engineer reported and the initial symptoms. 1-3 sentences.

## Steps Taken
List the key diagnostic steps, commands run, checks performed, and findings. Use bullet points.

## Resolution
What fixed the issue or what the final action was. Be specific and technical.

## Next Steps
Any follow-up items, monitoring to watch, or preventive measures. Write "None" if not applicable.

Rules:
- Be technical, concise, and factual
- Use markdown formatting (headers, bullet lists, bold for emphasis)
- Include specific technical details (commands, settings, error messages) where available
- Do NOT include greetings, sign-offs, or pleasantries
- Do NOT wrap output in code fences
- Output ONLY the four sections above, nothing else"""

    user_message_parts = [
        f"Session status: {session.status}",
        f"Time spent: {time_display}",
        f"Problem summary: {session.problem_summary or 'Not specified'}",
    ]
    if session.problem_domain:
        user_message_parts.append(f"Problem domain: {session.problem_domain}")
    if session.resolution_summary:
        user_message_parts.append(f"Resolution notes: {session.resolution_summary}")
    user_message_parts.append(f"\nSession conversation:\n{convo_text}")

    user_message = "\n".join(user_message_parts)

    provider = get_ai_provider(settings.get_model_for_action("quick_action"))

    # Use streaming if provider supports it (Anthropic), otherwise fall back
    try:
        async for chunk in provider.generate_text_stream(
            system_prompt=system_prompt,
            messages=[{"role": "user", "content": user_message}],
            max_tokens=1500,
        ):
            yield chunk
    except NotImplementedError:
        # Fallback for non-streaming providers (Gemini)
        text, _, _ = await provider.generate_text(
            system_prompt=system_prompt,
            messages=[{"role": "user", "content": user_message}],
            max_tokens=1500,
        )
        yield text


def _build_status_update_prompt(
    audience: str,
    length: str,
    context: str,
    client_name: str | None,
    engineer_name: str,
) -> str:
    """Build the system prompt for status update generation."""
    length_instruction = (
        "Keep it to 1-2 sentences maximum. Just the essentials."
        if length == "quick"
        else "Provide a full breakdown with steps completed, findings, and next steps."
    )

    context_labels = {
        "status": "mid-session progress update",
        "resolution": "resolution close-out summary",
        "escalation": "escalation handoff note",
    }
    context_label = context_labels.get(context, "status update")

    if audience == "ticket_notes":
        return f"""You are generating an internal {context_label} for a PSA ticket note.

Rules:
- Be technical, concise, and factual
- Use plain text with simple section headers (no markdown bold/bullets — PSA renders raw text)
- Structure as: current status paragraph, then "What We Know" section, then next steps
- "What We Know" should list confirmed findings, ruled-out causes, and open questions — keep each item to one line
- Do NOT soften language or add pleasantries
- Do NOT include greetings or sign-offs
- {length_instruction}
{"- Include root cause and resolution details since this is a close-out note" if context == "resolution" else ""}
{"- Include what was tried, what failed, and why this is being escalated" if context == "escalation" else ""}

Output ONLY the update text. No JSON, no markdown code fences, no preamble."""

    elif audience == "client_update":
        client_greeting = f"Address the client as '{client_name}'" if client_name else "Use a generic greeting like 'Hi'"
        context_guidance = {
            "status": "We're actively working on it. Describe progress made so far and what comes next without giving a timeline.",
            "resolution": "This is good news — the issue is resolved. Summarize what was wrong and what was done in plain language.",
            "escalation": "Be reassuring — explain that a specialist is being brought in to assist, not that something failed.",
        }.get(context, "")
        return f"""You are generating a brief client-facing {context_label}.

Rules:
- Be professional, reassuring, and non-technical
- NEVER use technical jargon (no "transport rules", "MX records", "DNS", "registry", "GPO", "connector", etc.)
- NEVER include server names, IP addresses, internal tool names, or ticket IDs
- Explain findings in plain language a non-technical business owner would understand
- {client_greeting}
- Sign off with: {engineer_name}
- {length_instruction}
- {context_guidance}

Output ONLY the update text. No JSON, no markdown code fences, no preamble."""

    elif audience == "request_info":
        client_greeting = f"Address the client as '{client_name}'" if client_name else "Use a generic greeting like 'Hi'"
        return f"""You are generating a brief, professional message requesting information from the client.

Rules:
- Be friendly, concise, and non-technical
- Start with one sentence explaining what you're currently working on (plain language, no jargon)
- Then list the specific questions you need answered, as a numbered list
- Each question should be clear and answerable by a non-technical user
- NEVER use technical jargon, server names, IP addresses, or internal tool names
- {client_greeting}
- Sign off with: {engineer_name}
- Keep it short — this is a targeted ask, not a status update

Output ONLY the message text. No JSON, no markdown code fences, no preamble."""

    else:  # email_draft
        client_greeting = f"Address the client as '{client_name}'" if client_name else "Use a generic greeting like 'Hi'"
        subject_hints = {
            "status": "Update: [brief issue description]",
            "resolution": "Resolved: [brief issue description]",
            "escalation": "Update: [brief issue description] — Specialist Assistance",
            "need_info": "Quick Question: [brief issue description]",
        }
        context_guidance = {
            "status": "We're actively working on it. Describe progress and next steps without giving a timeline.",
            "resolution": "This is good news — the issue is resolved. Summarize what was wrong and what was done in plain language.",
            "escalation": "Be reassuring — explain that a specialist is being brought in to assist, not that something failed.",
        }.get(context, "")
        return f"""You are generating a complete email draft for client communication.

Rules:
- Include a Subject: line at the very top
- Subject format: {subject_hints.get(context, "Update: [issue]")}
- {client_greeting}
- Be professional, reassuring, and non-technical
- NEVER use technical jargon, server names, IP addresses, or internal tool names
- Include a professional sign-off with: {engineer_name}
- {length_instruction}
- {context_guidance}

Output ONLY the email text (Subject + body). No JSON, no markdown code fences, no preamble."""


def _build_what_we_know(session: AISession) -> str:
    """Build a 'What We Know' summary from evidence_items (cockpit) or derived from steps.

    When the cockpit branch merges, session.evidence_items will be populated by the AI
    with confirmed/ruled_out/pending classifications. Until then, we derive findings
    from completed diagnostic steps.
    """
    evidence_items = getattr(session, 'evidence_items', None)
    if evidence_items:
        confirmed = [e['text'] for e in evidence_items if e.get('status') == 'confirmed']
        ruled_out = [e['text'] for e in evidence_items if e.get('status') == 'ruled_out']
        pending = [e['text'] for e in evidence_items if e.get('status') == 'pending']
        parts = []
        if confirmed:
            parts.append("Confirmed:\n" + "\n".join(f"  - {t}" for t in confirmed))
        if ruled_out:
            parts.append("Ruled out:\n" + "\n".join(f"  - {t}" for t in ruled_out))
        if pending:
            parts.append("Still investigating:\n" + "\n".join(f"  - {t}" for t in pending))
        return "\n".join(parts)

    # Derive from completed steps
    findings = []
    for step in sorted(session.steps or [], key=lambda s: s.step_order):
        content = step.content or {}
        if content.get("type") in ("resolution_suggestion", "briefing", "status_update"):
            continue
        description = content.get("text", "").strip()
        if not description or step.was_skipped:
            continue
        response = None
        if step.selected_option and step.options_presented:
            for opt in step.options_presented:
                if opt.get("value") == step.selected_option:
                    response = opt.get("label", step.selected_option)
                    break
            else:
                response = step.selected_option
        elif step.selected_option:
            response = step.selected_option
        elif step.free_text_input:
            response = step.free_text_input
        if response:
            findings.append(f"{description} — {response}")

    if not findings:
        return ""
    return "Findings so far:\n" + "\n".join(f"  - {f}" for f in findings)


def _build_status_update_context(
    session: AISession,
    steps_text: str,
    time_display: str,
    context: str,
    client_name: str | None,
) -> str:
    """Build the user message containing full session context for the AI."""
    parts = [
        f"Session status: {session.status}",
        f"Time spent: {time_display}",
        f"Problem summary: {session.problem_summary or 'Not yet determined'}",
    ]
    if session.problem_domain:
        parts.append(f"Problem domain: {session.problem_domain}")
    if client_name:
        parts.append(f"Client: {client_name}")
    if session.psa_ticket_id:
        parts.append(f"Ticket ID: {session.psa_ticket_id}")

    what_we_know = _build_what_we_know(session)
    if what_we_know:
        parts.append(f"\nWhat we know:\n{what_we_know}")

    parts.append(f"\nDiagnostic steps taken:\n{steps_text}")

    if context == "resolution" and session.resolution_summary:
        parts.append(f"\nResolution: {session.resolution_summary}")
    if context == "escalation" and session.escalation_reason:
        parts.append(f"\nEscalation reason: {session.escalation_reason}")

    return "\n".join(parts)


# ── Internal helpers ──

async def _load_session(
    session_id: UUID,
    user_id: UUID,
    db: AsyncSession,
    allow_team_access: bool = False,
    team_id: Optional[UUID] = None,
) -> AISession:
    """Load session with steps and user relationships, verifying ownership.

    Args:
        allow_team_access: If True, same-team users can access sessions in
            'requesting_escalation' status (for escalation pickup).
        team_id: Required when allow_team_access is True.
    """
    result = await db.execute(
        select(AISession)
        .options(
            selectinload(AISession.steps),
            selectinload(AISession.user),
            selectinload(AISession.escalated_to),
        )
        .where(AISession.id == session_id)
    )
    session = result.scalar_one_or_none()

    if not session:
        raise ValueError("Session not found")

    # Owner or escalation target always has access
    if session.user_id == user_id or session.escalated_to_id == user_id:
        return session

    # Engineer who picked up an escalated session has access
    pkg = session.escalation_package or {}
    if pkg.get("picked_up_by") == str(user_id):
        return session

    # Team-based access for escalation pickup
    if allow_team_access and team_id and session.team_id == team_id:
        if session.status == "requesting_escalation":
            return session

    raise PermissionError("Not authorized to access this session")


async def _classify_intake(intake_text: str) -> dict[str, Any]:
    """Quick LLM call to classify intake content."""
    try:
        provider = get_ai_provider(settings.get_model_for_action("quick_action"))
        raw, _, _ = await provider.generate_json(
            system_prompt=INTAKE_CLASSIFICATION_PROMPT,
            messages=[{"role": "user", "content": intake_text}],
            max_tokens=512,
        )
        return json.loads(raw.strip())
    except Exception as e:
        logger.warning("Intake classification failed: %s", e)
        return {
            "problem_summary": intake_text[:120],
            "problem_domain": "other",
            "key_symptoms": [],
            "urgency": "medium",
        }


def _extract_intake_text(intake_content: dict[str, Any]) -> str:
    """Extract searchable text from intake content."""
    parts = []
    if text := intake_content.get("text"):
        parts.append(text)
    if log := intake_content.get("log_content"):
        parts.append(f"Log output:\n{log}")
    if ticket := intake_content.get("ticket_data"):
        if isinstance(ticket, dict):
            parts.append(f"Ticket: {ticket.get('summary', '')}")
    return "\n\n".join(parts) if parts else str(intake_content)


def _format_intake_message(
    intake_content: dict[str, Any],
    classification: dict[str, Any],
) -> str:
    """Format intake + classification into the first user message."""
    parts = ["I need help troubleshooting an issue."]

    if text := intake_content.get("text"):
        parts.append(f"\n**Problem description:**\n{text}")

    if log := intake_content.get("log_content"):
        parts.append(f"\n**Log output:**\n```\n{log}\n```")

    if summary := classification.get("problem_summary"):
        parts.append(f"\n**Classified as:** {summary}")

    if domain := classification.get("problem_domain"):
        parts.append(f"**Domain:** {domain}")

    symptoms = classification.get("key_symptoms", [])
    if symptoms:
        parts.append(f"**Key symptoms:** {', '.join(symptoms)}")

    return "\n".join(parts)


def _format_engineer_response(request: StepResponseRequest) -> str:
    """Format engineer's step response into a conversation message."""
    if request.was_skipped:
        return "I can't check this right now / I don't know."

    parts = []
    if request.selected_option:
        parts.append(f"Selected: {request.selected_option}")

    if request.free_text_input:
        parts.append(request.free_text_input)

    if request.action_result:
        result = request.action_result
        success = "succeeded" if result.get("success") else "did not work"
        parts.append(f"Action {success}.")
        if details := result.get("details"):
            parts.append(f"Details: {details}")

    return "\n".join(parts) if parts else "No response provided."


def _create_step_from_parsed(
    session_id: UUID,
    step_order: int,
    parsed: dict[str, Any],
    input_tokens: int,
    output_tokens: int,
    branch_id: UUID | None = None,
) -> AISessionStep:
    """Create an AISessionStep from parsed LLM output."""
    step_type = parsed["type"]
    if step_type == "resolution_suggestion":
        step_type = "action"  # Store as action in DB, UI distinguishes via content

    # Build content dict (everything the UI needs to render)
    content = {
        "text": parsed.get("content", ""),
        "type": parsed["type"],
    }
    if parsed["type"] == "action":
        content["action_type"] = parsed.get("action_type", "instruction")
        content["expected_outcome"] = parsed.get("expected_outcome")
        # Script generation fields (populated when FlowPilot suggests a script)
        if parsed.get("template_id"):
            content["template_id"] = parsed["template_id"]
        if parsed.get("pre_filled_params"):
            content["pre_filled_params"] = parsed["pre_filled_params"]
        if parsed.get("instructions"):
            content["instructions"] = parsed["instructions"]
    elif parsed["type"] == "resolution_suggestion":
        content["resolution_summary"] = parsed.get("resolution_summary")
        content["follow_up_recommendations"] = parsed.get("follow_up_recommendations", [])
        content["allow_free_text"] = False
        content["allow_skip"] = False
    elif parsed["type"] == "fork":
        content["fork_reason"] = parsed.get("fork_reason", "")
        content["fork_options"] = parsed.get("options", [])
        content["allow_free_text"] = False
        content["allow_skip"] = False

    # Extract options for question type
    options = None
    if parsed["type"] == "question" and "options" in parsed:
        options = parsed["options"]
        content["allow_free_text"] = parsed.get("allow_free_text", True)
        content["allow_skip"] = parsed.get("allow_skip", True)

    return AISessionStep(
        id=uuid.uuid4(),
        session_id=session_id,
        branch_id=branch_id,
        step_order=step_order,
        step_type=step_type if parsed["type"] != "resolution_suggestion" else "action",
        content=content,
        context_message=parsed.get("context_message"),
        options_presented=options,
        confidence_at_step=parsed.get("confidence", 0.0),
        ai_reasoning=parsed.get("reasoning"),
        input_tokens=input_tokens,
        output_tokens=output_tokens,
    )


def _generate_documentation(session: AISession) -> SessionDocumentation:
    """Generate structured documentation from a session's steps."""
    diagnostic_steps = []
    follow_up_recommendations: list[str] = []

    for step in session.steps:
        content = step.content or {}
        description = content.get("text", "")

        # Determine engineer response
        engineer_response = None
        if step.was_skipped:
            engineer_response = "Skipped"
        elif step.selected_option:
            # Find the label for the selected option
            if step.options_presented:
                for opt in step.options_presented:
                    if opt.get("value") == step.selected_option:
                        engineer_response = opt.get("label", step.selected_option)
                        break
                else:
                    engineer_response = step.selected_option
            else:
                engineer_response = step.selected_option
        elif step.free_text_input:
            engineer_response = step.free_text_input

        # Determine outcome
        outcome = None
        if step.action_result:
            result = step.action_result
            outcome = "Succeeded" if result.get("success") else "Did not resolve"
            if details := result.get("details"):
                outcome += f" — {details}"

        diagnostic_steps.append(DocumentationStep(
            step_number=step.step_order + 1,
            step_type=step.step_type,
            description=description,
            engineer_response=engineer_response,
            outcome=outcome,
        ))

        # Collect follow-up recommendations from resolution suggestion steps
        if content.get("type") == "resolution_suggestion":
            recs = content.get("follow_up_recommendations", [])
            if isinstance(recs, list):
                follow_up_recommendations.extend(recs)

    # Calculate duration
    duration_display = None
    if session.resolved_at and session.created_at:
        delta = session.resolved_at - session.created_at
        minutes = int(delta.total_seconds() / 60)
        if minutes < 60:
            duration_display = f"{minutes}m"
        else:
            hours = minutes // 60
            remaining = minutes % 60
            duration_display = f"{hours}h {remaining}m"

    # Build intake summary
    intake = session.intake_content or {}
    intake_summary = intake.get("text", "")[:500]
    if not intake_summary:
        intake_summary = str(intake)[:500]

    return SessionDocumentation(
        problem_summary=session.problem_summary or "No summary available",
        problem_domain=session.problem_domain,
        intake_summary=intake_summary,
        diagnostic_steps=diagnostic_steps,
        resolution_summary=session.resolution_summary,
        escalation_reason=session.escalation_reason,
        follow_up_recommendations=follow_up_recommendations,
        total_steps=session.step_count,
        duration_display=duration_display,
        generated_at=datetime.now(timezone.utc),
    )


async def _push_to_psa(
    session: AISession,
    user_id: UUID,
    db: AsyncSession,
) -> dict[str, Any]:
    """Push documentation to PSA if session has a linked ticket.

    Returns dict with psa_push_status, psa_push_error, member_mapping_warning.
    """
    if not session.psa_ticket_id or not session.psa_connection_id:
        return {"psa_push_status": "no_psa", "psa_push_error": None, "member_mapping_warning": None}

    try:
        from app.services.psa_documentation_service import push_documentation
        return await push_documentation(session, user_id, db)
    except Exception as e:
        logger.warning("PSA documentation push failed for session %s: %s", session.id, e)
        return {
            "psa_push_status": "failed",
            "psa_push_error": str(e)[:200],
            "member_mapping_warning": None,
        }


async def _process_ticket_intake(
    psa_connection_id: UUID,
    psa_ticket_id: str,
    db: AsyncSession,
) -> tuple[Optional[str], Optional[dict[str, Any]], str]:
    """Fetch ticket context from PSA and format for AI prompt.

    Returns:
        (ticket_context_block, ticket_data_dict, psa_context_status)
        - ticket_context_block: formatted text for system prompt, or None on failure
        - ticket_data_dict: serialized TicketContext for storage, or None on failure
        - psa_context_status: "loaded" or "unavailable"
    """
    try:
        from app.services.psa.registry import get_provider_for_connection
        from app.services.psa.ticket_context import format_ticket_context_for_prompt

        provider = await get_provider_for_connection(psa_connection_id, db)
        ticket_context = await provider.get_ticket_context(
            int(psa_ticket_id), str(psa_connection_id)
        )
        ticket_prompt_block = format_ticket_context_for_prompt(ticket_context)
        ticket_data = ticket_context.model_dump(mode="json")
        return ticket_prompt_block, ticket_data, "loaded"
    except Exception as e:
        logger.warning(
            "Failed to fetch ticket context for ticket %s (connection %s): %s",
            psa_ticket_id, psa_connection_id, e,
        )
        return None, None, "unavailable"


async def _build_script_context(
    team_id: Optional[UUID],
    db: AsyncSession,
) -> Optional[str]:
    """Build script template context for the system prompt.

    Includes available script templates so FlowPilot can suggest
    script_generation actions with pre-filled parameters.
    """
    try:
        from app.models.script_template import ScriptTemplate

        result = await db.execute(
            select(ScriptTemplate)
            .where(
                ScriptTemplate.is_active.is_(True),
                or_(
                    ScriptTemplate.team_id.is_(None),
                    ScriptTemplate.team_id == team_id,
                ),
            )
            .order_by(ScriptTemplate.usage_count.desc())
            .limit(20)
        )
        templates = result.scalars().all()

        if not templates:
            return None

        lines = ["## AVAILABLE SCRIPTS"]
        lines.append("When the engineer needs to run a script, suggest an action with action_type='script_generation'.")
        lines.append("Include template_id and pre_filled_params based on the diagnostic context.\n")
        for t in templates:
            params = t.parameters_schema.get("parameters", [])
            param_keys = ", ".join(p.get("key", "") for p in params if p.get("key"))
            lines.append(f"- {t.name} (ID: {t.id}): {t.description or 'No description'}")
            if param_keys:
                lines.append(f"  Parameters: {param_keys}")

        return "\n".join(lines)
    except Exception as e:
        logger.warning("Failed to build script context: %s", e)
        return None


async def _build_escalation_package_enhanced(
    session: AISession,
    user_id: UUID,
) -> dict[str, Any]:
    """Build enhanced context package with LLM-generated hypotheses."""
    steps_tried = []
    for step in session.steps:
        content = step.content or {}
        entry = {
            "step_type": step.step_type,
            "description": content.get("text", ""),
        }
        if step.selected_option:
            entry["response"] = step.selected_option
        elif step.free_text_input:
            entry["response"] = step.free_text_input
        elif step.was_skipped:
            entry["response"] = "Skipped"
        if step.action_result:
            entry["action_result"] = step.action_result
        steps_tried.append(entry)

    package = {
        "original_user_id": str(user_id),
        "problem_summary": session.problem_summary,
        "problem_domain": session.problem_domain,
        "intake_content": session.intake_content,
        "confidence_at_escalation": session.confidence_score,
        "steps_tried": steps_tried,
        "escalation_reason": session.escalation_reason,
    }

    # LLM call for remaining hypotheses and suggested next steps (fast model)
    try:
        conversation_summary = "\n".join(
            f"- {s.get('description', '')} → {s.get('response', 'no response')}"
            for s in steps_tried
        )
        prompt = (
            "Based on this diagnostic conversation for an IT troubleshooting session:\n\n"
            f"Problem: {session.problem_summary}\n"
            f"Domain: {session.problem_domain}\n\n"
            f"Steps taken:\n{conversation_summary}\n\n"
            f"Escalation reason: {session.escalation_reason}\n\n"
            "Respond with ONLY a JSON object:\n"
            '{"remaining_hypotheses": ["hypothesis1", "hypothesis2"], '
            '"suggested_next_steps": ["step1", "step2"], '
            '"steps_ruled_out": ["ruled_out1"]}'
        )
        provider = get_ai_provider(settings.get_model_for_action("quick_action"))
        raw, _, _ = await provider.generate_json(
            system_prompt="You are an expert IT diagnostic assistant. Analyze the escalation context and provide concise insights.",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1024,
        )
        insights = json.loads(raw.strip().strip("`").lstrip("json\n"))
        package["remaining_hypotheses"] = insights.get("remaining_hypotheses", [])
        package["suggested_next_steps"] = insights.get("suggested_next_steps", [])
        package["steps_ruled_out"] = insights.get("steps_ruled_out", [])
    except Exception as e:
        logger.warning("Failed to generate escalation insights: %s", e)
        # Fall back gracefully — don't block the escalation

    return package