"""
Demo orchestrator (T6)  --  Fusion-less dry-run entrypoint for the UI.

Fusion is shared SaaS and configured by a human, so for a self-contained demo this
tiny endpoint stands in for Fusion's orchestration ONLY: it runs the full round trip
server-side so the static UI has one URL to call. In production Fusion does this; we
never build/host Fusion itself.

Chain (mirrors the 8 steps):
  {query}
    -> POST /tokenize        -> {deidentified_prompt, session_id}   (identity leaves as tokens)
    -> InvokeAgentRuntime    -> token-only talking points           (cloud sees tokens only)
    -> POST /restore         -> {final}                             (identity restored on-prem)
  return {final, deidentified_prompt, agent_tokenized, session_id}

The UI renders `final` (advisor view) vs `deidentified_prompt`+`agent_tokenized`
(what the cloud actually saw). Stdlib + botocore (bundled in Lambda) only.
"""
import json
import os
import secrets
import urllib.parse
import urllib.request

from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.session import Session

REGION = os.environ["REGION"]
TOKENIZE_URL = os.environ["TOKENIZE_URL"]
RESTORE_URL = os.environ["RESTORE_URL"]
API_KEY = os.environ.get("TOKENIZE_API_KEY", "")
RUNTIME_ARN = os.environ.get("AGENT_RUNTIME_ARN", "")

_creds = Session().get_credentials()


def _post_json(url, obj):
    req = urllib.request.Request(
        url, data=json.dumps(obj).encode(),
        headers={"Content-Type": "application/json", "x-api-key": API_KEY}, method="POST",
    )
    with urllib.request.urlopen(req, timeout=50) as r:
        return json.loads(r.read().decode())


def _invoke_runtime(prompt):
    """SigV4-signed InvokeAgentRuntime data-plane call; returns the agent's text."""
    url = (f"https://bedrock-agentcore.{REGION}.amazonaws.com"
           f"/runtimes/{urllib.parse.quote(RUNTIME_ARN, safe='')}/invocations?qualifier=DEFAULT")
    body = json.dumps({"prompt": prompt}).encode()
    headers = {
        "Content-Type": "application/json",
        "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": secrets.token_hex(20),
    }
    aws_req = AWSRequest(method="POST", url=url, data=body, headers=headers)
    SigV4Auth(_creds, "bedrock-agentcore", REGION).add_auth(aws_req)
    req = urllib.request.Request(url, data=body, headers=dict(aws_req.headers), method="POST")
    with urllib.request.urlopen(req, timeout=120) as r:
        raw = r.read().decode()
    try:
        obj = json.loads(raw)
        return obj.get("result", raw) if isinstance(obj, dict) else raw
    except json.JSONDecodeError:
        return raw


def _reply(status, obj):
    return {
        "statusCode": status,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(obj, ensure_ascii=False),
    }


def lambda_handler(event, _context):
    try:
        payload = json.loads(event.get("body") or "{}") if "body" in event else event
        query = (payload.get("query") or "").strip()
        if not query:
            return _reply(400, {"error": "query is required"})

        tok = _post_json(TOKENIZE_URL, {"query": query})
        deid = tok["deidentified_prompt"]
        session_id = tok["session_id"]

        # Cloud reasoning on tokens only (skip gracefully if no runtime configured).
        agent_tokenized = _invoke_runtime(deid) if RUNTIME_ARN else deid

        # Restore identity into the agent's answer (envelope + inline).
        res = _post_json(RESTORE_URL, {"session_id": session_id, "text": agent_tokenized})

        return _reply(200, {
            "final": res["final"],
            "deidentified_prompt": deid,
            "agent_tokenized": agent_tokenized,
            "session_id": session_id,
        })
    except Exception as e:
        return _reply(502, {"error": f"{type(e).__name__}: {e}"})