"""Offline HTTP replay for jevchat's choice-to-text loop."""

import json

import httpx

from jevchat.alphabet import Alphabet, Symbol
from jevchat.client import JevClient
from jevchat.config import Config
from jevchat.generate import Done, Step, generate


alphabet = Alphabet(
    name="two-letter",
    description="synthetic alphabet",
    symbols=(Symbol("a", "a"), Symbol("b", "b")),
)
requests = []
responses = []


def answer(request: httpx.Request) -> httpx.Response:
    body = json.loads(request.content)
    requests.append(body)
    prefix = body["state"]["answer_so_far"]
    desired_label = {"": "a", "a": "ab", "ab": "ab"}[prefix]
    labels = body["questions"]["next0"]["criteria"]
    assert desired_label in labels
    probabilities = {label: float(label == desired_label) for label in labels}
    payload = {
            "answers": {
                "next0": {
                    "type": "choice",
                    "choice": desired_label,
                    "confidence": 1.0,
                    "probabilities": probabilities,
                }
            },
            "usage": {"input_tokens": 7, "output_tokens": 2},
    }
    responses.append(payload)
    return httpx.Response(200, json=payload)


transport = httpx.MockTransport(answer)
http_client = httpx.Client(transport=transport)
client = JevClient("offline-key", client=http_client, max_retries=0)
config = Config(
    strategy="choice",
    presentation="hypothesis",
    shuffle_criteria=False,
    temperature=0.0,
    min_steps=0,
    max_steps=5,
    stop_bias=1.0,
    repetition_penalty=1.0,
)
events = list(generate(client, alphabet, config, "synthetic question"))
steps = [event for event in events if isinstance(event, Step)]
done = next(event for event in events if isinstance(event, Done))
print("POST /v1/systemone (MockTransport)")
print("first_response=" + json.dumps(responses[0], sort_keys=True))
for body in requests:
    print(
        json.dumps(
            {
                "answer_so_far": body["state"]["answer_so_far"],
                "criteria": list(body["questions"]["next0"]["criteria"]),
            },
            ensure_ascii=False,
        )
    )
print(
    json.dumps(
        {
            "step_keys": [step.key for step in steps],
            "text": done.text,
            "reason": done.reason,
            "http_calls": client.usage.calls,
            "input_tokens": client.usage.input_tokens,
        }
    )
)
assert [step.key for step in steps] == ["a", "b", "STOP"]
assert done.text == "ab" and done.reason == "stop"