"""VP trial chat — ephemeral sessions that must not write VP L2."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from niannian.virtual_persona.stage import get_stage

if TYPE_CHECKING:
    from niannian.chat.pipeline import ChatPipeline
    from niannian.chat.session import Session
    from niannian.virtual_persona.store import VirtualPersonaStore


class TrialChatForbiddenError(ValueError):
    """Trial chat is only allowed in deep_learn stage."""


class TrialSessionError(ValueError):
    """Invalid or mismatched trial session."""


def is_vp_trial_session(session_meta: dict[str, Any] | None) -> bool:
    return (session_meta or {}).get("trial") is True


def skip_vp_l2_ingest(session_meta: dict[str, Any] | None) -> bool:
    """Trial sessions must not ingest atoms/wiki/sticker_learn into VP L2."""
    return is_vp_trial_session(session_meta)


def assert_trial_chat_allowed(store: VirtualPersonaStore, vp_id: str) -> None:
    stage = get_stage(store, vp_id)
    if stage != "deep_learn":
        raise TrialChatForbiddenError(
            f"trial chat only allowed in deep_learn stage, got {stage!r}"
        )


def ensure_trial_session(
    pipeline: ChatPipeline,
    *,
    vp_id: str,
    user_id: str,
    session_id: str | None = None,
) -> Session:
    vp_id = str(vp_id).strip()
    if session_id:
        session = pipeline.sessions.require(session_id)
        if not is_vp_trial_session(session.meta):
            raise TrialSessionError("session is not a trial chat session")
        equipped = str(session.meta.get("equipped_vp_id") or "").strip()
        if equipped != vp_id:
            raise TrialSessionError("trial session equipped_vp_id mismatch")
        return session

    session = pipeline.create_session(project=f"vp:{vp_id}")
    session.meta["trial"] = True
    session.meta["equipped_vp_id"] = vp_id
    session.meta["user_id"] = user_id
    session.meta["chat_mode"] = "companion"
    return session
