"""Replay a logged prompt against an Azure OpenAI chat deployment.

Step 7 of PB-RAG-02-03 asks you to replay the prompt that the application sent, with the same
retrieved chunks, and then to replay it with the previous prompt or model version. This script
sends one set of messages to one deployment and prints the answer. Run it once for each variant
and compare the answers.

It is a test only: it calls the model and nothing else. It signs in with Microsoft Entra ID and
never uses an API key. The identity needs the Cognitive Services OpenAI User role on the Azure
OpenAI resource. Every call spends tokens: check the current price before you run many cases.

Usage:
  python replay_prompt.py --messages logged_messages.json [options]

The messages file is the chat messages that the application really sent, taken from its logs. It
holds a JSON list, or an object with a "messages" list:
  [{"role": "system", "content": "..."}, {"role": "user", "content": "..."}]
The messages must already contain the retrieved chunks, exactly as they were logged.

Variants to compare
  - Logged prompt, current deployment:   --messages logged.json
  - Previous prompt:                     --messages previous_prompt_logged.json
  - Previous model version:              --messages logged.json --deployment ${PREVIOUS_DEPLOYMENT}

Environment variables:
  MODEL_ENDPOINT    https://${OPENAI_RESOURCE}.openai.azure.com
  MODEL_DEPLOYMENT  the deployment to call (override it with --deployment)
"""
import argparse
import json
import os
import sys

import requests
from azure.identity import DefaultAzureCredential

SCOPE = "https://ai.azure.com/.default"  # token audience used in the Azure OpenAI Entra ID how-to


def get_token():
    return DefaultAzureCredential().get_token(SCOPE).token


def load_messages(path):
    with open(path, encoding="utf-8-sig") as handle:
        data = json.load(handle)
    messages = data.get("messages") if isinstance(data, dict) else data
    if not isinstance(messages, list) or not messages:
        sys.exit(f"'{path}' must hold a non-empty list of chat messages.")
    for position, message in enumerate(messages, start=1):
        if not isinstance(message, dict) or "role" not in message or "content" not in message:
            sys.exit(f"Message {position} in '{path}' needs 'role' and 'content'.")
    return messages


def replay(endpoint, deployment, messages, temperature, token):
    """Send the messages to the deployment and return the response JSON."""
    url = f"{endpoint}/openai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    body = {"model": deployment, "messages": messages}
    if temperature is not None:
        body["temperature"] = temperature
    resp = requests.post(url, headers=headers, json=body, timeout=120)
    if not resp.ok:
        sys.exit(f"Replay failed: HTTP {resp.status_code}: {resp.text[:500]}")
    return resp.json()


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--messages", required=True, help="JSON file with the logged chat messages")
    parser.add_argument("--deployment", help="deployment to call (default: MODEL_DEPLOYMENT)")
    parser.add_argument(
        "--temperature",
        type=float,
        help="temperature to send; leave it out to use the deployment's default",
    )
    parser.add_argument("--expect", help="text from the expected passage; reports if the answer has it")
    args = parser.parse_args()

    try:
        endpoint = os.environ["MODEL_ENDPOINT"].rstrip("/")
        deployment = args.deployment or os.environ["MODEL_DEPLOYMENT"]
    except KeyError as missing:
        sys.exit(f"Set the environment variable {missing}, or pass --deployment")

    result = replay(endpoint, deployment, load_messages(args.messages), args.temperature, get_token())
    answer = result["choices"][0]["message"]["content"]
    print(f"Deployment: {deployment}")
    print("Answer:")
    print(answer)
    usage = result.get("usage") or {}
    if usage.get("total_tokens") is not None:
        print(f"Tokens used: {usage['total_tokens']}")
    if args.expect:
        found = args.expect.lower() in (answer or "").lower()
        print(f"Answer contains the expected text: {'yes' if found else 'no'}")


if __name__ == "__main__":
    main()