from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.playground_session import PlaygroundSession
from app.schemas.ai import (
    PlaygroundSessionResponse,
    PlaygroundSessionState,
    PlaygroundSessionUpsertRequest,
)


def _to_response(record: PlaygroundSession) -> PlaygroundSessionResponse:
    return PlaygroundSessionResponse(
        id=str(record.id),
        session_key=record.session_key,
        title=record.title,
        state=PlaygroundSessionState.model_validate(record.state or {}),
        created_at=record.created_at.isoformat(),
        updated_at=record.updated_at.isoformat(),
    )


async def get_playground_session(
    db: AsyncSession,
    *,
    user_id: int,
    session_key: str = "default",
) -> PlaygroundSessionResponse | None:
    result = await db.execute(
        select(PlaygroundSession).where(
            PlaygroundSession.user_id == user_id,
            PlaygroundSession.session_key == session_key,
        )
    )
    record = result.scalar_one_or_none()
    if record is None:
        return None
    return _to_response(record)


async def upsert_playground_session(
    db: AsyncSession,
    *,
    user_id: int,
    payload: PlaygroundSessionUpsertRequest,
) -> PlaygroundSessionResponse:
    result = await db.execute(
        select(PlaygroundSession).where(
            PlaygroundSession.user_id == user_id,
            PlaygroundSession.session_key == payload.session_key,
        )
    )
    record = result.scalar_one_or_none()
    title = (payload.title or payload.state.title or "Playground 会话").strip()[:200] or "Playground 会话"

    if record is None:
        record = PlaygroundSession(
            user_id=user_id,
            session_key=payload.session_key,
            title=title,
            state=payload.state.model_dump(mode="json"),
        )
        db.add(record)
    else:
        record.title = title
        record.state = payload.state.model_dump(mode="json")

    await db.flush()
    await db.refresh(record)
    return _to_response(record)