#!/usr/bin/env python3
"""
DVeProto — эталон для разработчиков (standalone site, ./ рядом с index.php)

v1.0: после рукопожатия — JSON text frames (type: dve).
v1.1: после рукопожатия — binary frames (без JSON/Base64).

Рукопожатие (dve_hello / dve_client_ack) всегда JSON, совместимо с v1.
Крипто: X25519 + HKDF-SHA256 + AES-256-GCM (метки DVeProto-v1/c2s|s2c).

Зависимость: pip install cryptography
"""
from __future__ import annotations

import base64
import hashlib
import json
import os
import struct
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

PROTO_NAME = "DVeProto"
PROTO_HANDSHAKE_VERSION = 1
PROTO_VERSION = PROTO_HANDSHAKE_VERSION  # alias
WIRE_V10 = "1.0"
WIRE_V11 = "1.1"
DEFAULT_OFFER = (WIRE_V10, WIRE_V11)

INFO_C2S = b"DVeProto-v1/c2s"
INFO_S2C = b"DVeProto-v1/s2c"

BIN_VERSION = 0x11
HEADER_LEN = 14
GCM_TAG_LEN = 16
NONCE_LEN = 12

PTYPE_APP_JSON = 0x01
PTYPE_APP_BIN = 0x02
PTYPE_TEXT = 0x03
PTYPE_FILE_BEGIN = 0x10
PTYPE_FILE_CHUNK = 0x11
PTYPE_FILE_END = 0x12
PTYPE_FILE_ABORT = 0x13
PTYPE_FILE_ACK = 0x14
PTYPE_FILE_RESUME = 0x15

FILE_FLAG_SHA256 = 0x01
DEFAULT_CHUNK_SIZE = 256 * 1024

ACK_OK = 0
ACK_NEED_RESUME = 1
ACK_COMPLETE = 2
ACK_ERROR = 3


def _b64e(raw: bytes) -> str:
    return base64.b64encode(raw).decode("ascii")


def _b64d(s: str) -> bytes:
    return base64.b64decode(s.encode("ascii"), validate=True)


def derive_aes_keys(shared_secret: bytes) -> Tuple[bytes, bytes]:
    hkdf_c2s = HKDF(algorithm=hashes.SHA256(), length=32, salt=b"", info=INFO_C2S)
    hkdf_s2c = HKDF(algorithm=hashes.SHA256(), length=32, salt=b"", info=INFO_S2C)
    return hkdf_c2s.derive(shared_secret), hkdf_s2c.derive(shared_secret)


def normalize_wire(value: Any) -> Optional[str]:
    if value is None:
        return None
    s = str(value).strip()
    if s in ("1", "1.0", "v1", "v1.0", "json"):
        return WIRE_V10
    if s in ("1.1", "v1.1", "bin", "bin11", "binary"):
        return WIRE_V11
    return None


def pick_wire(offer: Iterable[str], select: Optional[str]) -> str:
    offered = set()
    for item in offer:
        w = normalize_wire(item)
        if w:
            offered.add(w)
    if not offered:
        offered.add(WIRE_V10)
    chosen = normalize_wire(select) if select is not None else None
    if chosen and chosen in offered:
        return chosen
    return WIRE_V10


@dataclass(frozen=True)
class ClientAck:
    client_pk: bytes
    select: str


@dataclass(frozen=True)
class DecodedFrame:
    ptype: int
    plaintext: bytes

    def as_json(self) -> Dict[str, Any]:
        if self.ptype != PTYPE_APP_JSON:
            raise ValueError("frame is not APP_JSON")
        return json.loads(self.plaintext.decode("utf-8"))

    def as_text(self) -> str:
        if self.ptype != PTYPE_TEXT:
            raise ValueError("frame is not TEXT")
        return self.plaintext.decode("utf-8")


def pack_file_begin(
    transfer_id: bytes,
    *,
    file_size: int,
    chunk_size: int,
    total_chunks: int,
    name: str,
    mime: str = "application/octet-stream",
    sha256: Optional[bytes] = None,
) -> bytes:
    if len(transfer_id) != 16:
        raise ValueError("transfer_id must be 16 bytes")
    name_b = name.encode("utf-8")
    mime_b = mime.encode("utf-8")
    flags = FILE_FLAG_SHA256 if sha256 is not None else 0
    parts = [
        transfer_id,
        struct.pack(">QII", int(file_size), int(chunk_size), int(total_chunks)),
        bytes([flags]),
        struct.pack(">H", len(name_b)),
        name_b,
        struct.pack(">H", len(mime_b)),
        mime_b,
    ]
    if sha256 is not None:
        if len(sha256) != 32:
            raise ValueError("sha256 must be 32 bytes")
        parts.append(sha256)
    return b"".join(parts)


def unpack_file_begin(pt: bytes) -> Dict[str, Any]:
    o = 0
    transfer_id = pt[o : o + 16]
    o += 16
    file_size, chunk_size, total_chunks = struct.unpack_from(">QII", pt, o)
    o += 16
    flags = pt[o]
    o += 1
    (name_len,) = struct.unpack_from(">H", pt, o)
    o += 2
    name = pt[o : o + name_len].decode("utf-8")
    o += name_len
    (mime_len,) = struct.unpack_from(">H", pt, o)
    o += 2
    mime = pt[o : o + mime_len].decode("utf-8")
    o += mime_len
    sha256 = pt[o : o + 32] if (flags & FILE_FLAG_SHA256) else None
    return {
        "transfer_id": transfer_id,
        "file_size": file_size,
        "chunk_size": chunk_size,
        "total_chunks": total_chunks,
        "flags": flags,
        "name": name,
        "mime": mime,
        "sha256": sha256,
    }


def pack_file_chunk(transfer_id: bytes, chunk_index: int, data: bytes) -> bytes:
    return transfer_id + struct.pack(">I", int(chunk_index)) + data


def unpack_file_chunk(pt: bytes) -> Dict[str, Any]:
    return {
        "transfer_id": pt[:16],
        "chunk_index": struct.unpack_from(">I", pt, 16)[0],
        "data": pt[20:],
    }


def pack_file_end(transfer_id: bytes, sha256: bytes) -> bytes:
    return transfer_id + sha256


def unpack_file_end(pt: bytes) -> Dict[str, Any]:
    return {"transfer_id": pt[:16], "sha256": pt[16:48]}


def pack_file_abort(transfer_id: bytes, reason_code: int = 0) -> bytes:
    return transfer_id + bytes([reason_code & 0xFF])


def unpack_file_abort(pt: bytes) -> Dict[str, Any]:
    return {"transfer_id": pt[:16], "reason_code": pt[16]}


def pack_file_ack(transfer_id: bytes, status: int, last_ok: int) -> bytes:
    return transfer_id + bytes([status & 0xFF]) + struct.pack(">I", int(last_ok) & 0xFFFFFFFF)


def unpack_file_ack(pt: bytes) -> Dict[str, Any]:
    return {
        "transfer_id": pt[:16],
        "status": pt[16],
        "last_ok": struct.unpack_from(">I", pt, 17)[0],
    }


def pack_file_resume(transfer_id: bytes, from_chunk: int) -> bytes:
    return transfer_id + struct.pack(">I", int(from_chunk) & 0xFFFFFFFF)


def unpack_file_resume(pt: bytes) -> Dict[str, Any]:
    return {
        "transfer_id": pt[:16],
        "from_chunk": struct.unpack_from(">I", pt, 16)[0],
    }


def decode_file_payload(ptype: int, plaintext: bytes) -> Dict[str, Any]:
    table = {
        PTYPE_FILE_BEGIN: unpack_file_begin,
        PTYPE_FILE_CHUNK: unpack_file_chunk,
        PTYPE_FILE_END: unpack_file_end,
        PTYPE_FILE_ABORT: unpack_file_abort,
        PTYPE_FILE_ACK: unpack_file_ack,
        PTYPE_FILE_RESUME: unpack_file_resume,
    }
    fn = table.get(ptype)
    if not fn:
        raise ValueError(f"not a file ptype: 0x{ptype:02x}")
    return fn(plaintext)


class DVeSession:
    """Серверная сессия: decrypt c2s, encrypt s2c."""

    __slots__ = ("_c2s", "_s2c", "wire_version")

    def __init__(
        self, c2s_key: bytes, s2c_key: bytes, wire_version: str = WIRE_V10
    ) -> None:
        self._c2s = AESGCM(c2s_key)
        self._s2c = AESGCM(s2c_key)
        self.wire_version = normalize_wire(wire_version) or WIRE_V10

    @property
    def is_binary(self) -> bool:
        return self.wire_version == WIRE_V11

    @classmethod
    def from_server_keys(
        cls,
        server_private: X25519PrivateKey,
        client_public_raw: bytes,
        wire_version: str = WIRE_V10,
    ) -> "DVeSession":
        client_pub = X25519PublicKey.from_public_bytes(client_public_raw)
        shared = server_private.exchange(client_pub)
        c2s, s2c = derive_aes_keys(shared)
        return cls(c2s, s2c, wire_version=wire_version)

    def encrypt_json(self, obj: Dict[str, Any]) -> str:
        nonce = os.urandom(NONCE_LEN)
        pt = json.dumps(obj, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        ct = self._s2c.encrypt(nonce, pt, None)
        return json.dumps(
            {
                "type": "dve",
                "proto": PROTO_NAME,
                "v": PROTO_HANDSHAKE_VERSION,
                "n": _b64e(nonce),
                "c": _b64e(ct),
            },
            separators=(",", ":"),
            ensure_ascii=False,
        )

    def decrypt_json(self, text: str) -> Dict[str, Any]:
        data = json.loads(text)
        if data.get("type") != "dve" or data.get("proto") != PROTO_NAME:
            raise ValueError("invalid DVe frame")
        if int(data.get("v", 0)) != PROTO_HANDSHAKE_VERSION:
            raise ValueError("unsupported DVe version")
        pt = self._c2s.decrypt(_b64d(data["n"]), _b64d(data["c"]), None)
        return json.loads(pt.decode("utf-8"))

    def _seal(self, aead: AESGCM, ptype: int, plaintext: bytes) -> bytes:
        nonce = os.urandom(NONCE_LEN)
        ct = aead.encrypt(nonce, plaintext, None)
        out = bytearray(HEADER_LEN + len(ct))
        out[0] = BIN_VERSION
        out[1] = ptype & 0xFF
        out[2:14] = nonce
        out[14:] = ct
        return bytes(out)

    def _open(self, aead: AESGCM, frame: bytes) -> DecodedFrame:
        raw = bytes(frame)
        if len(raw) < HEADER_LEN + GCM_TAG_LEN or raw[0] != BIN_VERSION:
            raise ValueError("invalid binary DVe frame")
        pt = aead.decrypt(raw[2:14], raw[14:], None)
        return DecodedFrame(ptype=raw[1], plaintext=pt)

    def encrypt_bin(self, ptype: int, plaintext: bytes) -> bytes:
        return self._seal(self._s2c, ptype, plaintext)

    def decrypt_bin(self, frame: bytes) -> DecodedFrame:
        return self._open(self._c2s, frame)

    def encrypt(self, obj: Dict[str, Any]) -> Union[str, bytes]:
        if self.is_binary:
            pt = json.dumps(obj, separators=(",", ":"), ensure_ascii=False).encode(
                "utf-8"
            )
            return self.encrypt_bin(PTYPE_APP_JSON, pt)
        return self.encrypt_json(obj)

    def decrypt(self, message: Union[str, bytes, bytearray]) -> Dict[str, Any]:
        if self.is_binary:
            if isinstance(message, str):
                raise ValueError("DVeProto 1.1 rejects text frames")
            frame = self.decrypt_bin(message)
            if frame.ptype != PTYPE_APP_JSON:
                raise ValueError(f"expected APP_JSON, got 0x{frame.ptype:02x}")
            return frame.as_json()
        if isinstance(message, (bytes, bytearray)):
            raise ValueError("DVeProto 1.0 expects text JSON frames")
        return self.decrypt_json(message)


def server_hello_payload(
    server_private: X25519PrivateKey, offer: Iterable[str] = DEFAULT_OFFER
) -> str:
    pub = server_private.public_key().public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    offer_list: List[str] = []
    for item in offer:
        w = normalize_wire(item)
        if w and w not in offer_list:
            offer_list.append(w)
    if not offer_list:
        offer_list = [WIRE_V10]
    return json.dumps(
        {
            "type": "dve_hello",
            "proto": PROTO_NAME,
            "v": PROTO_HANDSHAKE_VERSION,
            "server_pk": _b64e(pub),
            "offer": offer_list,
        },
        separators=(",", ":"),
        ensure_ascii=False,
    )


def parse_client_ack(
    text: str, offer: Iterable[str] = DEFAULT_OFFER
) -> Optional[ClientAck]:
    try:
        data = json.loads(text)
    except (json.JSONDecodeError, TypeError):
        return None
    if data.get("type") != "dve_client_ack" or data.get("proto") != PROTO_NAME:
        return None
    if int(data.get("v", 0)) != PROTO_HANDSHAKE_VERSION:
        return None
    try:
        client_pk = _b64d(data["client_pk"])
    except (KeyError, ValueError):
        return None
    if len(client_pk) != 32:
        return None
    return ClientAck(client_pk=client_pk, select=pick_wire(offer, data.get("select")))


class DVeSeq:
    __slots__ = ("_n",)

    def __init__(self) -> None:
        self._n = 0

    def next(self) -> int:
        self._n += 1
        return self._n


def dve_app_message(
    op: str,
    body: Optional[Dict[str, Any]] = None,
    *,
    seq: Optional[int] = None,
) -> Dict[str, Any]:
    msg: Dict[str, Any] = {"dve_op": op}
    if seq is not None:
        msg["dve_seq"] = int(seq)
    if body:
        msg.update(body)
    return msg


class DVeClientSession:
    """Клиент: исходящие c2s, входящие s2c."""

    __slots__ = ("_c2s", "_s2c", "wire_version")

    def __init__(
        self, c2s_key: bytes, s2c_key: bytes, wire_version: str = WIRE_V10
    ) -> None:
        self._c2s = AESGCM(c2s_key)
        self._s2c = AESGCM(s2c_key)
        self.wire_version = normalize_wire(wire_version) or WIRE_V10

    @property
    def is_binary(self) -> bool:
        return self.wire_version == WIRE_V11

    @classmethod
    def from_server_hello_text(
        cls, hello_text: str, prefer: str = WIRE_V11
    ) -> Tuple["DVeClientSession", str]:
        data = json.loads(hello_text)
        if data.get("type") != "dve_hello" or data.get("proto") != PROTO_NAME:
            raise ValueError("ожидался dve_hello DVeProto")
        if int(data.get("v", 0)) != PROTO_HANDSHAKE_VERSION:
            raise ValueError("неподдерживаемая версия DVeProto")
        offer = data.get("offer") or [WIRE_V10]
        selected = pick_wire(offer, prefer)
        server_pub = X25519PublicKey.from_public_bytes(_b64d(data["server_pk"]))
        client_priv = X25519PrivateKey.generate()
        shared = client_priv.exchange(server_pub)
        c2s, s2c = derive_aes_keys(shared)
        client_pub = client_priv.public_key().public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )
        ack_obj: Dict[str, Any] = {
            "type": "dve_client_ack",
            "proto": PROTO_NAME,
            "v": PROTO_HANDSHAKE_VERSION,
            "client_pk": _b64e(client_pub),
        }
        if selected == WIRE_V11:
            ack_obj["select"] = WIRE_V11
        elif "offer" in data:
            ack_obj["select"] = WIRE_V10
        ack = json.dumps(ack_obj, separators=(",", ":"), ensure_ascii=False)
        return cls(c2s, s2c, wire_version=selected), ack

    def _seal(self, aead: AESGCM, ptype: int, plaintext: bytes) -> bytes:
        nonce = os.urandom(NONCE_LEN)
        ct = aead.encrypt(nonce, plaintext, None)
        out = bytearray(HEADER_LEN + len(ct))
        out[0] = BIN_VERSION
        out[1] = ptype & 0xFF
        out[2:14] = nonce
        out[14:] = ct
        return bytes(out)

    def _open(self, aead: AESGCM, frame: bytes) -> DecodedFrame:
        raw = bytes(frame)
        if len(raw) < HEADER_LEN + GCM_TAG_LEN or raw[0] != BIN_VERSION:
            raise ValueError("invalid binary DVe frame")
        pt = aead.decrypt(raw[2:14], raw[14:], None)
        return DecodedFrame(ptype=raw[1], plaintext=pt)

    def pack_outgoing(self, obj: Dict[str, Any]) -> Union[str, bytes]:
        pt = json.dumps(obj, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        if self.is_binary:
            return self._seal(self._c2s, PTYPE_APP_JSON, pt)
        nonce = os.urandom(NONCE_LEN)
        ct = self._c2s.encrypt(nonce, pt, None)
        return json.dumps(
            {
                "type": "dve",
                "proto": PROTO_NAME,
                "v": PROTO_HANDSHAKE_VERSION,
                "n": _b64e(nonce),
                "c": _b64e(ct),
            },
            separators=(",", ":"),
            ensure_ascii=False,
        )

    def unpack_incoming(
        self, message: Union[str, bytes, bytearray]
    ) -> Dict[str, Any]:
        if self.is_binary:
            if isinstance(message, str):
                raise ValueError("DVeProto 1.1 rejects text frames")
            frame = self._open(self._s2c, message)
            if frame.ptype != PTYPE_APP_JSON:
                raise ValueError(f"expected APP_JSON, got 0x{frame.ptype:02x}")
            return frame.as_json()
        if isinstance(message, (bytes, bytearray)):
            raise ValueError("DVeProto 1.0 expects text JSON frames")
        data = json.loads(message)
        if data.get("type") != "dve" or data.get("proto") != PROTO_NAME:
            raise ValueError("invalid DVe frame")
        if int(data.get("v", 0)) != PROTO_HANDSHAKE_VERSION:
            raise ValueError("unsupported DVe version")
        pt = self._s2c.decrypt(_b64d(data["n"]), _b64d(data["c"]), None)
        return json.loads(pt.decode("utf-8"))

    def pack_bin(self, ptype: int, plaintext: bytes) -> bytes:
        if not self.is_binary:
            raise ValueError("binary pack requires wire 1.1")
        return self._seal(self._c2s, ptype, plaintext)

    def unpack_bin(self, frame: bytes) -> DecodedFrame:
        if not self.is_binary:
            raise ValueError("binary unpack requires wire 1.1")
        return self._open(self._s2c, frame)

    def pack_file_begin(self, **kwargs: Any) -> bytes:
        return self.pack_bin(PTYPE_FILE_BEGIN, pack_file_begin(**kwargs))

    def pack_file_chunk(
        self, transfer_id: bytes, chunk_index: int, data: bytes
    ) -> bytes:
        return self.pack_bin(
            PTYPE_FILE_CHUNK, pack_file_chunk(transfer_id, chunk_index, data)
        )

    def pack_file_end(self, transfer_id: bytes, sha256: bytes) -> bytes:
        return self.pack_bin(PTYPE_FILE_END, pack_file_end(transfer_id, sha256))

    def pack_file_resume(self, transfer_id: bytes, from_chunk: int) -> bytes:
        return self.pack_bin(
            PTYPE_FILE_RESUME, pack_file_resume(transfer_id, from_chunk)
        )


def developer_server_session_from_ack(
    server_private: X25519PrivateKey, client_ack_text: str
) -> DVeSession:
    ack = parse_client_ack(client_ack_text)
    if ack is None:
        raise ValueError("некорректный dve_client_ack")
    return DVeSession.from_server_keys(
        server_private, ack.client_pk, wire_version=ack.select
    )


def iter_file_chunks(
    data: bytes, chunk_size: int = DEFAULT_CHUNK_SIZE
) -> List[bytes]:
    if chunk_size <= 0:
        raise ValueError("chunk_size must be > 0")
    return [data[i : i + chunk_size] for i in range(0, len(data), chunk_size)] or [b""]


def file_transfer_plan(
    data: bytes,
    name: str,
    mime: str = "application/octet-stream",
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    transfer_id: Optional[bytes] = None,
) -> Dict[str, Any]:
    """Метаданные + чанки для отправки через pack_file_*."""
    tid = transfer_id or os.urandom(16)
    chunks = iter_file_chunks(data, chunk_size)
    digest = hashlib.sha256(data).digest()
    return {
        "transfer_id": tid,
        "file_size": len(data),
        "chunk_size": chunk_size,
        "total_chunks": len(chunks),
        "name": name,
        "mime": mime,
        "sha256": digest,
        "chunks": chunks,
    }