Cortex-Inara/cortex/routers/onboarding.py

"""
Onboarding router — invite-based setup + persona creation + model connect.

Routes:
  GET  /setup/{token}      → show password setup form (step 1)
  POST /setup/{token}      → set password, redirect to persona step
  GET  /setup/persona      → show persona creation form (step 2, requires auth)
  POST /setup/persona      → create persona, redirect to /setup/model
  GET  /setup/model        → OpenRouter quick-connect (step 3, also standalone)
  POST /setup/model        → save host + model + assign to chat role, redirect to chat
"""

import logging
import re
from pathlib import Path

from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse, RedirectResponse

from auth_utils import (
    COOKIE_NAME, validate_invite, consume_invite,
    set_password, create_token,
)
from persona_template import create_persona
from persona import list_user_personas, validate as validate_persona
import model_registry

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/setup")

_STATIC = Path(__file__).parent.parent / "static"
_SLUG_RE = re.compile(r"^[a-z_][a-z0-9_-]{0,31}$")


def _setup_page(error: str = "", step: int = 1) -> str:
    html = (_STATIC / "setup.html").read_text()
    if error:
        html = html.replace(
            "<!-- ERROR -->",
            f'<p class="error">{error}</p>',
        )
    if step == 2:
        html = html.replace("location.search)", "location.search)", 1)  # noop, handled by ?step=2
    return html


# ---------------------------------------------------------------------------
# Step 2 — persona creation (requires active session)
# IMPORTANT: must be registered before /{token} so "/persona" literal wins
# ---------------------------------------------------------------------------

@router.get("/persona", include_in_schema=False)
async def persona_page(request: Request):
    from auth_utils import decode_token
    import jwt
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return RedirectResponse("/login", status_code=302)
    try:
        decode_token(token)
    except jwt.InvalidTokenError:
        return RedirectResponse("/login", status_code=302)

    html = (_STATIC / "setup.html").read_text()
    # Show step 2 directly — inject ?step=2 behaviour inline
    html = html.replace(
        "if (params.get('step') === '2') {",
        "if (true || params.get('step') === '2') {",
    )
    return HTMLResponse(html)


@router.post("/persona", include_in_schema=False)
async def persona_submit(
    request: Request,
    step: str = Form(...),
    persona_name: str = Form(...),
    display_name: str = Form(...),
    user_real_name: str = Form(...),
    emoji: str = Form(default="✨"),
    description: str = Form(default=""),
):
    from auth_utils import decode_token
    import jwt

    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return RedirectResponse("/login", status_code=302)
    try:
        username = decode_token(token)
    except jwt.InvalidTokenError:
        return RedirectResponse("/login", status_code=302)

    # Validate persona slug
    if not _SLUG_RE.match(persona_name):
        html = (_STATIC / "setup.html").read_text().replace(
            "if (params.get('step') === '2') {",
            "if (true || params.get('step') === '2') {",
        ).replace("<!-- ERROR -->", '<p class="error">Invalid persona name. Use lowercase letters, digits, _ or - only.</p>')
        return HTMLResponse(html, status_code=422)

    # Check for collision
    existing = list_user_personas(username)
    if persona_name in existing:
        html = (_STATIC / "setup.html").read_text().replace(
            "if (params.get('step') === '2') {",
            "if (true || params.get('step') === '2') {",
        ).replace("<!-- ERROR -->", f'<p class="error">Persona "{persona_name}" already exists.</p>')
        return HTMLResponse(html, status_code=422)

    create_persona(
        username=username,
        persona_name=persona_name,
        display_name=display_name.strip() or persona_name.capitalize(),
        user_real_name=user_real_name.strip() or username.capitalize(),
        emoji=emoji or "✨",
        description=description.strip(),
    )
    logger.info("persona created: %s/%s", username, persona_name)
    # Step 3: guided model setup before entering the chat
    resp = RedirectResponse("/setup/model", status_code=302)
    # Remember which persona to land on after model setup
    resp.set_cookie("cx_setup_persona", f"{username}/{persona_name}", max_age=3600, httponly=True, samesite="lax")
    return resp


# ---------------------------------------------------------------------------
# Step 1 — invite token → set password
# IMPORTANT: registered after /persona so the literal path wins above
# ---------------------------------------------------------------------------

@router.get("/{token}", include_in_schema=False)
async def setup_page(token: str, request: Request):
    """Show the password setup page for a valid invite token."""
    username = validate_invite(token)
    if not username:
        return HTMLResponse(
            "<h1 style='font-family:sans-serif;padding:2rem'>This link is invalid or has expired.</h1>",
            status_code=400,
        )
    return HTMLResponse(_setup_page())


@router.get("/{token}/persona", include_in_schema=False)
async def setup_persona_via_token(token: str, request: Request):
    """After password setup, redirect to the generic /setup/persona page."""
    # Cookie is already set — just redirect. Token is consumed so this is safe.
    return RedirectResponse("/setup/persona", status_code=302)


@router.post("/{token}", include_in_schema=False)
async def setup_submit(
    token: str,
    step: str = Form(...),
    password: str = Form(default=""),
    confirm: str = Form(default=""),
):
    username = validate_invite(token)
    if not username:
        return HTMLResponse(
            "<h1 style='font-family:sans-serif;padding:2rem'>This link is invalid or has expired.</h1>",
            status_code=400,
        )

    if step == "password":
        if len(password) < 8:
            return HTMLResponse(_setup_page("Password must be at least 8 characters."))
        if password != confirm:
            return HTMLResponse(_setup_page("Passwords do not match."))

        set_password(username, password)
        consume_invite(username)
        logger.info("setup complete (password): %s", username)

        # Log them in and move to persona step
        resp = RedirectResponse(f"/setup/{token}/persona", status_code=302)
        resp.set_cookie(
            COOKIE_NAME,
            create_token(username),
            max_age=30 * 86400,
            httponly=True,
            samesite="lax",
            secure=False,
        )
        return resp

    return HTMLResponse(_setup_page("Unknown step."), status_code=400)


# ---------------------------------------------------------------------------
# Step 3 — model connect (OpenRouter quick-connect, also standalone)
# ---------------------------------------------------------------------------

# Curated model list shown in the Step 3 dropdown.
_OPENROUTER_MODELS = [
    ("anthropic/claude-3-5-haiku-20241022",  "Claude 3.5 Haiku — Fast & affordable"),
    ("anthropic/claude-3-7-sonnet-20250219", "Claude 3.7 Sonnet — Smarter Claude"),
    ("google/gemini-2.0-flash-001",          "Gemini 2.0 Flash — Fast Google model"),
    ("meta-llama/llama-3.3-70b-instruct",    "Llama 3.3 70B — Open source"),
]


def _model_page(error: str = "", from_setup: bool = False) -> str:
    html = (_STATIC / "setup.html").read_text()
    # Hide steps 1 and 2 inline; show step 3
    html = html.replace('<div id="step-password">', '<div id="step-password" style="display:none">')
    html = html.replace('<div id="step-persona" style="display:none">', '<div id="step-persona" style="display:none">')
    html = html.replace('<div id="step-model" style="display:none">', '<div id="step-model">')
    if from_setup:
        html = html.replace("<!-- SETUP_STEP3_LABEL -->", "Step 3 of 3")
    if error:
        html = html.replace("<!-- ERROR_MODEL -->", f'<p class="error">{error}</p>')
    return html


@router.post("/model/skip", include_in_schema=False)
async def model_skip(request: Request):
    """Skip model setup — redirect to the remembered persona or user root."""
    from auth_utils import decode_token
    import jwt
    token = request.cookies.get(COOKIE_NAME)
    username = None
    if token:
        try:
            username = decode_token(token)
        except jwt.InvalidTokenError:
            pass

    dest_cookie = request.cookies.get("cx_setup_persona", "")
    dest = f"/{dest_cookie}" if dest_cookie else (f"/{username}" if username else "/")
    resp = RedirectResponse(dest, status_code=302)
    resp.delete_cookie("cx_setup_persona")
    return resp


@router.get("/model", include_in_schema=False)
async def model_page(request: Request):
    from auth_utils import decode_token
    import jwt
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return RedirectResponse("/login", status_code=302)
    try:
        decode_token(token)
    except jwt.InvalidTokenError:
        return RedirectResponse("/login", status_code=302)

    from_setup = bool(request.cookies.get("cx_setup_persona"))
    return HTMLResponse(_model_page(from_setup=from_setup))


@router.post("/model", include_in_schema=False)
async def model_submit(
    request: Request,
    api_key: str = Form(...),
    model_name: str = Form(...),
):
    from auth_utils import decode_token
    import jwt
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return RedirectResponse("/login", status_code=302)
    try:
        username = decode_token(token)
    except jwt.InvalidTokenError:
        return RedirectResponse("/login", status_code=302)

    api_key = api_key.strip()
    model_name = model_name.strip()

    if not api_key:
        from_setup = bool(request.cookies.get("cx_setup_persona"))
        return HTMLResponse(_model_page("API key is required.", from_setup=from_setup), status_code=422)

    # Save OpenRouter as a host
    host_id = model_registry.save_host(
        username=username,
        host_id=None,
        label="OpenRouter",
        api_url="https://openrouter.ai/api/v1",
        api_key=api_key,
        host_type="openai",
    )

    # Find label for selected model
    label = next((lbl for mn, lbl in _OPENROUTER_MODELS if mn == model_name), model_name)
    label = label.split(" — ")[0]  # keep just the model name part

    # Save model entry
    mid = model_registry.save_model(
        username=username,
        model_id=None,
        host_id=host_id,
        label=label,
        model_name=model_name,
        context_k=128,
        tools=True,
    )

    # Assign as chat role primary
    model_registry.set_role(username, "chat", "primary", mid)
    logger.info("openrouter setup complete: %s → %s", username, model_name)

    # Redirect to chat (use remembered persona, or user root)
    dest_cookie = request.cookies.get("cx_setup_persona", "")
    dest = f"/{dest_cookie}" if dest_cookie else f"/{username}"

    resp = RedirectResponse(dest, status_code=302)
    resp.delete_cookie("cx_setup_persona")
    return resp