"""Shared payloads and wire helpers for the status-code follow-up audit cells.

Every cell drives a real proxy against the scripted upstream and reads the outbound provider body back from the
upstream's observation feed, so a cell asserts both what the caller got and what the provider received."""

from __future__ import annotations

import json
import uuid
from collections.abc import Iterator, Mapping, Sequence
from contextlib import contextmanager
from pathlib import Path
from typing import Final

import httpx
from integration._support.client import JSON_OBJECT, Gateway, Scenario, gateway_from_environment, object_value
from integration._support.process import owned_proxy_process
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from pydantic import JsonValue

from tests.integration.cost_calculation.cost_tracking_case import (
    JsonResponse,
    SseResponse,
    StoredResponse,
)

ROUTER_DEFAULTS: Final[dict[str, JsonValue]] = {
    "prompt": "router default prompt",
    "input": "router default input",
    "max_tokens": 32,
}
CHAT: Final[dict[str, JsonValue]] = {
    "id": "chatcmpl-$UNIQUE_ID",
    "object": "chat.completion",
    "created": 1,
    "model": "gpt-4o-mini",
    "choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}],
    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
RESPONSE: Final[dict[str, JsonValue]] = {
    "id": "resp_$UNIQUE_ID",
    "object": "response",
    "created_at": 1,
    "status": "completed",
    "model": "gpt-4o-mini",
    "output": [
        {
            "type": "message",
            "id": "msg_$UNIQUE_ID",
            "status": "completed",
            "role": "assistant",
            "content": [{"type": "output_text", "text": "scripted", "annotations": []}],
        }
    ],
    "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
MESSAGE: Final[dict[str, JsonValue]] = {
    "id": "msg_$UNIQUE_ID",
    "type": "message",
    "role": "assistant",
    "model": "claude-haiku-4-5",
    "content": [{"type": "text", "text": "scripted"}],
    "stop_reason": "end_turn",
    "stop_sequence": None,
    "usage": {"input_tokens": 1, "output_tokens": 1},
}
MODERATION: Final[dict[str, JsonValue]] = {
    "id": "modr-$UNIQUE_ID",
    "model": "omni-moderation-latest",
    "results": [{"flagged": False, "categories": {}, "category_scores": {}}],
}
CHAT_FRAMES: Final = (
    'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
    '"choices":[{"index":0,"delta":{"role":"assistant","content":"streamed "},"finish_reason":null}]}',
    'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
    '"choices":[{"index":0,"delta":{"content":"response"},"finish_reason":null}]}',
    'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
    '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}',
    "data: [DONE]",
)
RESPONSES_FRAMES: Final = (
    "event: response.created\n"
    'data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_$UNIQUE_ID","object":"response",'
    '"created_at":1,"status":"in_progress","model":"gpt-4o-mini","output":[],"usage":null}}',
    "event: response.output_text.delta\n"
    'data: {"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_$UNIQUE_ID","output_index":0,'
    '"content_index":0,"delta":"streamed response"}',
    "event: response.completed\n"
    'data: {"type":"response.completed","sequence_number":2,"response":{"id":"resp_$UNIQUE_ID","object":"response",'
    '"created_at":1,"status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_$UNIQUE_ID",'
    '"status":"completed","role":"assistant","content":[{"type":"output_text","text":"streamed response",'
    '"annotations":[]}]}],"usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}}',
)
MESSAGE_FRAMES: Final = (
    "event: message_start\n"
    'data: {"type":"message_start","message":{"id":"msg_$UNIQUE_ID","type":"message","role":"assistant",'
    '"model":"claude-haiku-4-5","content":[],"stop_reason":null,"stop_sequence":null,'
    '"usage":{"input_tokens":1,"output_tokens":0}}}',
    "event: content_block_start\n"
    'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
    "event: content_block_delta\n"
    'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"streamed response"}}',
    'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}',
    "event: message_delta\n"
    'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}',
    'event: message_stop\ndata: {"type":"message_stop"}',
)
UNSET: Final = object()
USER_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Hello"}]
STREAM_TERMINALS: Final[Mapping[str, str]] = {
    "chat": "data: [DONE]",
    "responses": '"type":"response.completed"',
    "messages": "event: message_stop",
}
ENDPOINT_PATHS: Final[Mapping[str, str]] = {
    "chat": "/v1/chat/completions",
    "responses": "/v1/responses",
    "messages": "/v1/messages",
}
ENDPOINT_BODIES: Final[Mapping[str, dict[str, JsonValue]]] = {
    "chat": {"messages": USER_MESSAGES},
    "responses": {"input": "Hello", "store": False},
    "messages": {"messages": USER_MESSAGES, "max_tokens": 16},
}


def json_response(body: dict[str, JsonValue], status: int = 200) -> JsonResponse:
    return JsonResponse(content_type="application/json", body=body, status=status)


def sse_response(frames: tuple[str, ...]) -> SseResponse:
    return SseResponse(content_type="text/event-stream", frames=frames)


def openai_error(message: str, param: str) -> JsonResponse:
    return json_response(
        {"error": {"message": message, "type": "invalid_request_error", "param": param, "code": None}}, 400
    )


class Upstream:
    def __init__(self, url: str) -> None:
        self.url = url.rstrip("/")

    def reset(self) -> None:
        with httpx.Client(timeout=10, trust_env=False) as client:
            client.delete(f"{self.url}/__observations").raise_for_status()

    def observed(self) -> tuple[dict[str, JsonValue], ...]:
        with httpx.Client(timeout=10, trust_env=False) as client:
            payload: Final = JSON_OBJECT.validate_python(
                client.get(f"{self.url}/__observations?include_method=true").json()
            )
        requests: Final = payload.get("requests")
        assert isinstance(requests, list)
        return tuple(object_value(item) for item in requests if isinstance(item, dict))

    def calls(self, identity: str) -> tuple[dict[str, JsonValue], ...]:
        return tuple(
            item
            for item in self.observed()
            if f"/{identity}/" in str(item.get("path")) and not str(item.get("path", "")).endswith("/models")
        )


def assert_no_provider_call(gateway: Gateway, *identities: str) -> None:
    upstream: Final = Upstream(gateway.upstream_url)
    assert {identity: upstream.calls(identity) for identity in identities} == dict.fromkeys(identities, ())


def one_outbound(gateway: Gateway, identity: str) -> dict[str, JsonValue]:
    calls: Final = Upstream(gateway.upstream_url).calls(identity)
    assert len(calls) == 1, calls
    return object_value(calls[0]["body"])


def post(
    gateway: Gateway,
    path: str,
    body: Mapping[str, JsonValue],
    *,
    key: str | None = None,
    headers: Mapping[str, str] | None = None,
) -> httpx.Response:
    return gateway.client.post(
        path,
        json=dict(body),
        headers={"Authorization": f"Bearer {gateway.key if key is None else key}", **(headers or {})},
        timeout=30,
    )


def stream_lines(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, tuple[str, ...]]:
    with gateway.client.stream(
        "POST", path, json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"}, timeout=30
    ) as response:
        return response.status_code, tuple(response.iter_lines())


def error_body(response: httpx.Response) -> dict[str, JsonValue]:
    """The error object of an OpenAI-shaped or Anthropic-shaped error response."""
    return object_value(JSON_OBJECT.validate_python(response.json())["error"])


def invalid_request(response: httpx.Response) -> dict[str, JsonValue]:
    assert response.status_code == 400, response.text
    error: Final = error_body(response)
    assert error.get("type") == "invalid_request_error", response.text
    return error


def register(scenario: Scenario, prefix: str, response: StoredResponse) -> tuple[str, ScenarioHandle]:
    identity: Final = f"{prefix}-{uuid.uuid4().hex}"
    handle: Final = register_scenario(identity, response)
    scenario.cleanups.callback(delete_scenario, handle)
    return identity, handle


def deployment(
    scenario: Scenario, name: str, model: str, response: StoredResponse, **extra: JsonValue
) -> tuple[str, dict[str, JsonValue]]:
    identity, handle = register(scenario, name, response)
    return identity, {
        "model_name": name,
        "litellm_params": {"model": model, "api_base": handle.api_base(), "api_key": identity, **extra},
    }


@contextmanager
def owned_gateway(
    directory: Path,
    config: Mapping[str, JsonValue],
    *,
    workers: int = 1,
    environment: Mapping[str, str] | None = None,
) -> Iterator[Gateway]:
    """A proxy booted from this checkout with `config` and extra `environment`, sharing the rig's database, Redis
    and upstream."""
    path: Final = directory / "config.yaml"
    path.write_text(json.dumps(dict(config)), encoding="utf-8")
    with (
        gateway_from_environment() as gateway,
        owned_proxy_process(gateway, directory, environment or {}, config=path, workers=workers) as owned,
    ):
        yield owned.gateway


def stream_finished(endpoint: str, lines: Sequence[str]) -> bool:
    """Whether the stream reached the endpoint's terminal frame."""
    return any(STREAM_TERMINALS[endpoint] in line for line in lines)


def _frame_text(frame: Mapping[str, JsonValue]) -> str:
    """The text one chat, responses or messages delta frame carries, or "" for any other frame."""
    choices: Final = frame.get("choices")
    first: Final = choices[0] if isinstance(choices, list) and choices else None
    chat_delta: Final = first.get("delta") if isinstance(first, dict) else None
    if isinstance(chat_delta, dict) and isinstance(chat_delta.get("content"), str):
        return str(chat_delta["content"])
    delta: Final = frame.get("delta")
    if frame.get("type") == "response.output_text.delta" and isinstance(delta, str):
        return delta
    if frame.get("type") == "content_block_delta" and isinstance(delta, dict) and delta.get("type") == "text_delta":
        return str(delta.get("text"))
    return ""


def assembled_text(lines: Sequence[str]) -> str:
    """The text a chat, responses or messages stream carried, joined across its delta frames."""
    return "".join(
        _frame_text(JSON_OBJECT.validate_json(line.removeprefix("data: ")))
        for line in lines
        if line.startswith("data: ") and line != "data: [DONE]"
    )