import base64
import hashlib
import hmac
import json
from typing import Any, Dict

from config import settings

PROTOCOL = "gradllm-hmac-sha256-v1"


def _key() -> bytes:
    value = getattr(settings, "GradLlmMessageHmacKey", None)
    if not isinstance(value, str) or len(value.encode("utf-8")) < 32:
        raise RuntimeError("GradLlmMessageHmacKey must contain at least 32 UTF-8 bytes")
    return value.encode("utf-8")


def _signature(direction: str, exchange: str, routing_key: str, envelope: Dict[str, Any]) -> str:
    fields = (
        PROTOCOL,
        direction,
        exchange,
        routing_key,
        envelope["reply_key"],
        str(envelope["sequence"]),
        "1" if envelope["final"] else "0",
        envelope["payload_base64"],
    )
    digest = hmac.new(_key(), "\n".join(fields).encode("utf-8"), hashlib.sha256).digest()
    return base64.b64encode(digest).decode("ascii")


def verify_request(envelope: Dict[str, Any], exchange: str, routing_key: str) -> Dict[str, Any]:
    if not isinstance(envelope, dict) or envelope.get("protocol") != PROTOCOL:
        raise ValueError("missing GradLLM HMAC envelope")
    if envelope.get("sequence") != 0 or envelope.get("final") is not False:
        raise ValueError("invalid GradLLM request envelope")
    expected = _signature("request", exchange, routing_key, envelope)
    if not hmac.compare_digest(str(envelope.get("signature", "")), expected):
        raise ValueError("invalid GradLLM request HMAC")
    payload = base64.b64decode(envelope["payload_base64"], validate=True)
    decoded = json.loads(payload.decode("utf-8"))
    if not isinstance(decoded, dict):
        raise ValueError("GradLLM request payload must be an object")
    decoded["reply_key"] = envelope["reply_key"]
    return decoded


def create_response(exchange: str, reply_key: str, sequence: int, payload: Any, final: bool = False) -> Dict[str, Any]:
    payload_bytes = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    envelope = {
        "protocol": PROTOCOL,
        "reply_key": reply_key,
        "sequence": sequence,
        "final": final,
        "payload_base64": base64.b64encode(payload_bytes).decode("ascii"),
    }
    envelope["signature"] = _signature("response", exchange, reply_key, envelope)
    return envelope