"""Strict generic Native Messaging wire validation (stdlib-only)."""

from __future__ import annotations

import json
import re
import struct
import time
import unicodedata
from datetime import datetime
from typing import Any, BinaryIO, Iterable

from .constants import (
    AUDIT_ID_RE,
    COOKIE_ACCESS_REASONS,
    EXTENSION_BUILD,
    HOST_BUILD,
    JOB_ID_RE,
    HANDOFF_ID_RE,
    MESSAGE_ID_RE,
    PRESTART_ABORT_CODES,
    MAX_COOKIE_COUNT,
    MAX_INPUT_FRAME,
    MAX_OUTPUT_FRAME,
    MAX_SAFE_INTEGER,
    SCHEMA_VERSION,
    canonical_url,
    duration_tolerance_ms,
    stable_job_id,
    stable_successor_job_id,
    UPPER_SHA256_RE,
    target_path,
    validate_bvid,
    validate_creator_uid,
    validate_prepareless_terminal,
)

_NONCE_RE = re.compile(r"[0-9a-f]{32}\Z")
_PREPARE_ID_RE = re.compile(r"[0-9a-f]{32}\Z")
_LEASE_ID_RE = re.compile(r"[0-9a-f]{32}\Z")
_RELOAD_TOKEN_RE = re.compile(r"[0-9a-f]{32}\Z")
_STORE_RE = re.compile(r"[0-9]{1,8}\Z")
_PARENT_RE = re.compile(r"--parent-window=[0-9]+\Z")
_COOKIE_SAME_SITE = {"no_restriction", "lax", "strict", "unspecified"}


class ProtocolError(Exception):
    def __init__(self, code: str = "E_PROTOCOL") -> None:
        super().__init__(code)
        self.code = code


def _reject_constant(_: str) -> None:
    raise ProtocolError()


def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    result: dict[str, Any] = {}
    for key, value in pairs:
        if key in result:
            raise ProtocolError()
        result[key] = value
    return result


def strict_json_loads(payload: bytes) -> dict[str, Any]:
    try:
        value = json.loads(
            payload.decode("utf-8", errors="strict"),
            object_pairs_hook=_unique_object,
            parse_constant=_reject_constant,
        )
    except ProtocolError:
        raise
    except (UnicodeError, json.JSONDecodeError) as exc:
        raise ProtocolError() from exc
    if not isinstance(value, dict):
        raise ProtocolError()
    return value


def encode_json(value: dict[str, Any], limit: int = MAX_OUTPUT_FRAME) -> bytes:
    payload = json.dumps(
        value, ensure_ascii=True, allow_nan=False, separators=(",", ":"), sort_keys=True
    ).encode("utf-8")
    if len(payload) > limit:
        raise ProtocolError()
    return payload


def read_frame(stream: BinaryIO, limit: int = MAX_INPUT_FRAME) -> bytes | None:
    header = stream.read(4)
    if header == b"":
        return None
    if len(header) != 4:
        raise ProtocolError()
    (length,) = struct.unpack("<I", header)
    if length == 0 or length > limit:
        raise ProtocolError()
    payload = stream.read(length)
    if len(payload) != length:
        raise ProtocolError()
    return payload


def write_frame(stream: BinaryIO, value: dict[str, Any]) -> None:
    payload = encode_json(value)
    stream.write(struct.pack("<I", len(payload)))
    stream.write(payload)
    stream.flush()


def _exact_keys(value: dict[str, Any], expected: Iterable[str]) -> None:
    if set(value) != set(expected):
        raise ProtocolError()


def _integer(value: Any, minimum: int, maximum: int) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum:
        raise ProtocolError()
    return value


def _boolean(value: Any) -> bool:
    if not isinstance(value, bool):
        raise ProtocolError()
    return value


def _safe_string(value: Any, minimum: int, maximum_utf8: int) -> str:
    if not isinstance(value, str):
        raise ProtocolError()
    encoded = value.encode("utf-8")
    if len(encoded) < minimum or len(encoded) > maximum_utf8:
        raise ProtocolError()
    if any(ch in "\x00\t\r\n" or unicodedata.category(ch).startswith("C") for ch in value):
        raise ProtocolError()
    return value


def _schema_type(value: dict[str, Any], message_type: str, keys: Iterable[str]) -> None:
    _exact_keys(value, keys)
    if _integer(value.get("schema"), SCHEMA_VERSION, SCHEMA_VERSION) != SCHEMA_VERSION:
        raise ProtocolError()
    if value.get("type") != message_type:
        raise ProtocolError()


def validate_origin_argv(arguments: list[str], expected_origin: str) -> None:
    if len(arguments) not in (1, 2) or arguments[0] != expected_origin:
        raise ProtocolError("E_ORIGIN")
    if len(arguments) == 2 and not _PARENT_RE.fullmatch(arguments[1]):
        raise ProtocolError("E_ORIGIN")


def validate_job(value: Any) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ProtocolError("E_JOB")
    base_keys = {"job_id", "bvid", "creator_uid", "canonical_url", "expected_duration_ms", "published_at", "title"}
    successor = "lineage" in value
    _exact_keys(value, base_keys | ({"lineage"} if successor else set()))
    try:
        bvid = validate_bvid(value["bvid"])
        creator = validate_creator_uid(value["creator_uid"])
    except ValueError as exc:
        raise ProtocolError("E_JOB") from exc
    if not isinstance(value["job_id"], str) or not JOB_ID_RE.fullmatch(value["job_id"]):
        raise ProtocolError("E_JOB")
    if value["canonical_url"] != canonical_url(bvid):
        raise ProtocolError("E_JOB")
    if successor:
        lineage = value["lineage"]
        lineage_keys = {
            "predecessor_job_id", "retry_generation", "predecessor_terminal_error_code",
            "authorization_message_id", "authorization_handoff_id", "authorization_sha256",
            "repair_review_result_message_id", "repair_audit_id", "repair_audit_bytes",
            "repair_audit_sha256",
        }
        if not isinstance(lineage, dict):
            raise ProtocolError("E_JOB")
        _exact_keys(lineage, lineage_keys)
        try:
            expected_job_id = stable_successor_job_id(
                creator, bvid, lineage["predecessor_job_id"], lineage["retry_generation"],
                lineage["predecessor_terminal_error_code"], lineage["authorization_message_id"],
                lineage["authorization_handoff_id"], lineage["authorization_sha256"],
                lineage["repair_review_result_message_id"], lineage["repair_audit_id"],
                lineage["repair_audit_bytes"], lineage["repair_audit_sha256"],
            )
        except (KeyError, ValueError) as exc:
            raise ProtocolError("E_JOB") from exc
        if (
            not JOB_ID_RE.fullmatch(lineage["predecessor_job_id"])
            or not MESSAGE_ID_RE.fullmatch(lineage["authorization_message_id"])
            or not HANDOFF_ID_RE.fullmatch(lineage["authorization_handoff_id"])
            or not UPPER_SHA256_RE.fullmatch(lineage["authorization_sha256"])
            or not MESSAGE_ID_RE.fullmatch(lineage["repair_review_result_message_id"])
            or not AUDIT_ID_RE.fullmatch(lineage["repair_audit_id"])
            or not UPPER_SHA256_RE.fullmatch(lineage["repair_audit_sha256"])
        ):
            raise ProtocolError("E_JOB")
    else:
        expected_job_id = stable_job_id(creator, bvid)
    if value["job_id"] != expected_job_id:
        raise ProtocolError("E_JOB")
    _integer(value["expected_duration_ms"], 1_000, 86_400_000)
    title = _safe_string(value["title"], 1, 600)
    if not title.strip():
        raise ProtocolError("E_JOB")
    published = _safe_string(value["published_at"], 1, 64)
    try:
        parsed = datetime.fromisoformat(published)
    except ValueError as exc:
        raise ProtocolError("E_JOB") from exc
    if parsed.utcoffset() is None:
        raise ProtocolError("E_JOB")
    return value


def validate_hello(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(value, "hello", {"schema", "type", "extension_build"})
    _safe_string(value["extension_build"], 1, 128)
    return value


def validate_poll(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(value, "poll", {"schema", "type"})
    return value


def validate_foreground(value: dict[str, Any]) -> dict[str, Any]:
    """Validate one pre-secret request to foreground the bound Chrome window.

    Only opaque Chrome identifiers and screen geometry cross the wire.  The
    Native Host derives the canonical target URL from the already-claimed job;
    callers cannot supply an arbitrary command, executable, URL, title, or
    profile path.
    """

    _schema_type(
        value,
        "foreground",
        {"schema", "type", "job_id", "lease_id", "tab_id", "window_id", "window_bounds"},
    )
    if not isinstance(value["job_id"], str) or not JOB_ID_RE.fullmatch(value["job_id"]):
        raise ProtocolError("E_JOB")
    _lease(value["lease_id"])
    _integer(value["tab_id"], 0, MAX_SAFE_INTEGER)
    _integer(value["window_id"], 0, MAX_SAFE_INTEGER)
    bounds = value["window_bounds"]
    if not isinstance(bounds, dict):
        raise ProtocolError("E_FOREGROUND")
    _exact_keys(bounds, {"left", "top", "width", "height"})
    _integer(bounds["left"], -1_000_000, 1_000_000)
    _integer(bounds["top"], -1_000_000, 1_000_000)
    _integer(bounds["width"], 1, 1_000_000)
    _integer(bounds["height"], 1, 1_000_000)
    return value


def validate_reject(value: dict[str, Any]) -> dict[str, Any]:
    try:
        _schema_type(
            value,
            "reject",
            {"schema", "type", "job_id", "lease_id", "error_code", "diagnostic"},
        )
    except ProtocolError as exc:
        raise ProtocolError("E_REJECT") from exc
    if not isinstance(value["job_id"], str) or not JOB_ID_RE.fullmatch(value["job_id"]):
        raise ProtocolError("E_JOB")
    _lease(value["lease_id"])
    try:
        validate_prepareless_terminal(value["error_code"], value["diagnostic"])
    except ValueError as exc:
        raise ProtocolError("E_REJECT") from exc
    return value


def validate_abort_prepare(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(
        value,
        "abort_prepare",
        {"schema", "type", "job_id", "lease_id", "prepare_id", "error_code", "error_reason"},
    )
    if not isinstance(value["job_id"], str) or not JOB_ID_RE.fullmatch(value["job_id"]):
        raise ProtocolError("E_JOB")
    _lease(value["lease_id"])
    if not isinstance(value["prepare_id"], str) or not _PREPARE_ID_RE.fullmatch(value["prepare_id"]):
        raise ProtocolError("E_PREPARE")
    if value["error_code"] not in PRESTART_ABORT_CODES:
        raise ProtocolError("E_PREPARE")
    if value["error_reason"] not in COOKIE_ACCESS_REASONS:
        raise ProtocolError("E_PREPARE")
    return value


def validate_reload_begin(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(value, "reload_begin", {"schema", "type", "extension_build", "reload_token"})
    _safe_string(value["extension_build"], 1, 128)
    if not isinstance(value["reload_token"], str) or not _RELOAD_TOKEN_RE.fullmatch(value["reload_token"]):
        raise ProtocolError("E_RELOAD")
    return value


def _lease(value: Any) -> str:
    if not isinstance(value, str) or not _LEASE_ID_RE.fullmatch(value):
        raise ProtocolError("E_LEASE")
    return value


def validate_status(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(value, "status", {"schema", "type", "job_id", "lease_id"})
    if not isinstance(value["job_id"], str) or not JOB_ID_RE.fullmatch(value["job_id"]):
        raise ProtocolError("E_JOB")
    _lease(value["lease_id"])
    return value


def validate_cancel(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(value, "cancel", {"schema", "type", "job_id", "lease_id", "task_nonce"})
    if not isinstance(value["job_id"], str) or not JOB_ID_RE.fullmatch(value["job_id"]):
        raise ProtocolError("E_JOB")
    _lease(value["lease_id"])
    if not isinstance(value["task_nonce"], str) or not _NONCE_RE.fullmatch(value["task_nonce"]):
        raise ProtocolError()
    return value


def validate_prepare(value: dict[str, Any], now_ms: int | None = None) -> dict[str, Any]:
    _schema_type(
        value,
        "prepare",
        {"schema", "type", "extension_build", "lease_id", "prepare_id", "job", "page_proof"},
    )
    if value["extension_build"] != EXTENSION_BUILD:
        raise ProtocolError("E_BUILD")
    _lease(value["lease_id"])
    if not isinstance(value["prepare_id"], str) or not _PREPARE_ID_RE.fullmatch(value["prepare_id"]):
        raise ProtocolError("E_PREPARE")
    job = validate_job(value["job"])
    validate_page_proof(value["page_proof"], job, now_ms=now_ms)
    return value


def validate_worker_prepare(value: dict[str, Any]) -> dict[str, Any]:
    _schema_type(
        value, "worker_prepare", {"schema", "type", "job", "lease_id", "recovery_mode"}
    )
    validate_job(value["job"])
    _lease(value["lease_id"])
    if value["recovery_mode"] not in {"NONE", "EXACT_PUBLISHED_PAIR"}:
        raise ProtocolError("E_PREPARE")
    return value


_MEDIA_COMPLETE_KEYS = {
    "formal_filename", "mapping_filename", "media_bytes", "media_sha256",
    "mapping_bytes", "mapping_sha256", "duration_milliseconds",
    "video_codec", "audio_codec",
}


def validate_media_complete_identity(value: Any, job: dict[str, Any]) -> dict[str, Any]:
    """Validate the exact, non-secret identity durably ACKed before postprocess."""

    if not isinstance(job, dict):
        raise ProtocolError("E_MEDIA_COMPLETE")
    try:
        bvid = validate_bvid(job.get("bvid"))
    except (TypeError, ValueError) as exc:
        raise ProtocolError("E_MEDIA_COMPLETE") from exc
    if not isinstance(value, dict):
        raise ProtocolError("E_MEDIA_COMPLETE")
    _exact_keys(value, _MEDIA_COMPLETE_KEYS)
    if (
        value["formal_filename"] != f"{bvid}.mkv"
        or value["mapping_filename"] != f"{bvid}.download.json"
    ):
        raise ProtocolError("E_MEDIA_COMPLETE")
    _integer(value["media_bytes"], 1, MAX_SAFE_INTEGER)
    _integer(value["mapping_bytes"], 1, 65_536)
    _integer(value["duration_milliseconds"], 1, MAX_SAFE_INTEGER)
    for name in ("media_sha256", "mapping_sha256"):
        if not isinstance(value[name], str) or UPPER_SHA256_RE.fullmatch(value[name]) is None:
            raise ProtocolError("E_MEDIA_COMPLETE")
    for name in ("video_codec", "audio_codec"):
        if (
            not isinstance(value[name], str)
            or re.fullmatch(r"[A-Za-z0-9_.-]{1,64}", value[name]) is None
        ):
            raise ProtocolError("E_MEDIA_COMPLETE")
    return value


def validate_media_complete_ack(
    value: dict[str, Any], job: dict[str, Any], lease_id: str,
    media: dict[str, Any],
) -> dict[str, Any]:
    """Validate the Host's typed durable-state acknowledgement."""

    _schema_type(
        value, "media_complete_ack",
        {"schema", "type", "job_id", "lease_id", "media"},
    )
    validated_job = validate_job(job)
    _lease(value["lease_id"])
    if value["job_id"] != validated_job["job_id"] or value["lease_id"] != lease_id:
        raise ProtocolError("E_MEDIA_COMPLETE")
    validated_media = validate_media_complete_identity(value["media"], validated_job)
    if validated_media != media:
        raise ProtocolError("E_MEDIA_COMPLETE")
    return value


def validate_page_proof(value: Any, job: dict[str, Any], now_ms: int | None = None) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ProtocolError()
    _exact_keys(
        value,
        {
            "job_id", "bvid", "creator_uid", "canonical_url", "task_nonce",
            "observed_at_unix_ms", "observed_duration_ms", "video_width",
            "video_height", "ready_state", "eme_present", "metadata_source",
        },
    )
    for key in ("job_id", "bvid", "creator_uid", "canonical_url"):
        if value[key] != job[key]:
            raise ProtocolError("E_PAGE_PROOF")
    if not isinstance(value["task_nonce"], str) or not _NONCE_RE.fullmatch(value["task_nonce"]):
        raise ProtocolError()
    observed_at = _integer(value["observed_at_unix_ms"], 1, MAX_SAFE_INTEGER)
    observed_duration = _integer(value["observed_duration_ms"], 1, MAX_SAFE_INTEGER)
    _integer(value["video_width"], 1, 7680)
    _integer(value["video_height"], 1, 4320)
    ready_state = _integer(value["ready_state"], 1, 4)
    if value["metadata_source"] not in {"HTML_MEDIA_ELEMENT", "BILIBILI_INITIAL_STATE"}:
        raise ProtocolError("E_PAGE_PROOF")
    if value["metadata_source"] == "BILIBILI_INITIAL_STATE" and ready_state != 1:
        raise ProtocolError("E_PAGE_PROOF")
    if _boolean(value["eme_present"]):
        raise ProtocolError("E_DRM")
    if abs(observed_duration - job["expected_duration_ms"]) > duration_tolerance_ms(job["expected_duration_ms"]):
        raise ProtocolError("E_DURATION")
    current = int(time.time() * 1000) if now_ms is None else now_ms
    if observed_at < current - 60_000 or observed_at > current + 5_000:
        raise ProtocolError("E_PAGE_PROOF")
    return value


def _cookie_path_matches(path: str, bvid: str) -> bool:
    wanted = target_path(bvid)
    return path == "/" or (
        wanted.startswith(path) and (path.endswith("/") or len(path) == len(wanted) or wanted[len(path)] == "/")
    )


def validate_cookie(value: Any, store_id: str, observed_at_unix_ms: int, bvid: str) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ProtocolError("E_SECRET_INPUT")
    try:
        _exact_keys(
            value,
            {"name", "value", "domain", "host_only", "path", "secure", "http_only",
             "same_site", "session", "expiration_unix", "store_id", "partition_key"},
        )
        _safe_string(value["name"], 1, 256)
        _safe_string(value["value"], 1, 4096)
        host_only = _boolean(value["host_only"])
        if value["domain"] != ("www.bilibili.com" if host_only else ".bilibili.com"):
            raise ProtocolError()
        path = _safe_string(value["path"], 1, 1024)
        if not path.startswith("/") or not _cookie_path_matches(path, bvid):
            raise ProtocolError()
        _boolean(value["secure"])
        _boolean(value["http_only"])
        session = _boolean(value["session"])
        if value["same_site"] not in _COOKIE_SAME_SITE or value["store_id"] != store_id or value["partition_key"] is not None:
            raise ProtocolError()
        if session:
            if value["expiration_unix"] is not None:
                raise ProtocolError()
        elif _integer(value["expiration_unix"], 1, MAX_SAFE_INTEGER) <= observed_at_unix_ms // 1000:
            raise ProtocolError()
    except ProtocolError as exc:
        raise ProtocolError("E_SECRET_INPUT") from exc
    return value


def validate_start(value: dict[str, Any], now_ms: int | None = None) -> dict[str, Any]:
    _schema_type(
        value,
        "start",
        {"schema", "type", "extension_build", "lease_id", "prepare_id", "job",
         "cookie_store_id", "page_proof", "cookies"},
    )
    if value["extension_build"] != EXTENSION_BUILD:
        raise ProtocolError("E_BUILD")
    _lease(value["lease_id"])
    if not isinstance(value["prepare_id"], str) or not _PREPARE_ID_RE.fullmatch(value["prepare_id"]):
        raise ProtocolError("E_PREPARE")
    job = validate_job(value["job"])
    store_id = value["cookie_store_id"]
    if not isinstance(store_id, str) or not _STORE_RE.fullmatch(store_id):
        raise ProtocolError("E_SECRET_INPUT")
    proof = validate_page_proof(value["page_proof"], job, now_ms=now_ms)
    cookies = value["cookies"]
    if not isinstance(cookies, list) or not 1 <= len(cookies) <= MAX_COOKIE_COUNT:
        raise ProtocolError("E_SECRET_INPUT")
    for cookie in cookies:
        validate_cookie(cookie, store_id, proof["observed_at_unix_ms"], job["bvid"])
    return value


def validate_message(value: dict[str, Any], now_ms: int | None = None) -> dict[str, Any]:
    message_type = value.get("type")
    if message_type == "hello":
        return validate_hello(value)
    if message_type == "poll":
        return validate_poll(value)
    if message_type == "foreground":
        return validate_foreground(value)
    if message_type == "reject":
        return validate_reject(value)
    if message_type == "abort_prepare":
        return validate_abort_prepare(value)
    if message_type == "reload_begin":
        return validate_reload_begin(value)
    if message_type == "status":
        return validate_status(value)
    if message_type == "cancel":
        return validate_cancel(value)
    if message_type == "prepare":
        return validate_prepare(value, now_ms=now_ms)
    if message_type == "start":
        return validate_start(value, now_ms=now_ms)
    raise ProtocolError()


def maintenance_state(
    *, required_extension_build: str = EXTENSION_BUILD, reload_required: bool = False,
    reload_token: str | None = None, retry_after_unix_ms: int = 0,
) -> dict[str, Any]:
    return {
        "required_extension_build": required_extension_build,
        "reload_required": bool(reload_required),
        "reload_token": reload_token,
        "retry_after_unix_ms": max(0, int(retry_after_unix_ms)),
    }


def safe_response(
    message_type: str,
    phase: str,
    *,
    progress: int = 0,
    error_code: str | None = None,
    formal_filename: str | None = None,
    mapping_filename: str | None = None,
    prepare_id: str | None = None,
    job: dict[str, Any] | None = None,
    lease_id: str | None = None,
    maintenance: dict[str, Any] | None = None,
) -> dict[str, Any]:
    return {
        "schema": SCHEMA_VERSION,
        "type": message_type,
        "host_build": HOST_BUILD,
        "phase": phase,
        "progress": max(0, min(100, int(progress))),
        "error_code": error_code,
        "formal_filename": formal_filename,
        "mapping_filename": mapping_filename,
        "prepare_id": prepare_id,
        "job": job,
        "lease_id": lease_id,
        "maintenance": maintenance if maintenance is not None else maintenance_state(),
    }