from __future__ import annotations

import hmac
from dataclasses import dataclass, field
from enum import IntEnum

from k1link.device_plugins.xgrids_k1.protocol.protobuf_wire import (
    ProtobufWireError,
    ProtoField,
    iter_fields,
)

# The values below are scoped to the reviewed LixelGO/K1 protocol profile.  This
# module deliberately has no MQTT dependency: it can build and validate bytes,
# but it cannot send them to equipment.
MAX_CONTROL_PAYLOAD_BYTES = 64 * 1024
MAX_HEADER_BYTES = 4 * 1024
MAX_TEXT_BYTES = 4 * 1024
MODELING_SESSION_SUFFIX = ":ModelingRequest"

OPENAPI_RESULT_BASE = 302_252_032
OPENAPI_SUCCESS = OPENAPI_RESULT_BASE + 1
MODELING_STATE_BASE = 302_252_032
SYSTEM_ERROR_STATE_BASE = 0x3204_0000


class ModelingProtocolError(ValueError):
    """A bounded K1 modeling-control payload violated the recovered schema."""


class ModelingEncodeError(ModelingProtocolError):
    """A command could not be encoded without guessing vendor-owned input."""


class ModelingResponseCorrelationError(ModelingProtocolError):
    """A response did not match the exact command identity and action."""


class ModelingCommandRejected(ModelingProtocolError):
    """The device returned a non-success result for a correlated command."""

    def __init__(self, code: int, description: str) -> None:
        # Device descriptions are retained for a caller that needs structured
        # diagnostics, but are not interpolated into the exception string.
        super().__init__(f"K1 modeling command rejected with result code {code}")
        self.code = code
        self.description = description


class ModelingAction(IntEnum):
    START = 1
    STOP = 2


class RecordMode(IntEnum):
    CALCULATE_ONLY = 0
    RECORD_ONLY = 1
    RECORD_AND_CALCULATE = 2


class ScanMode(IntEnum):
    POINT_CLOUD = 0
    LCC = 1
    PORTRAIT = 2


class MountType(IntEnum):
    HANDHELD = 0
    VEHICLE = 1
    UAV = 2
    BACKPACK = 3


class SessionState(IntEnum):
    OTHER_STATUS = 0
    READY = 300
    SCAN_STARTING = 301
    SCANNING = 302
    SCAN_STOPPING = 303
    SCAN_OVER = 304
    DISK_ERROR = 305
    SAVE_ERROR = 306
    ALGORITHM_ERROR = 307
    NO_CONTINUE_ERROR = 308
    NO_SPACE_ERROR = 309
    CAMERA_ERROR = 310
    CONTINUE_SUCCESS = 311
    CONTINUE_FAIL = 312
    USB_DISK = 313
    SD_SPACE_NOT_ENOUGH = 314
    MEMORY_NOT_ENOUGH = 315
    LIDAR_DATA_ERROR = 316
    MAPPING_ERROR = 317

    @property
    def is_fault(self) -> bool:
        return self in _FAULT_SESSION_STATES


_FAULT_SESSION_STATES = frozenset(
    {
        SessionState.DISK_ERROR,
        SessionState.SAVE_ERROR,
        SessionState.ALGORITHM_ERROR,
        SessionState.NO_CONTINUE_ERROR,
        SessionState.NO_SPACE_ERROR,
        SessionState.CAMERA_ERROR,
        SessionState.CONTINUE_FAIL,
        SessionState.SD_SPACE_NOT_ENOUGH,
        SessionState.MEMORY_NOT_ENOUGH,
        SessionState.LIDAR_DATA_ERROR,
        SessionState.MAPPING_ERROR,
    }
)


@dataclass(frozen=True, slots=True)
class CommandHeaderIdentity:
    """Exact recovered modeling identity; only the credential stays caller-provided."""

    device_id: str = field(repr=False)
    openapi_key: str = field(repr=False)

    def __post_init__(self) -> None:
        _validate_required_ascii_identity(self.device_id, "device_id", MAX_TEXT_BYTES)
        _validate_required_ascii_identity(self.openapi_key, "openapi_key", MAX_TEXT_BYTES)
        if len(self.session_id.encode("ascii")) > MAX_TEXT_BYTES:
            raise ModelingEncodeError("derived modeling session_id exceeds configured limit")

    @property
    def session_id(self) -> str:
        """Return the literal relation proven in the retained wire evidence."""

        return f"{self.device_id}{MODELING_SESSION_SUFFIX}"

    def matches(self, observed: ObservedHeader) -> bool:
        """Compare every recovered identity field without exposing its value."""

        if (
            observed.device_id is None
            or observed.session_id is None
            or observed.openapi_key is None
        ):
            return False
        return all(
            hmac.compare_digest(expected.encode("ascii"), actual.encode("utf-8"))
            for expected, actual in (
                (self.device_id, observed.device_id),
                (self.session_id, observed.session_id),
                (self.openapi_key, observed.openapi_key),
            )
        )


@dataclass(frozen=True, slots=True)
class ObservedHeader:
    """Identity subset decoded from an inbound header; repr stays redacted."""

    device_id: str | None = field(default=None, repr=False)
    session_id: str | None = field(default=None, repr=False)
    openapi_key: str | None = field(default=None, repr=False)


@dataclass(frozen=True, slots=True)
class EncodedModelingCommand:
    """An inert command envelope.  No transport operation exists in this type."""

    action: ModelingAction
    header: CommandHeaderIdentity = field(repr=False)
    payload: bytes = field(repr=False)


@dataclass(frozen=True, slots=True)
class ModelingResponseError:
    code: int
    description: str = field(repr=False)


@dataclass(frozen=True, slots=True)
class ModelingResponse:
    header: ObservedHeader = field(repr=False)
    action: ModelingAction
    error: ModelingResponseError


@dataclass(frozen=True, slots=True)
class DeviceStatusReport:
    """Recovered acquisition subset of ``DeviceStatusReport``.

    Vendor identity and project strings are intentionally redacted from repr.
    Nested system/RTK messages are bounded by the outer parser but are not
    interpreted by this modeling state machine.
    """

    header: ObservedHeader | None = field(repr=False)
    modeling_state_code: int
    session_state: SessionState | None
    device_sn: str | None = field(repr=False)
    project_id: str | None = field(repr=False)
    time_zone: str | None = field(repr=False)
    init_ready: bool
    system_status_present: bool
    rtk_status_present: bool


@dataclass(frozen=True, slots=True)
class SystemErrorReport:
    """Bounded K1 system-error report with identity kept out of repr."""

    header: ObservedHeader | None = field(repr=False)
    error_code: int
    session_state: SessionState | None


def encode_modeling_start(
    header: CommandHeaderIdentity,
    *,
    project_name: str,
    record_mode: RecordMode,
    scan_mode: ScanMode,
    mount_type: MountType,
    pre_project_id: str | None = None,
) -> EncodedModelingCommand:
    """Encode the recovered ``ModelingRequest`` start shape.

    All modes are mandatory so that a future caller cannot inherit an implicit
    vendor default.  Proto3 zero-valued enums are canonically omitted, matching
    the reviewed application serializer.
    """

    _expect_exact_enum(record_mode, RecordMode, "record_mode")
    _expect_exact_enum(scan_mode, ScanMode, "scan_mode")
    _expect_exact_enum(mount_type, MountType, "mount_type")
    _validate_required_text(project_name, "project_name", MAX_TEXT_BYTES)
    if pre_project_id is not None:
        _validate_required_text(pre_project_id, "pre_project_id", MAX_TEXT_BYTES)

    fields = [
        _bytes_field(1, _encode_command_header(header)),
        _varint_field(2, ModelingAction.START),
        _text_field(3, project_name),
    ]
    if record_mode:
        fields.append(_varint_field(4, record_mode))
    if scan_mode:
        fields.append(_varint_field(5, scan_mode))
    if mount_type:
        fields.append(_varint_field(6, mount_type))
    if pre_project_id is not None:
        fields.append(_text_field(7, pre_project_id))

    payload = b"".join(fields)
    _ensure_encoded_bound(payload)
    return EncodedModelingCommand(ModelingAction.START, header, payload)


def encode_modeling_stop(header: CommandHeaderIdentity) -> EncodedModelingCommand:
    """Encode the recovered stop shape: header plus action, and nothing else."""

    payload = b"".join(
        (
            _bytes_field(1, _encode_command_header(header)),
            _varint_field(2, ModelingAction.STOP),
        )
    )
    _ensure_encoded_bound(payload)
    return EncodedModelingCommand(ModelingAction.STOP, header, payload)


def decode_modeling_response(
    payload: bytes, *, max_payload_bytes: int = MAX_CONTROL_PAYLOAD_BYTES
) -> ModelingResponse:
    """Decode a bounded response without treating its description as success."""

    _check_payload_bound(payload, max_payload_bytes, "ModelingResponse")
    header: ObservedHeader | None = None
    action: ModelingAction | None = None
    response_error: ModelingResponseError | None = None
    seen: set[int] = set()
    try:
        for proto_field in iter_fields(payload, max_fields=64):
            if proto_field.number == 1:
                _mark_once(seen, 1, "response.header")
                header = _decode_observed_header(
                    _bytes_value(proto_field, "response.header"), require_identity=True
                )
            elif proto_field.number == 2:
                _mark_once(seen, 2, "response.action")
                raw_action = _uint_value(proto_field, "response.action", max_value=0xFFFF_FFFF)
                try:
                    action = ModelingAction(raw_action)
                except ValueError as exc:
                    raise ModelingProtocolError("response.action is not start or stop") from exc
            elif proto_field.number == 15:
                _mark_once(seen, 15, "response.error")
                response_error = _decode_response_error(_bytes_value(proto_field, "response.error"))
    except ProtobufWireError as exc:
        raise ModelingProtocolError(f"invalid ModelingResponse: {exc}") from exc

    if header is None:
        raise ModelingProtocolError("ModelingResponse has no header")
    if action is None:
        raise ModelingProtocolError("ModelingResponse has no action")
    if response_error is None:
        raise ModelingProtocolError("ModelingResponse has no error result")
    return ModelingResponse(header, action, response_error)


def correlate_modeling_response(
    payload: bytes,
    expected: EncodedModelingCommand,
    *,
    max_payload_bytes: int = MAX_CONTROL_PAYLOAD_BYTES,
) -> ModelingResponse:
    """Require identity, action, and the physically observed success code."""

    response = decode_modeling_response(payload, max_payload_bytes=max_payload_bytes)
    if not expected.header.matches(response.header):
        raise ModelingResponseCorrelationError("response header identity mismatch")
    if response.action is not expected.action:
        raise ModelingResponseCorrelationError("response action mismatch")
    if response.error.code != OPENAPI_SUCCESS:
        raise ModelingCommandRejected(response.error.code, response.error.description)
    return response


def decode_device_status_report(
    payload: bytes, *, max_payload_bytes: int = MAX_CONTROL_PAYLOAD_BYTES
) -> DeviceStatusReport:
    """Decode the fields that bound the K1 acquisition lifecycle."""

    _check_payload_bound(payload, max_payload_bytes, "DeviceStatusReport")
    header: ObservedHeader | None = None
    modeling_state_code: int | None = None
    device_sn: str | None = None
    project_id: str | None = None
    time_zone: str | None = None
    init_ready = False
    system_status_present = False
    rtk_status_present = False
    seen: set[int] = set()
    try:
        for proto_field in iter_fields(payload, max_fields=64):
            if proto_field.number == 1:
                _mark_once(seen, 1, "device_status.header")
                header = _decode_observed_header(
                    _bytes_value(proto_field, "device_status.header"), require_identity=False
                )
            elif proto_field.number == 2:
                _mark_once(seen, 2, "device_status.modeling_state")
                modeling_state_code = _uint_value(
                    proto_field, "device_status.modeling_state", max_value=0xFFFF_FFFF
                )
            elif proto_field.number == 3:
                _mark_once(seen, 3, "device_status.device_sn")
                device_sn = _text_value(proto_field, "device_status.device_sn")
            elif proto_field.number == 4:
                _mark_once(seen, 4, "device_status.project_id")
                project_id = _text_value(proto_field, "device_status.project_id")
            elif proto_field.number == 5:
                _mark_once(seen, 5, "device_status.time_zone")
                time_zone = _text_value(proto_field, "device_status.time_zone")
            elif proto_field.number == 6:
                _mark_once(seen, 6, "device_status.init_ready")
                raw_ready = _uint_value(proto_field, "device_status.init_ready", max_value=1)
                init_ready = bool(raw_ready)
            elif proto_field.number == 7:
                _mark_once(seen, 7, "device_status.system_status")
                _bytes_value(proto_field, "device_status.system_status")
                system_status_present = True
            elif proto_field.number == 8:
                _mark_once(seen, 8, "device_status.rtk_status")
                _bytes_value(proto_field, "device_status.rtk_status")
                rtk_status_present = True
    except ProtobufWireError as exc:
        raise ModelingProtocolError(f"invalid DeviceStatusReport: {exc}") from exc

    if modeling_state_code is None:
        raise ModelingProtocolError("DeviceStatusReport has no modeling_state")
    return DeviceStatusReport(
        header=header,
        modeling_state_code=modeling_state_code,
        session_state=session_state_from_code(modeling_state_code),
        device_sn=device_sn,
        project_id=project_id,
        time_zone=time_zone,
        init_ready=init_ready,
        system_status_present=system_status_present,
        rtk_status_present=rtk_status_present,
    )


def session_state_from_code(code: int) -> SessionState | None:
    """Map the captured base-offset status code, preserving unknown values."""

    if code < MODELING_STATE_BASE:
        return None
    try:
        return SessionState(code - MODELING_STATE_BASE)
    except ValueError:
        return None


def decode_system_error_report(
    payload: bytes, *, max_payload_bytes: int = MAX_CONTROL_PAYLOAD_BYTES
) -> SystemErrorReport:
    """Decode the observed ``report/system_error`` envelope without guessing text."""

    _check_payload_bound(payload, max_payload_bytes, "SystemErrorReport")
    header: ObservedHeader | None = None
    error_code: int | None = None
    seen: set[int] = set()
    try:
        for proto_field in iter_fields(payload, max_fields=64):
            if proto_field.number == 1:
                _mark_once(seen, 1, "system_error.header")
                header = _decode_observed_header(
                    _bytes_value(proto_field, "system_error.header"), require_identity=False
                )
            elif proto_field.number == 15:
                _mark_once(seen, 15, "system_error.error")
                nested_seen: set[int] = set()
                for nested in iter_fields(
                    _bytes_value(proto_field, "system_error.error"), max_fields=8
                ):
                    if nested.number == 1:
                        _mark_once(nested_seen, 1, "system_error.error.code")
                        error_code = _uint_value(
                            nested,
                            "system_error.error.code",
                            max_value=0xFFFF_FFFF,
                        )
    except ProtobufWireError as exc:
        raise ModelingProtocolError(f"invalid SystemErrorReport: {exc}") from exc
    if error_code is None:
        raise ModelingProtocolError("SystemErrorReport has no error code")
    return SystemErrorReport(
        header=header,
        error_code=error_code,
        session_state=system_error_state_from_code(error_code),
    )


def system_error_state_from_code(code: int) -> SessionState | None:
    """Map the observed system-error namespace while preserving unknown codes."""

    if code < SYSTEM_ERROR_STATE_BASE:
        return None
    try:
        return SessionState(code - SYSTEM_ERROR_STATE_BASE)
    except ValueError:
        return None


def _encode_command_header(header: CommandHeaderIdentity) -> bytes:
    # seq/stamp/scaler were default-valued and omitted in the reviewed command
    # requests.  The three identities must be supplied explicitly by a caller.
    payload = b"".join(
        (
            _text_field(4, header.device_id),
            _text_field(5, header.session_id),
            _text_field(6, header.openapi_key),
        )
    )
    if len(payload) > MAX_HEADER_BYTES:
        raise ModelingEncodeError("encoded command header exceeds configured limit")
    return payload


def _decode_observed_header(payload: bytes, *, require_identity: bool) -> ObservedHeader:
    if len(payload) > MAX_HEADER_BYTES:
        raise ModelingProtocolError("header exceeds configured limit")
    identities: dict[int, str] = {}
    seen: set[int] = set()
    try:
        for proto_field in iter_fields(payload, max_fields=64):
            if proto_field.number in {4, 5, 6}:
                _mark_once(seen, proto_field.number, "header identity")
                identities[proto_field.number] = _identity_text_value(
                    proto_field,
                    "header identity",
                )
    except ProtobufWireError as exc:
        raise ModelingProtocolError(f"invalid header: {exc}") from exc

    header = ObservedHeader(
        device_id=identities.get(4),
        session_id=identities.get(5),
        openapi_key=identities.get(6),
    )
    if require_identity and (
        header.device_id is None or header.session_id is None or header.openapi_key is None
    ):
        raise ModelingProtocolError("response header identity is incomplete")
    return header


def _decode_response_error(payload: bytes) -> ModelingResponseError:
    if len(payload) > MAX_HEADER_BYTES:
        raise ModelingProtocolError("response error exceeds configured limit")
    code: int | None = None
    description = ""
    seen: set[int] = set()
    try:
        for proto_field in iter_fields(payload, max_fields=16):
            if proto_field.number == 1:
                _mark_once(seen, 1, "response.error.code")
                code = _uint_value(proto_field, "response.error.code", max_value=0xFFFF_FFFF)
            elif proto_field.number == 2:
                _mark_once(seen, 2, "response.error.description")
                description = _text_value(proto_field, "response.error.description")
    except ProtobufWireError as exc:
        raise ModelingProtocolError(f"invalid response error: {exc}") from exc
    if code is None:
        raise ModelingProtocolError("response error has no code")
    return ModelingResponseError(code, description)


def _expect_exact_enum(value: object, enum_type: type[IntEnum], name: str) -> None:
    if type(value) is not enum_type:
        raise ModelingEncodeError(f"{name} must be an explicit {enum_type.__name__}")


def _validate_required_text(value: object, name: str, max_bytes: int) -> None:
    if not isinstance(value, str) or not value:
        raise ModelingEncodeError(f"{name} must be a non-empty string")
    try:
        encoded = value.encode("utf-8")
    except UnicodeEncodeError as exc:
        raise ModelingEncodeError(f"{name} is not valid UTF-8 text") from exc
    if len(encoded) > max_bytes:
        raise ModelingEncodeError(f"{name} exceeds configured limit")
    if any(ord(character) < 0x20 or ord(character) == 0x7F for character in value):
        raise ModelingEncodeError(f"{name} contains a control character")


def _validate_required_ascii_identity(value: object, name: str, max_bytes: int) -> None:
    if not isinstance(value, str) or not value:
        raise ModelingEncodeError(f"{name} must be a non-empty string")
    try:
        encoded = value.encode("ascii")
    except UnicodeEncodeError as exc:
        raise ModelingEncodeError(f"{name} must use printable ASCII") from exc
    if len(encoded) > max_bytes:
        raise ModelingEncodeError(f"{name} exceeds configured limit")
    if any(byte <= 0x20 or byte > 0x7E for byte in encoded):
        raise ModelingEncodeError(f"{name} must use printable ASCII without spaces")


def _mark_once(seen: set[int], number: int, name: str) -> None:
    if number in seen:
        raise ModelingProtocolError(f"{name} is duplicated")
    seen.add(number)


def _check_payload_bound(payload: bytes, max_payload_bytes: int, name: str) -> None:
    if max_payload_bytes < 1:
        raise ValueError("max_payload_bytes must be positive")
    if len(payload) > max_payload_bytes:
        raise ModelingProtocolError(f"{name} exceeds configured limit")


def _ensure_encoded_bound(payload: bytes) -> None:
    if len(payload) > MAX_CONTROL_PAYLOAD_BYTES:
        raise ModelingEncodeError("encoded ModelingRequest exceeds configured limit")


def _uint_value(proto_field: ProtoField, name: str, *, max_value: int) -> int:
    if proto_field.wire_type != 0 or not isinstance(proto_field.value, int):
        raise ModelingProtocolError(f"{name} has the wrong protobuf wire type")
    if proto_field.value > max_value:
        raise ModelingProtocolError(f"{name} exceeds its recovered integer range")
    return proto_field.value


def _bytes_value(proto_field: ProtoField, name: str) -> bytes:
    if proto_field.wire_type != 2 or not isinstance(proto_field.value, bytes):
        raise ModelingProtocolError(f"{name} has the wrong protobuf wire type")
    return proto_field.value


def _text_value(proto_field: ProtoField, name: str) -> str:
    raw = _bytes_value(proto_field, name)
    if len(raw) > MAX_TEXT_BYTES:
        raise ModelingProtocolError(f"{name} exceeds configured limit")
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ModelingProtocolError(f"{name} is not valid UTF-8") from exc


def _identity_text_value(proto_field: ProtoField, name: str) -> str:
    value = _text_value(proto_field, name)
    try:
        encoded = value.encode("ascii")
    except UnicodeEncodeError as exc:
        raise ModelingProtocolError(f"{name} must use printable ASCII") from exc
    if not encoded or any(byte <= 0x20 or byte > 0x7E for byte in encoded):
        raise ModelingProtocolError(f"{name} must use printable ASCII without spaces")
    return value


def _varint(value: int) -> bytes:
    if value < 0 or value > 0xFFFF_FFFF_FFFF_FFFF:
        raise ModelingEncodeError("protobuf varint is outside uint64")
    encoded = bytearray()
    while value > 0x7F:
        encoded.append((value & 0x7F) | 0x80)
        value >>= 7
    encoded.append(value)
    return bytes(encoded)


def _key(number: int, wire_type: int) -> bytes:
    if number < 1:
        raise ModelingEncodeError("protobuf field number must be positive")
    return _varint((number << 3) | wire_type)


def _varint_field(number: int, value: int) -> bytes:
    return _key(number, 0) + _varint(int(value))


def _bytes_field(number: int, value: bytes) -> bytes:
    return _key(number, 2) + _varint(len(value)) + value


def _text_field(number: int, value: str) -> bytes:
    return _bytes_field(number, value.encode("utf-8"))
