"""Wire encodings: JSON, compact (kez:z1:), markdown, DNS (Spec §6)."""

from __future__ import annotations

import base64
import json
from typing import Any

import zstandard

from .envelope import COMPACT_CHAIN_PREFIX, COMPACT_PROOF_PREFIX

_ZSTD_LEVEL = 3


def _b64url_nopad(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(s: str) -> bytes:
    pad = "=" * (-len(s) % 4)
    return base64.urlsafe_b64decode(s + pad)


def _zstd_compress(data: bytes) -> bytes:
    return zstandard.ZstdCompressor(level=_ZSTD_LEVEL).compress(data)


def _zstd_decompress(data: bytes) -> bytes:
    # decompressobj handles frames that omit the content-size header, which
    # some encoders (e.g. Node's zstd) produce.
    dobj = zstandard.ZstdDecompressor().decompressobj()
    return dobj.decompress(data) + dobj.flush()


def to_pretty_json(envelope: dict[str, Any]) -> str:
    return json.dumps(envelope, indent=2, ensure_ascii=False)


def to_compact_json(envelope: dict[str, Any]) -> str:
    return json.dumps(envelope, separators=(",", ":"), ensure_ascii=False)


def to_compact(envelope: dict[str, Any]) -> str:
    raw = to_compact_json(envelope).encode("utf-8")
    return COMPACT_PROOF_PREFIX + _b64url_nopad(_zstd_compress(raw))


def from_compact(value: str) -> dict[str, Any]:
    trimmed = value.strip()
    if not trimmed.startswith(COMPACT_PROOF_PREFIX):
        raise ValueError("compact proof missing kez:z1: prefix")
    body = trimmed[len(COMPACT_PROOF_PREFIX) :]
    raw = _zstd_decompress(_b64url_decode(body))
    return json.loads(raw.decode("utf-8"))


def to_markdown(envelope: dict[str, Any]) -> str:
    payload = envelope["payload"]
    return (
        "# KEZ Proof\n\n"
        "This account publishes a signed KEZ identity claim.\n\n"
        f"- Primary: `{payload['primary']}`\n"
        f"- Subject: `{payload['subject']}`\n"
        f"- Created: `{payload['created_at']}`\n\n"
        "```kez\n"
        f"{to_pretty_json(envelope)}\n"
        "```\n"
    )


def extract_markdown_proof(markdown: str) -> dict[str, Any]:
    fence = "```kez"
    start = markdown.find(fence)
    if start < 0:
        raise ValueError("missing ```kez proof block")
    body_start = start + len(fence)
    end = markdown.find("```", body_start)
    if end < 0:
        raise ValueError("unterminated ```kez proof block")
    return json.loads(markdown[body_start:end].strip())


# ── Sigchain compact bundle (kez:zc1:) ──────────────────────────────────────


def chain_to_jsonl(events: list[dict[str, Any]]) -> str:
    if not events:
        return ""
    return "\n".join(json.dumps(e, separators=(",", ":"), ensure_ascii=False) for e in events) + "\n"


def chain_from_jsonl(text: str) -> list[dict[str, Any]]:
    return [json.loads(line) for line in text.splitlines() if line.strip()]


def chain_to_compact_bundle(events: list[dict[str, Any]]) -> str:
    raw = chain_to_jsonl(events).encode("utf-8")
    return COMPACT_CHAIN_PREFIX + _b64url_nopad(_zstd_compress(raw))


def chain_from_compact_bundle(value: str) -> list[dict[str, Any]]:
    trimmed = value.strip()
    if not trimmed.startswith(COMPACT_CHAIN_PREFIX):
        raise ValueError("compact chain missing kez:zc1: prefix")
    body = trimmed[len(COMPACT_CHAIN_PREFIX) :]
    raw = _zstd_decompress(_b64url_decode(body))
    return chain_from_jsonl(raw.decode("utf-8"))


# ── DNS TXT helpers ─────────────────────────────────────────────────────────


def dns_txt_name(subject) -> str:
    from .identity import Identity

    ident = subject if isinstance(subject, Identity) else Identity.parse(str(subject))
    if ident.scheme != "dns":
        raise ValueError("DNS TXT proof requires a dns: subject")
    return f"_kez.{ident.value}"


def quote_dns_txt_value(value: str) -> str:
    chunks = [value[i : i + 240] for i in range(0, len(value), 240)]
    quoted = []
    for chunk in chunks:
        escaped = chunk.replace("\\", "\\\\").replace('"', '\\"')
        quoted.append(f'"{escaped}"')
    return " ".join(quoted)