"""Proof parsing across the four wire encodings (Spec §6)."""

from __future__ import annotations

import json
from typing import Any

from .encodings import extract_markdown_proof, from_compact
from .envelope import COMPACT_PROOF_PREFIX

_B64URL = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_")


def extract_compact_token(text: str) -> str | None:
    idx = text.find(COMPACT_PROOF_PREFIX)
    if idx < 0:
        return None
    body_chars = []
    for ch in text[idx + len(COMPACT_PROOF_PREFIX) :]:
        if ch in _B64URL:
            body_chars.append(ch)
        else:
            break
    if not body_chars:
        return None
    return COMPACT_PROOF_PREFIX + "".join(body_chars)


def parse_proof(raw: str) -> dict[str, Any]:
    trimmed = raw.strip()

    # Markdown fence is the most specific marker — check it first.
    if "```kez" in trimmed:
        return extract_markdown_proof(trimmed)
    # Raw JSON envelope.
    if trimmed.startswith("{"):
        return json.loads(trimmed)
    # Compact: extract the kez:z1:<base64url> token anywhere in the input.
    token = extract_compact_token(trimmed)
    if token is not None:
        return from_compact(token)
    raise ValueError("unknown KEZ proof format")