"""The hosted playground's sign-in through WorkOS, and the session cookie it
leaves behind.

Two programs have to agree on this: the app's login gate, and the wake
function that stands in for the app while no task is running. Whichever of
the two the load balancer routes a request to has to be able to complete a
sign-in, so both serve the same three routes from the functions here, and
the app accepts a cookie the function set without a second sign-in. The wake
function's package is built from this file and its own, so everything here
is standard library only, the call to WorkOS included.

WorkOS proves who a person is and which organization they belong to. The
playground decides who may enter: only members of the one organization named
in settings. The WorkOS access token is read once, when the sign-in
completes, and thrown away; what the playground keeps is its own cookie.

A cookie value and a sign-in ``state`` share one format: a JSON record,
base64url encoded, a dot, and an HMAC-SHA256 over the encoded record under
the session secret. The record's ``kind`` keeps the two apart, so neither can
stand in for the other. A cookie's expiry is fixed when it is issued; using
the playground does not extend it. When it lapses while the WorkOS session is
still alive, the person goes to WorkOS and straight back without typing
anything.

``/login`` answers with a page rather than a bare redirect. A link scraper
never signs in, so the page a pasted link unfurls from is this one, and it
carries the preview tags; a browser moves on to WorkOS at once.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import time
import urllib.error
import urllib.request
from dataclasses import asdict, dataclass
from html import escape
from typing import Any, Optional
from urllib.parse import urlencode

COOKIE = "quber_session"

# What a link to the playground unfurls into when it is pasted into a chat or
# a social post. A link scraper never signs in, so the page it reads is the
# sign-in page, and that page is where these tags live. The card image has to
# be reachable without a session in both states of the service, so the app
# serves it from its static files and the wake function serves the same
# bytes from its own package.
SITE_URL = "https://playground.qubera.ai"
LINK_CARD_PATH = "/static/link-card.png"
LINK_CARD_WIDTH = 1200
LINK_CARD_HEIGHT = 630
LINK_TITLE = "Qubera — extraction playground"
LINK_DESCRIPTION = (
    "Upload a financial filing and ask a question. Every figure in the answer "
    "is cited to the page and the table cell it came from."
)
SESSION_HOURS = 2
SESSION_SECONDS = SESSION_HOURS * 3600
# How long a sign-in may take between leaving for WorkOS and coming back.
STATE_SECONDS = 15 * 60
WORKOS_API = "https://api.workos.com"
WORKOS_TIMEOUT_SECONDS = 8


@dataclass(frozen=True)
class WorkOS:
    """What both programs need to run a sign-in: the WorkOS environment's API
    key and client ID, the one organization whose members may enter, and the
    secret that signs the cookie and the ``state``."""

    api_key: str
    client_id: str
    organization_id: str
    session_secret: str


@dataclass(frozen=True)
class Identity:
    """Who a session belongs to, as WorkOS reported it at sign-in. ``sid`` is
    the WorkOS session id, which sign-out needs to end that session too."""

    user_id: str
    email: str
    organization_id: str
    role: str
    sid: str


@dataclass(frozen=True)
class SignIn:
    """How a return from WorkOS ends. Accepted: a redirect to the page the
    person asked for, with a cookie. Refused or failed: a page, no cookie."""

    status: int
    location: Optional[str] = None
    token: Optional[str] = None
    page: Optional[str] = None
    identity: Optional[Identity] = None


class SignInError(Exception):
    """WorkOS did not confirm the sign-in."""


def signature(secret: str, body: str) -> str:
    return hmac.new(secret.encode(), body.encode(), hashlib.sha256).hexdigest()


def b64decode(text: str) -> bytes:
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


def encode(secret: str, record: dict[str, Any]) -> str:
    body = base64.urlsafe_b64encode(json.dumps(record, separators=(",", ":")).encode()).decode().rstrip("=")
    return f"{body}.{signature(secret, body)}"


def decode(
    secret: str, token: Optional[str], kind: str, now: Optional[float] = None
) -> Optional[dict[str, Any]]:
    """The record inside ``token`` when its signature holds, it is of ``kind``,
    and it has not expired. ``now=None`` checks against the clock; pass
    ``float("-inf")`` to ignore the expiry."""
    if not token or token.count(".") != 1:
        return None
    body, given = token.split(".")
    if not hmac.compare_digest(signature(secret, body), given):
        return None
    try:
        record = json.loads(b64decode(body))
    except ValueError:
        return None
    if not isinstance(record, dict) or record.get("kind") != kind:
        return None
    expires = record.get("expires")
    if not isinstance(expires, int) or expires <= (now if now is not None else time.time()):
        return None
    return record


def issue(secret: str, identity: Identity, now: Optional[float] = None) -> str:
    expires = int(now if now is not None else time.time()) + SESSION_SECONDS
    return encode(secret, {"kind": "session", **asdict(identity), "expires": expires})


def identity_of(record: Optional[dict[str, Any]]) -> Optional[Identity]:
    if record is None:
        return None
    try:
        return Identity(
            user_id=str(record["user_id"]),
            email=str(record["email"]),
            organization_id=str(record["organization_id"]),
            role=str(record["role"]),
            sid=str(record["sid"]),
        )
    except KeyError:
        return None


def verify(
    secret: str, organization_id: str, token: Optional[str], now: Optional[float] = None
) -> Optional[Identity]:
    """The identity in a live session cookie for ``organization_id``, or None.
    A cookie issued under another organization is refused like an expired one."""
    identity = identity_of(decode(secret, token, "session", now))
    if identity is None or not hmac.compare_digest(identity.organization_id, organization_id):
        return None
    return identity


def signed_identity(secret: str, token: Optional[str]) -> Optional[Identity]:
    """The identity in a cookie whose signature holds, expired or not. Sign-out
    reads the WorkOS session id this way, so a person whose cookie has just
    lapsed can still end their WorkOS session."""
    return identity_of(decode(secret, token, "session", float("-inf")))


def state_for(secret: str, next_path: Optional[str], now: Optional[float] = None) -> str:
    expires = int(now if now is not None else time.time()) + STATE_SECONDS
    return encode(secret, {"kind": "state", "next": safe_next(next_path), "expires": expires})


def next_from_state(secret: str, state: Optional[str], now: Optional[float] = None) -> Optional[str]:
    """The return page a ``state`` carries, or None when it is forged or stale."""
    record = decode(secret, state, "state", now)
    if record is None:
        return None
    return safe_next(str(record.get("next", "")))


def safe_next(raw: Optional[str]) -> str:
    """The page to land on after sign-in: a path on this site, never elsewhere."""
    if raw and raw.startswith("/") and not raw.startswith("//") and "\\" not in raw:
        return raw
    return "/"


def base_url(scheme: str, host: str) -> str:
    """The site a request came in on. The callback and the sign-out return
    both hang off it, so the hosted playground and a developer machine each
    get their own; WorkOS accepts only the ones registered in its dashboard."""
    return f"{scheme}://{host}"


def authorize_url(config: WorkOS, base: str, next_path: Optional[str], now: Optional[float] = None) -> str:
    query = urlencode(
        {
            "client_id": config.client_id,
            "redirect_uri": base + "/callback",
            "response_type": "code",
            "provider": "authkit",
            "state": state_for(config.session_secret, next_path, now),
        }
    )
    return f"{WORKOS_API}/user_management/authorize?{query}"


def logout_url(sid: str, base: str) -> str:
    """Where to send a browser to end its WorkOS session. WorkOS then returns it
    to ``base``, which has to be a registered sign-out URI."""
    query = urlencode({"session_id": sid, "return_to": base + "/"})
    return f"{WORKOS_API}/user_management/sessions/logout?{query}"


def authenticate(config: WorkOS, code: str) -> dict[str, Any]:
    """Trade the code WorkOS sent back for the person's identity."""
    body = json.dumps(
        {
            "client_id": config.client_id,
            "client_secret": config.api_key,
            "grant_type": "authorization_code",
            "code": code,
        }
    ).encode()
    request = urllib.request.Request(
        f"{WORKOS_API}/user_management/authenticate",
        data=body,
        headers={"Content-Type": "application/json", "Accept": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=WORKOS_TIMEOUT_SECONDS) as response:
            answer = json.loads(response.read())
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")[:300]
        raise SignInError(f"WorkOS answered {exc.code}: {detail}") from exc
    except (urllib.error.URLError, TimeoutError, ValueError) as exc:
        raise SignInError(f"WorkOS could not be reached: {exc}") from exc
    if not isinstance(answer, dict):
        raise SignInError("WorkOS answered with something other than an object")
    return answer


def claims(access_token: str) -> dict[str, Any]:
    """The claims inside a WorkOS access token, read without checking its
    signature. That is safe here only because the token arrived straight from
    WorkOS over HTTPS in the same call."""
    parts = access_token.split(".")
    if len(parts) != 3:
        return {}
    try:
        found = json.loads(b64decode(parts[1]))
    except ValueError:
        return {}
    return found if isinstance(found, dict) else {}


def complete_sign_in(
    config: WorkOS,
    base: str,
    code: Optional[str],
    state: Optional[str],
    mark_src: str,
    now: Optional[float] = None,
) -> SignIn:
    """Finish a sign-in WorkOS sent back to ``/callback``.

    The organization check happens here: an account outside the configured
    organization is refused, and gets no cookie. WorkOS has already signed
    that account in by then, so the refusal page offers to end the WorkOS
    session; otherwise the next sign-in would come straight back as the same
    account until WorkOS times it out.
    """
    target = next_from_state(config.session_secret, state, now)
    if target is None or not code:
        return SignIn(400, page=failed_page(mark_src))
    try:
        answer = authenticate(config, code)
    except SignInError:
        return SignIn(400, page=failed_page(mark_src))
    raw_user = answer.get("user")
    user: dict[str, Any] = raw_user if isinstance(raw_user, dict) else {}
    token_claims = claims(str(answer.get("access_token", "")))
    email = str(user.get("email", ""))
    sid = str(token_claims.get("sid", ""))
    if answer.get("organization_id") != config.organization_id:
        switch = logout_url(sid, base) if sid else "/login"
        return SignIn(403, page=refusal_page(email, switch, mark_src))
    identity = Identity(
        user_id=str(user.get("id", "")),
        email=email,
        organization_id=config.organization_id,
        role=str(token_claims.get("role", "")),
        sid=sid,
    )
    return SignIn(303, location=target, token=issue(config.session_secret, identity, now), identity=identity)


def cookie_header(token: str, *, secure: bool) -> str:
    """The Set-Cookie value for a session, as the app and the wake function
    both send it: HttpOnly, SameSite=Lax, and Secure on HTTPS, which is every
    hosted request."""
    return cookie_value(token, SESSION_SECONDS, secure=secure)


def clear_cookie_header(*, secure: bool) -> str:
    return cookie_value("", 0, secure=secure)


def cookie_value(token: str, max_age: int, *, secure: bool) -> str:
    parts = [f"{COOKIE}={token}", f"Max-Age={max_age}", "Path=/", "HttpOnly", "SameSite=Lax"]
    if secure:
        parts.append("Secure")
    return "; ".join(parts)


_PAGE = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1" />
<title>{title}</title>
{head}
<link rel="preconnect" href="https://fonts.googleapis.com" />
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500&amp;family=IBM+Plex+Sans:wght@400;600&amp;family=IBM+Plex+Serif:wght@400;600&amp;display=swap" />
<style>
  /* Values follow the playground's shared token set. */
  * {{ box-sizing: border-box; }}
  body {{ margin: 0; min-height: 100vh; display: grid; place-items: center; padding: 24px 16px;
         background: #f5f8fb; color: #15222e;
         font-family: "IBM Plex Sans", system-ui, -apple-system, "Segoe UI", sans-serif;
         -webkit-font-smoothing: antialiased; }}
  main {{ width: min(420px, 100%); padding: 36px; background: #ffffff; border: 1px solid #e2e8ef;
          border-radius: 14px; box-shadow: 0 4px 14px rgba(12, 32, 50, .08), 0 18px 50px rgba(12, 32, 50, .10);
          display: flex; flex-direction: column; gap: 18px; }}
  .brand {{ display: flex; align-items: center; gap: 12px; }}
  .brand img {{ width: 32px; height: 32px; display: block; }}
  .brand span {{ font-family: "IBM Plex Serif", Georgia, serif; font-weight: 600; font-size: 22px;
                 letter-spacing: -0.01em; }}
  .copy {{ display: flex; flex-direction: column; gap: 10px; }}
  .eyebrow {{ margin: 0; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11px;
              letter-spacing: 0.15em; text-transform: uppercase; color: #74828f; }}
  h1 {{ margin: 0; font-family: "IBM Plex Serif", Georgia, serif; font-weight: 400; font-size: 26px;
        line-height: 1.15; letter-spacing: -0.015em; }}
  p.message {{ margin: 0; font-size: 15px; line-height: 1.6; color: #475563; overflow-wrap: anywhere; }}
  a.action {{ align-self: flex-start; min-height: 44px; padding: 0 18px; display: inline-flex; align-items: center;
              border-radius: 999px; background: #1E4E78; color: #ffffff; font-size: 14px; font-weight: 600;
              text-decoration: none; }}
  a.action:hover {{ background: #163b5c; }}
  a.action:focus-visible {{ outline: none; box-shadow: 0 0 0 3px rgba(191, 139, 58, .25); }}
</style>
</head>
<body>
<main>
  <div class="brand"><img src="{mark}" alt="" /><span>Qubera</span></div>
  <div class="copy">
    <p class="eyebrow">{eyebrow}</p>
    <h1>{heading}</h1>
    <p class="message">{message}</p>
  </div>
  <a class="action" href="{href}">{action}</a>
</main>
</body>
</html>
"""


def page(
    *,
    title: str,
    eyebrow: str,
    heading: str,
    message: str,
    action: str,
    href: str,
    mark_src: str,
    head: str = "",
) -> str:
    """One card on the paper background, with the Qubera mark and wordmark,
    in the design system's type and colors. Self-contained apart from the
    fonts and ``mark_src``, because the wake function serves it while nothing
    under the app's /static/ answers; each program passes a mark it can serve."""
    return _PAGE.format(
        title=escape(title),
        head=head,
        mark=escape(mark_src, quote=True),
        eyebrow=escape(eyebrow),
        heading=escape(heading),
        message=escape(message),
        href=escape(href, quote=True),
        action=escape(action),
    )


def login_page(config: WorkOS, base: str, next_path: Optional[str], mark_src: str) -> str:
    """The page ``/login`` answers with: preview tags for link scrapers, and an
    immediate move to WorkOS for a browser."""
    target = authorize_url(config, base, next_path)
    refresh = f'<meta http-equiv="refresh" content="0; url={escape(target, quote=True)}" />'
    return page(
        title="Qubera — sign in",
        eyebrow="Sign in",
        heading="Taking you to sign in",
        message="The playground is for invited members. Your browser should move on by itself.",
        action="Continue to sign in",
        href=target,
        mark_src=mark_src,
        head=link_preview_tags() + "\n" + refresh,
    )


def refusal_page(email: str, switch_href: str, mark_src: str) -> str:
    who = email or "This account"
    return page(
        title="Qubera — no access",
        eyebrow="Access",
        heading="This account can’t open the playground",
        message=f"{who} is not a member of the Qubera organization. Only members can sign in.",
        action="Sign in with a different account",
        href=switch_href,
        mark_src=mark_src,
    )


def failed_page(mark_src: str) -> str:
    return page(
        title="Qubera — sign-in did not finish",
        eyebrow="Sign in",
        heading="The sign-in didn’t finish",
        message="It could not be confirmed, or it took longer than 15 minutes. Start it again.",
        action="Sign in again",
        href="/login",
        mark_src=mark_src,
    )


def link_preview_tags() -> str:
    """The Open Graph and Twitter Card tags, one per line, for a page's head."""
    image = SITE_URL + LINK_CARD_PATH
    title = escape(LINK_TITLE, quote=True)
    description = escape(LINK_DESCRIPTION, quote=True)
    tags = [
        ("name", "description", description),
        ("property", "og:type", "website"),
        ("property", "og:site_name", "Qubera"),
        ("property", "og:url", SITE_URL + "/"),
        ("property", "og:title", title),
        ("property", "og:description", description),
        ("property", "og:image", image),
        ("property", "og:image:secure_url", image),
        ("property", "og:image:type", "image/png"),
        ("property", "og:image:width", str(LINK_CARD_WIDTH)),
        ("property", "og:image:height", str(LINK_CARD_HEIGHT)),
        ("property", "og:image:alt", title),
        ("name", "twitter:card", "summary_large_image"),
        ("name", "twitter:title", title),
        ("name", "twitter:description", description),
        ("name", "twitter:image", image),
        ("name", "twitter:image:alt", title),
    ]
    return "\n".join(f'<meta {attr}="{key}" content="{value}" />' for attr, key, value in tags)