planet/backend/app/services/playground_chat_service.py

from __future__ import annotations

import asyncio
from collections.abc import Sequence
from datetime import UTC, datetime
from time import perf_counter
from uuid import uuid4

from fastapi import HTTPException, status
from sqlalchemy import func, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.session import async_session_factory
from app.models.playground_message import PlaygroundMessage
from app.models.playground_session import PlaygroundSession
from app.schemas.ai import (
    PlaygroundMessageEditRequest,
    PlaygroundMessageActionResponse,
    PlaygroundMessageCreateRequest,
    PlaygroundMessageRecord,
    PlaygroundMessageResendRequest,
    PlaygroundMessageStopRequest,
    PlaygroundSessionResponse,
    PlaygroundSessionState,
    PlaygroundSessionUpsertRequest,
    PlaygroundThreadResponse,
    SituationalAnalysisRequest,
)
from app.services.ai_client import AIProviderClient
from app.services.playground_session_store import _to_response as session_to_response
from app.services.playground_session_store import upsert_playground_session

STREAM_CHUNK_SIZE = 24
STREAM_INTERVAL_SECONDS = 0.08
THINKING_PREVIEW_SECONDS = 2.6
ORPHANED_RUN_MESSAGE = "后台生成任务已中断，请点击上一条用户消息的重试按钮重新生成。"


class _ActiveRun:
    def __init__(self, task: asyncio.Task[None]) -> None:
        self.task = task
        self.stop_requested = asyncio.Event()


_ACTIVE_RUNS: dict[str, _ActiveRun] = {}


async def _get_session_by_key(
    db: AsyncSession,
    *,
    user_id: int,
    session_key: str,
) -> PlaygroundSession | None:
    result = await db.execute(
        select(PlaygroundSession).where(
            PlaygroundSession.user_id == user_id,
            PlaygroundSession.session_key == session_key,
        )
    )
    return result.scalar_one_or_none()


async def _require_session(
    db: AsyncSession,
    *,
    user_id: int,
    session_key: str,
) -> PlaygroundSession:
    session = await _get_session_by_key(db, user_id=user_id, session_key=session_key)
    if session is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Playground session not found")
    return session


async def _require_visible_message(
    db: AsyncSession,
    *,
    user_id: int,
    public_id: str,
    role: str | None = None,
) -> PlaygroundMessage:
    conditions = [
        PlaygroundMessage.user_id == user_id,
        PlaygroundMessage.public_id == public_id,
        PlaygroundMessage.is_visible.is_(True),
    ]
    if role is not None:
        conditions.append(PlaygroundMessage.role == role)

    result = await db.execute(select(PlaygroundMessage).where(*conditions))
    message = result.scalar_one_or_none()
    if message is None:
        detail = "User message not found" if role == "user" else "Playground message not found"
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=detail)
    return message


def _message_to_record(message: PlaygroundMessage, parent_public_id: str | None = None) -> PlaygroundMessageRecord:
    return PlaygroundMessageRecord(
        id=message.public_id,
        role=message.role,
        kind=message.kind,
        status=message.status,
        title=message.title,
        content=message.content or "",
        thinking_content=message.thinking_content or "",
        meta=list(message.meta or []),
        markdown=message.role != "system",
        provider=message.provider,
        model=message.model,
        request_id=message.request_id,
        raw_response=dict(message.raw_response or {}),
        content_blocks=list(message.content_blocks or []),
        text_blocks=list(message.text_blocks or []),
        thinking_blocks=list(message.thinking_blocks or []),
        parent_message_id=parent_public_id,
        created_at=message.created_at.isoformat(),
        updated_at=message.updated_at.isoformat(),
    )


async def _ensure_session(
    db: AsyncSession,
    *,
    user_id: int,
    session_key: str,
    title: str,
    state: PlaygroundSessionState | None = None,
) -> PlaygroundSession:
    result = await db.execute(
        select(PlaygroundSession).where(
            PlaygroundSession.user_id == user_id,
            PlaygroundSession.session_key == session_key,
        )
    )
    session = result.scalar_one_or_none()
    if session is not None:
        if title:
            session.title = title[:200]
        if state is not None:
            session.state = state.model_dump(mode="json")
        await db.flush()
        await db.refresh(session)
        return session

    payload = PlaygroundSessionUpsertRequest(
        session_key=session_key,
        title=title[:200],
        state=state or PlaygroundSessionState(title=title[:200]),
    )
    await upsert_playground_session(db, user_id=user_id, payload=payload)
    result = await db.execute(
        select(PlaygroundSession).where(
            PlaygroundSession.user_id == user_id,
            PlaygroundSession.session_key == session_key,
        )
    )
    return result.scalar_one()


async def _list_visible_messages(
    db: AsyncSession,
    *,
    session_id: int,
) -> list[PlaygroundMessage]:
    result = await db.execute(
        select(PlaygroundMessage)
        .where(
            PlaygroundMessage.session_id == session_id,
            PlaygroundMessage.is_visible.is_(True),
        )
        .order_by(PlaygroundMessage.sort_order.asc(), PlaygroundMessage.id.asc())
    )
    return list(result.scalars().all())


async def _build_thread_response(
    db: AsyncSession,
    *,
    session: PlaygroundSession,
) -> PlaygroundThreadResponse:
    messages = await _list_visible_messages(db, session_id=session.id)
    messages = await _reconcile_orphaned_active_messages(db, messages)
    id_map = {item.id: item.public_id for item in messages}
    return PlaygroundThreadResponse(
        session=session_to_response(session),
        messages=[_message_to_record(item, id_map.get(item.parent_message_id)) for item in messages],
    )


async def _reconcile_orphaned_active_messages(
    db: AsyncSession,
    messages: list[PlaygroundMessage],
) -> list[PlaygroundMessage]:
    changed = False
    for item in messages:
        if item.status not in {"pending", "thinking", "answering"}:
            continue
        if item.public_id in _ACTIVE_RUNS:
            continue
        item.status = "error"
        item.content = item.content or ORPHANED_RUN_MESSAGE
        orphan_meta = "错误: 后台任务已中断"
        if orphan_meta not in (item.meta or []):
            item.meta = [*(item.meta or []), orphan_meta]
        changed = True
    if changed:
        await db.flush()
        await db.commit()
        for item in messages:
            await db.refresh(item)
    return messages


async def get_thread(
    db: AsyncSession,
    *,
    user_id: int,
    session_key: str,
) -> PlaygroundThreadResponse | None:
    session = await _get_session_by_key(db, user_id=user_id, session_key=session_key)
    if session is None:
        return None
    return await _build_thread_response(db, session=session)


async def _build_action_response(
    db: AsyncSession,
    *,
    session: PlaygroundSession,
    active_message_id: str | None = None,
) -> PlaygroundMessageActionResponse:
    thread = await _build_thread_response(db, session=session)
    return PlaygroundMessageActionResponse(
        session=thread.session,
        messages=thread.messages,
        active_message_id=active_message_id,
    )


def _collect_constraints(raw_constraints: str) -> list[str]:
    return [item.strip() for item in raw_constraints.split("\n") if item.strip()]


async def _next_sort_order(db: AsyncSession, session_id: int) -> int:
    result = await db.execute(
        select(func.max(PlaygroundMessage.sort_order)).where(PlaygroundMessage.session_id == session_id)
    )
    current = result.scalar_one_or_none()
    return int(current or 0)


async def _set_session_state(
    db: AsyncSession,
    *,
    session: PlaygroundSession,
    payload: PlaygroundMessageCreateRequest,
) -> PlaygroundSession:
    session.state = PlaygroundSessionState(
        messages=[],
        selectedPresetKey=payload.selected_preset_key,
        title=payload.title,
        objective=payload.objective,
        constraints=payload.constraints,
        inputValue="",
        analysis=None,
        latestAnalysisMessageId=None,
        analysisMeta={},
        helpExpanded=payload.help_expanded,
    ).model_dump(mode="json")
    session.title = payload.title[:200]
    await db.flush()
    await db.refresh(session)
    return session


def _spawn_assistant_run(
    *,
    user_id: int,
    session_id: int,
    session_key: str,
    user_message_id: int,
    assistant_message_id: int,
    assistant_public_id: str,
    payload: PlaygroundMessageCreateRequest,
    provider_client: AIProviderClient,
) -> None:
    task = asyncio.create_task(
        _run_assistant_message(
            user_id=user_id,
            session_id=session_id,
            session_key=session_key,
            user_message_id=user_message_id,
            assistant_message_id=assistant_message_id,
            payload=payload,
            provider_client=provider_client,
        )
    )
    _ACTIVE_RUNS[assistant_public_id] = _ActiveRun(task)


async def create_turn(
    db: AsyncSession,
    *,
    user_id: int,
    payload: PlaygroundMessageCreateRequest,
    provider_client: AIProviderClient,
) -> PlaygroundMessageActionResponse:
    session = await _ensure_session(
        db,
        user_id=user_id,
        session_key=payload.session_key,
        title=payload.title,
        state=PlaygroundSessionState(
            selectedPresetKey=payload.selected_preset_key,
            title=payload.title,
            objective=payload.objective,
            constraints=payload.constraints,
            inputValue="",
            helpExpanded=payload.help_expanded,
        ),
    )
    session = await _set_session_state(db, session=session, payload=payload)
    base_order = await _next_sort_order(db, session.id)

    user_message = PlaygroundMessage(
        public_id=uuid4().hex,
        session_id=session.id,
        user_id=user_id,
        role="user",
        kind="message",
        status="done",
        title=payload.selected_preset_key,
        content=payload.input,
        meta=[payload.title],
        sort_order=base_order + 10,
    )
    assistant_message = PlaygroundMessage(
        public_id=uuid4().hex,
        session_id=session.id,
        user_id=user_id,
        parent_message_id=None,
        role="assistant",
        kind="thinking",
        status="pending",
        title="AI 回应",
        content="",
        thinking_content="",
        meta=[],
        sort_order=base_order + 20,
    )
    db.add(user_message)
    await db.flush()
    assistant_message.parent_message_id = user_message.id
    db.add(assistant_message)
    await db.flush()
    await db.refresh(user_message)
    await db.refresh(assistant_message)
    await db.commit()
    await db.refresh(session)
    await db.refresh(user_message)
    await db.refresh(assistant_message)

    _spawn_assistant_run(
        user_id=user_id,
        session_id=session.id,
        session_key=payload.session_key,
        user_message_id=user_message.id,
        assistant_message_id=assistant_message.id,
        assistant_public_id=assistant_message.public_id,
        payload=payload,
        provider_client=provider_client,
    )

    return await _build_action_response(
        db,
        session=session,
        active_message_id=assistant_message.public_id,
    )


async def _create_assistant_retry_turn(
    db: AsyncSession,
    *,
    user_id: int,
    session: PlaygroundSession,
    user_message: PlaygroundMessage,
    payload: PlaygroundMessageCreateRequest,
    provider_client: AIProviderClient,
) -> PlaygroundMessageActionResponse:
    base_order = await _next_sort_order(db, session.id)
    assistant_message = PlaygroundMessage(
        public_id=uuid4().hex,
        session_id=session.id,
        user_id=user_id,
        parent_message_id=user_message.id,
        role="assistant",
        kind="thinking",
        status="pending",
        title="AI 回应",
        content="",
        thinking_content="",
        meta=[],
        sort_order=base_order + 10,
    )
    db.add(assistant_message)
    await db.flush()
    await db.commit()
    await db.refresh(session)
    await db.refresh(assistant_message)

    _spawn_assistant_run(
        user_id=user_id,
        session_id=session.id,
        session_key=payload.session_key,
        user_message_id=user_message.id,
        assistant_message_id=assistant_message.id,
        assistant_public_id=assistant_message.public_id,
        payload=payload,
        provider_client=provider_client,
    )

    return await _build_action_response(
        db,
        session=session,
        active_message_id=assistant_message.public_id,
    )


async def stop_message(
    db: AsyncSession,
    *,
    user_id: int,
    payload: PlaygroundMessageStopRequest,
) -> PlaygroundMessageActionResponse:
    session = await _require_session(db, user_id=user_id, session_key=payload.session_key)
    message = await _require_visible_message(db, user_id=user_id, public_id=payload.message_id)

    if message.status not in {"pending", "thinking", "answering"}:
        return await _build_action_response(db, session=session)

    active_run = _ACTIVE_RUNS.get(message.public_id)
    if active_run is not None:
        active_run.stop_requested.set()
        active_run.task.cancel()

    message.status = "stopped"
    if "已手动停止生成" not in (message.meta or []):
        message.meta = [*(message.meta or []), "已手动停止生成"]
    await db.flush()
    await db.commit()
    await db.refresh(message)

    return await _build_action_response(db, session=session)


async def resend_turn(
    db: AsyncSession,
    *,
    user_id: int,
    payload: PlaygroundMessageResendRequest,
    provider_client: AIProviderClient,
) -> PlaygroundMessageActionResponse:
    session = await _require_session(db, user_id=user_id, session_key=payload.session_key)
    user_message = await _require_visible_message(
        db,
        user_id=user_id,
        public_id=payload.user_message_id,
        role="user",
    )

    later_messages = await db.execute(
        select(PlaygroundMessage).where(
            PlaygroundMessage.session_id == session.id,
            PlaygroundMessage.sort_order > user_message.sort_order,
            PlaygroundMessage.is_visible.is_(True),
        )
    )
    for item in later_messages.scalars().all():
        item.is_visible = False
        if item.status in {"pending", "thinking", "answering"}:
            active_run = _ACTIVE_RUNS.get(item.public_id)
            if active_run is not None:
                active_run.stop_requested.set()
                active_run.task.cancel()
    await db.flush()
    await db.commit()

    session_state = PlaygroundSessionState.model_validate(session.state or {})
    create_payload = PlaygroundMessageCreateRequest(
        session_key=payload.session_key,
        title=session_state.title or session.title,
        objective=session_state.objective or "继续当前对话",
        constraints=session_state.constraints or "",
        input=user_message.content,
        selected_preset_key=session_state.selectedPresetKey or "bgp-brief",
        help_expanded=session_state.helpExpanded,
    )
    return await _create_assistant_retry_turn(
        db,
        user_id=user_id,
        session=session,
        user_message=user_message,
        payload=create_payload,
        provider_client=provider_client,
    )


async def edit_user_message(
    db: AsyncSession,
    *,
    user_id: int,
    payload: PlaygroundMessageEditRequest,
) -> PlaygroundMessageActionResponse:
    session = await _require_session(db, user_id=user_id, session_key=payload.session_key)
    user_message = await _require_visible_message(
        db,
        user_id=user_id,
        public_id=payload.user_message_id,
        role="user",
    )

    user_message.content = payload.content.strip()
    await db.flush()
    await db.commit()
    await db.refresh(user_message)

    return await _build_action_response(db, session=session)


async def _append_meta_if_missing(db: AsyncSession, message_id: int, meta_line: str) -> None:
    result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == message_id))
    message = result.scalar_one_or_none()
    if message is None:
        return
    if meta_line not in (message.meta or []):
        message.meta = [*(message.meta or []), meta_line]
        await db.flush()


async def _should_stop(message_public_id: str) -> bool:
    active_run = _ACTIVE_RUNS.get(message_public_id)
    return active_run.stop_requested.is_set() if active_run is not None else False


async def _mark_message_state(
    db: AsyncSession,
    *,
    message_id: int,
    **updates,
) -> PlaygroundMessage:
    result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == message_id))
    message = result.scalar_one()
    for key, value in updates.items():
        setattr(message, key, value)
    await db.flush()
    await db.refresh(message)
    return message


def _build_conversation_history(messages: Sequence[PlaygroundMessage], current_user_message_id: int) -> list[dict]:
    history: list[dict] = []
    for item in messages:
        if item.id >= current_user_message_id:
            break
        if item.role == "system":
            continue
        history.append(
            {
                "role": item.role,
                "kind": item.kind or "message",
                "title": item.title,
                "content": item.content or "",
            }
        )
    return history[-8:]


def _format_run_exception(exc: Exception) -> str:
    if isinstance(exc, HTTPException):
        detail = exc.detail
        if isinstance(detail, str):
            return detail
        return str(detail)
    return str(exc) or type(exc).__name__


async def _run_assistant_message(
    *,
    user_id: int,
    session_id: int,
    session_key: str,
    user_message_id: int,
    assistant_message_id: int,
    payload: PlaygroundMessageCreateRequest,
    provider_client: AIProviderClient,
) -> None:
    request_id = str(uuid4())
    started_at = perf_counter()
    assistant_public_id: str | None = None
    try:
        async with async_session_factory() as db:
            session = await db.get(PlaygroundSession, session_id)
            user_message = await db.get(PlaygroundMessage, user_message_id)
            assistant_message = await db.get(PlaygroundMessage, assistant_message_id)
            if session is None or user_message is None or assistant_message is None:
                return
            assistant_public_id = assistant_message.public_id

            visible_messages = await _list_visible_messages(db, session_id=session_id)
            conversation_history = _build_conversation_history(visible_messages, user_message_id)

            request_payload = SituationalAnalysisRequest(
                title=payload.title,
                objective=payload.objective,
                observations=[item.strip() for item in payload.input.split("\n") if item.strip()],
                constraints=_collect_constraints(payload.constraints),
                context={
                    "source": "playground",
                    "preset": payload.selected_preset_key,
                    "conversation_history": conversation_history,
                    "history_size": len(conversation_history),
                },
                thinking={"type": "enabled"},
            )

        analysis = await provider_client.analyze(request_payload, request_id=request_id)

        async with async_session_factory() as db:
            assistant_message = await _mark_message_state(
                db,
                message_id=assistant_message_id,
                status="thinking" if analysis.thinking_blocks else "answering",
                title=f"{analysis.provider} / {analysis.model}",
                provider=analysis.provider,
                model=analysis.model,
                request_id=request_id,
                raw_response=analysis.raw_response,
                content_blocks=[item.model_dump(mode="json") for item in analysis.content_blocks],
                text_blocks=analysis.text_blocks,
                thinking_blocks=analysis.thinking_blocks,
                thinking_content="\n\n".join(analysis.thinking_blocks).strip(),
            )
            session = await db.get(PlaygroundSession, session_id)
            if session is not None:
                session_state = PlaygroundSessionState.model_validate(session.state or {})
                session.state = session_state.model_copy(
                    update={
                        "latestAnalysisMessageId": assistant_message.public_id,
                        "analysis": analysis.model_dump(mode="json"),
                    }
                ).model_dump(mode="json")
                await db.flush()
            await db.commit()

        if assistant_public_id and analysis.thinking_blocks:
            await asyncio.sleep(THINKING_PREVIEW_SECONDS)
            if await _should_stop(assistant_public_id):
                return

        content = analysis.content or ""
        cursor = 0
        while cursor < len(content):
            if assistant_public_id and await _should_stop(assistant_public_id):
                return
            cursor = min(len(content), cursor + STREAM_CHUNK_SIZE)
            async with async_session_factory() as db:
                await _mark_message_state(
                    db,
                    message_id=assistant_message_id,
                    status="answering",
                    content=content[:cursor],
                )
                await db.commit()
            await asyncio.sleep(STREAM_INTERVAL_SECONDS)

        duration_ms = round((perf_counter() - started_at) * 1000)
        async with async_session_factory() as db:
            assistant_message = await _mark_message_state(
                db,
                message_id=assistant_message_id,
                status="done",
                content=content,
                meta=[
                    f"Request ID: {request_id}",
                    f"耗时: {duration_ms} ms",
                    f"完成时间: {datetime.now(UTC).astimezone().isoformat()}",
                ],
            )
            session = await db.get(PlaygroundSession, session_id)
            if session is not None:
                session_state = PlaygroundSessionState.model_validate(session.state or {})
                session.state = session_state.model_copy(
                    update={
                        "latestAnalysisMessageId": assistant_message.public_id,
                        "analysis": analysis.model_dump(mode="json"),
                        "analysisMeta": {
                            "requestId": request_id,
                            "durationMs": duration_ms,
                            "completedAt": datetime.now().isoformat(),
                        },
                    }
                ).model_dump(mode="json")
                await db.flush()
            await db.commit()
    except asyncio.CancelledError:
        async with async_session_factory() as db:
            result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == assistant_message_id))
            message = result.scalar_one_or_none()
            if message is not None and message.status in {"pending", "thinking", "answering"}:
                message.status = "stopped"
                if "已手动停止生成" not in (message.meta or []):
                    message.meta = [*(message.meta or []), "已手动停止生成"]
                await db.flush()
                await db.commit()
        raise
    except Exception as exc:
        error_message = _format_run_exception(exc)
        async with async_session_factory() as db:
            result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == assistant_message_id))
            message = result.scalar_one_or_none()
            if message is not None:
                message.status = "error"
                message.content = message.content or f"分析失败：{error_message}"
                message.meta = [
                    *(message.meta or []),
                    f"Request ID: {request_id}",
                    f"错误: {error_message}",
                ]
                await db.flush()
                await db.commit()
    finally:
        if assistant_public_id:
            _ACTIVE_RUNS.pop(assistant_public_id, None)